import os
import gradio as gr
from config import DATA_DIR, VECTOR_DIR, UPLOADED_PDF_PATH, GROQ_API_KEY
from ingestion import ingest_pdf
from rag import answer, reset_memory

# Prepare runtime folders
os.makedirs(DATA_DIR, exist_ok=True)
os.makedirs(VECTOR_DIR, exist_ok=True)

def do_ingest(pdf_file):
    if GROQ_API_KEY.strip() == "":
        return "⚠️ GROQ_API_KEY is missing. Add it in your environment / HF Space secrets."

    if pdf_file is None:
        return "Please upload a PDF."

    # Save uploaded file
    with open(UPLOADED_PDF_PATH, "wb") as f:
        f.write(pdf_file.read())

    # Ingest and reset memory
    msg = ingest_pdf(UPLOADED_PDF_PATH)
    reset_memory()
    return f"✅ {msg}"

def chat_fn(user_message, history):
    # Ensure vector index exists
    if not os.path.exists(VECTOR_DIR) or not os.listdir(VECTOR_DIR):
        return "Please upload a PDF and click 'Ingest PDF' first.", history

    ans = answer(user_message)
    history.append((user_message, ans))
    return "", history

with gr.Blocks() as demo:
    gr.Markdown("## 📚 Chat with Your PDF — FAISS + Groq + Memory (LangChain)")

    with gr.Row():
        pdf = gr.File(label="Upload PDF", file_types=[".pdf"], type="file")
        ingest_btn = gr.Button("Ingest PDF")

    status = gr.Textbox(label="Status", interactive=False)

    chatbot = gr.Chatbot(label="Chat")
    msg = gr.Textbox(label="Ask a question about the PDF and press Enter")
    clear = gr.Button("Clear Chat")

    ingest_btn.click(do_ingest, inputs=pdf, outputs=status)
    msg.submit(chat_fn, [msg, chatbot], [msg, chatbot])
    clear.click(lambda: None, None, chatbot, queue=False)

if __name__ == "__main__":
    demo.launch()