import gradio as gr
from rag_pipeline import build_rag_chain
from dotenv import load_dotenv

load_dotenv()

chain = None

def upload_file(files):
    global chain
    if files is None:
        return "No file uploaded."
    if isinstance(files, str):
        paths = [files]
    else:
        paths = [f if isinstance(f, str) else f.name for f in files] if isinstance(files, list) else [files.name]
    
    full_chain, _ = build_rag_chain(paths)
    chain = full_chain
    return f"✅ {len(paths)} document(s) loaded. Ask away!"

def ask_question(message, history):
    if chain is None:
        return "Upload a document first."
    
    history_text = ""
    for h in history:
        if isinstance(h, (list, tuple)) and len(h) == 2:
            history_text += f"User: {h[0]}\nAssistant: {h[1]}\n"
    
    full_query = f"{history_text}Question: {message}" if history_text else message
    
    result = chain.invoke(full_query)
    answer = result["answer"]
    pages = sorted(set(
        d.metadata.get("page", 0) + 1
        for d in result["source_documents"]
    ))
    return f"{answer}\n\n📄 Source pages: {pages}"

with gr.Blocks() as full_app:
    gr.Markdown("# 📄 RAG Document Q&A")
    file_input = gr.File(label="Upload PDF(s)", file_types=[".pdf"], file_count="multiple")
    status = gr.Textbox(label="Status", interactive=False)
    file_input.upload(upload_file, inputs=file_input, outputs=status)
    gr.Markdown("---")
    question = gr.Textbox(label="Ask a question")
    answer = gr.Textbox(label="Answer", interactive=False, lines=10)
    btn = gr.Button("Ask")
    btn.click(fn=ask_question, inputs=[question, answer], outputs=answer)

full_app.launch()