import gradio as gr
from rag_pipeline import RAGPipeline
from generator import QuestionGenerator
from utils import load_text
from pathlib import Path

# Initialize pipeline and generator
rag = RAGPipeline(db_dir="./chroma_store")
gen = QuestionGenerator()

# Ensure data folder exists
DATA_DIR = Path("./data")
DATA_DIR.mkdir(exist_ok=True)

# Load a sample document at startup (optional)
sample_doc_path = DATA_DIR / "sample_doc.txt"
if sample_doc_path.exists():
    text = load_text(sample_doc_path)
    rag.add_document(text, doc_id="sample_doc")

# Functions for Gradio interface
def upload_file(file):
    text = load_text(file.name)
    rag.add_document(text, doc_id=file.name)
    return f"Document '{file.name}' added successfully."

def ask_question(query, top_k=3):
    retrieved_docs = rag.retrieve(query, top_k=top_k)
    if not retrieved_docs:
        context = "No relevant documents found."
    else:
        context = "\n".join(retrieved_docs)
    answer = gen.generate(context, query)
    return answer

# Gradio UI
with gr.Blocks(title="RAG Question Generator") as demo:

    gr.Markdown("## Upload Documents")
    with gr.Row():
        file_upload = gr.File(label="Upload TXT / PDF / DOCX", file_types=[".txt", ".pdf", ".docx"])
        upload_btn = gr.Button("Upload")
        upload_status = gr.Textbox(label="Status", interactive=False)

    upload_btn.click(fn=upload_file, inputs=file_upload, outputs=upload_status)

    gr.Markdown("## Ask Questions")
    query_input = gr.Textbox(label="Your Question")
    top_k_slider = gr.Slider(1, 10, value=3, step=1, label="Number of retrieved chunks (top_k)")
    ask_btn = gr.Button("Ask")
    answer_output = gr.Textbox(label="Generated Answer", interactive=False)

    ask_btn.click(fn=ask_question, inputs=[query_input, top_k_slider], outputs=answer_output)
    
demo.launch(share=True)