import os
import tempfile
from pathlib import Path
import gradio as gr
from pypdf import PdfReader

from agents import GroqClient, SerpClient
from multi_agent import MultiAgentOrchestrator
from docx_builder import (
    build_question_paper_docx,
    build_answers_docx,
    build_obe_docx,
)

# ---------------------------
# Utility: Extract PDF text
# ---------------------------
def extract_text_from_pdf(file_obj) -> str:
    try:
        reader = PdfReader(file_obj.name)
        pages = [p.extract_text() or "" for p in reader.pages]
        return "\n".join(pages)
    except:
        return ""

# ---------------------------
# Initialize Clients
# ---------------------------
def init_clients():
    groq_key = os.getenv("GROQ_API_KEY")
    serp_key = os.getenv("SERPAPI_KEY")
    groq = GroqClient(api_key=groq_key)
    serp = SerpClient(api_key=serp_key)
    return groq, serp

groq_client, serp_client = init_clients()
orchestrator = MultiAgentOrchestrator(groq_client, serp_client)

# ---------------------------
# Main Pipeline Trigger
# ---------------------------
def run_system(subject, stream, partA, partB, partC, syl_file, ref_file):

    if syl_file is None:
        return None, None, None, "Upload syllabus first."

    syllabus_text = extract_text_from_pdf(syl_file)
    ref_text = extract_text_from_pdf(ref_file) if ref_file else ""

    output = orchestrator.run_pipeline(
        subject=subject,
        stream=stream,
        partA=int(partA),
        partB=int(partB),
        partC=int(partC),
        syllabus_text=syllabus_text,
        ref_qp_text=ref_text,
    )

    final_json = output.get("final", {})
    generator_raw = output.get("generator_raw", "")

    tmpdir = Path(tempfile.mkdtemp())
    qp_path = tmpdir / f"{subject}_QP.docx"
    ans_path = tmpdir / f"{subject}_Answers.docx"
    obe_path = tmpdir / f"{subject}_OBE.docx"

    build_question_paper_docx(qp_path, final_json, generator_raw, subject)
    build_answers_docx(ans_path, final_json, subject)
    build_obe_docx(obe_path, final_json, subject)

    return qp_path, ans_path, obe_path, "Done!"

# ---------------------------
# Gradio UI
# ---------------------------
with gr.Blocks() as app:
    gr.Markdown("# Multi-Agent Question Paper Generator (Groq + SerpAPI)")

    with gr.Row():
        subject = gr.Textbox(label="Subject Name")
        stream = gr.Dropdown(["CSE", "Non-CSE"], value="CSE", label="Stream")

    with gr.Row():
        partA = gr.Number(value=10, precision=0, label="Part A Count")
        partB = gr.Number(value=5, precision=0, label="Part B Count")
        partC = gr.Number(value=1, precision=0, label="Part C Count")

    syllabus = gr.File(label="Upload Syllabus PDF")
    ref_qp = gr.File(label="Upload Reference QP (Optional)")

    btn = gr.Button("Generate Question Paper")

    qp_file = gr.File(label="Question Paper")
    ans_file = gr.File(label="Answer Key")
    obe_file = gr.File(label="OBE Summary")
    status = gr.Markdown("Status: Idle")

    btn.click(
        run_system,
        inputs=[subject, stream, partA, partB, partC, syllabus, ref_qp],
        outputs=[qp_file, ans_file, obe_file, status]
    )

app.launch()