# app.py
import os
import json
import time
import tempfile

import gradio as gr
from fastapi import FastAPI, HTTPException, UploadFile, File
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel, Field, validator

from quiz import QuizManager, SUBJECTS, DIFFICULTIES

# ── Startup check ─────────────────────────────────────────────
api_key = os.getenv("GEMINI_API_KEY")
if not api_key:
    raise ValueError("GEMINI_API_KEY environment variable not set")

# ── Shared quiz manager (one instance, used by BOTH API & UI) ─
#    This is the key fix — no HTTP calls needed
quiz_manager = QuizManager()

# ── FastAPI ───────────────────────────────────────────────────
fastapi_app = FastAPI(
    title="NEET/JEE Quiz Generator API",
    description="RAG-powered MCQ generator for competitive exam prep",
    version="2.0.0"
)

fastapi_app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# ── Pydantic models ───────────────────────────────────────────
class TextRequest(BaseModel):
    text: str = Field(..., min_length=100)
    subject: str = Field(...)
    chapter: str = Field(...)
    topic: str = Field(...)
    exam_type: str = Field("NEET")
    difficulty: str = Field("Medium")
    num_questions: int = Field(10, ge=1, le=30)

    @validator("exam_type")
    def valid_exam(cls, v):
        if v not in ["NEET", "JEE_MAINS", "JEE_ADVANCED"]:
            raise ValueError("Invalid exam type")
        return v

    @validator("difficulty")
    def valid_diff(cls, v):
        if v not in ["Easy", "Medium", "Hard", "Mixed"]:
            raise ValueError("Invalid difficulty")
        return v

class QuizResponse(BaseModel):
    success: bool
    session_id: str
    total_questions: int
    time_limit_minutes: int
    quiz: dict
    generated_in_seconds: float

class HealthResponse(BaseModel):
    status: str
    model: str
    version: str

# ═════════════════════════════════════════════════════════════
# CORE LOGIC  (called by BOTH API endpoints AND Gradio directly)
# ═════════════════════════════════════════════════════════════

def core_generate_from_text(
    text: str,
    subject: str,
    chapter: str,
    topic: str,
    exam_type: str,
    difficulty: str,
    num_questions: int
) -> dict:
    """
    Pure Python function — no HTTP involved.
    Returns a dict that both the API endpoint and Gradio can use.
    Raises ValueError / RuntimeError on bad input.
    """
    start      = time.time()
    num_chunks = quiz_manager.load_and_index(text)

    if num_chunks < 2:
        raise ValueError(
            "Notes too short — please provide more content."
        )

    session = quiz_manager.create_quiz(
        topic         = topic,
        subject       = subject,
        chapter       = chapter,
        exam_type     = exam_type,
        difficulty    = difficulty,
        num_questions = num_questions
    )

    return {
        "success"             : True,
        "session_id"          : session.session_id,
        "total_questions"     : session.total_questions,
        "time_limit_minutes"  : session.time_limit_minutes,
        "quiz"                : session.to_dict(),
        "generated_in_seconds": round(time.time() - start, 2)
    }


def core_generate_from_pdf(
    pdf_path: str,
    subject: str,
    chapter: str,
    topic: str,
    exam_type: str,
    difficulty: str,
    num_questions: int
) -> dict:
    """Same idea — direct Python call, no HTTP."""
    start = time.time()
    quiz_manager.load_and_index(pdf_path)

    session = quiz_manager.create_quiz(
        topic         = topic,
        subject       = subject,
        chapter       = chapter,
        exam_type     = exam_type,
        difficulty    = difficulty,
        num_questions = num_questions
    )

    return {
        "success"             : True,
        "session_id"          : session.session_id,
        "total_questions"     : session.total_questions,
        "time_limit_minutes"  : session.time_limit_minutes,
        "quiz"                : session.to_dict(),
        "generated_in_seconds": round(time.time() - start, 2)
    }


def core_score_quiz(
    answers: dict,
    correct_map: dict,
    exam_type: str = "NEET"
) -> dict:
    """Scoring logic — also pure Python."""
    marks_correct  = 4
    marks_negative = 1
    score = correct = wrong = skipped = 0
    details = {}

    for qid, correct_ans in correct_map.items():
        user_ans = answers.get(qid)
        if user_ans is None:
            skipped += 1
            details[qid] = {"status": "skipped", "marks": 0}
        elif user_ans == correct_ans:
            correct += 1
            score   += marks_correct
            details[qid] = {
                "status"     : "correct",
                "marks"      : marks_correct,
                "your_answer": user_ans
            }
        else:
            wrong += 1
            score -= marks_negative
            details[qid] = {
                "status"        : "wrong",
                "marks"         : -marks_negative,
                "your_answer"   : user_ans,
                "correct_answer": correct_ans
            }

    total_possible = len(correct_map) * marks_correct
    percentage     = round(
        (score / total_possible) * 100, 1
    ) if total_possible else 0

    return {
        "score"         : score,
        "total_possible": total_possible,
        "percentage"    : percentage,
        "correct"       : correct,
        "wrong"         : wrong,
        "skipped"       : skipped,
        "rank_estimate" : _estimate_rank(percentage, exam_type),
        "details"       : details
    }


def _estimate_rank(percentage: float, exam_type: str) -> str:
    if exam_type == "NEET":
        if percentage >= 90: return "Top 1000 (AIR)"
        if percentage >= 80: return "Top 10,000"
        if percentage >= 70: return "Top 50,000"
        if percentage >= 60: return "Top 1,00,000"
        return "Below cut-off range"
    else:
        if percentage >= 85: return "IIT Top-10 branch eligible"
        if percentage >= 70: return "IIT eligible"
        if percentage >= 55: return "NIT Top-tier eligible"
        if percentage >= 40: return "NIT eligible"
        return "Below JEE cut-off range"


# ═════════════════════════════════════════════════════════════
# FASTAPI ENDPOINTS  (thin wrappers around core_ functions)
# ═════════════════════════════════════════════════════════════

@fastapi_app.get("/health", response_model=HealthResponse)
async def health():
    return HealthResponse(
        status  = "healthy",
        model   = "nvidia/nemotron-3-super-120b-a12b:free + BGE-small",
        version = "2.0.0"
    )


@fastapi_app.post("/generate/from-text", response_model=QuizResponse)
async def api_generate_from_text(req: TextRequest):
    try:
        result = core_generate_from_text(
            text          = req.text,
            subject       = req.subject,
            chapter       = req.chapter,
            topic         = req.topic,
            exam_type     = req.exam_type,
            difficulty    = req.difficulty,
            num_questions = req.num_questions
        )
        return QuizResponse(**result)
    except ValueError as e:
        raise HTTPException(status_code=400, detail=str(e))
    except RuntimeError as e:
        raise HTTPException(status_code=500, detail=str(e))
    except Exception as e:
        raise HTTPException(
            status_code=500,
            detail=f"Generation failed: {str(e)}"
        )


@fastapi_app.post("/generate/from-pdf", response_model=QuizResponse)
async def api_generate_from_pdf(
    file         : UploadFile = File(...),
    subject      : str = "Biology",
    chapter      : str = "Chapter",
    topic        : str = "Main topic",
    exam_type    : str = "NEET",
    difficulty   : str = "Medium",
    num_questions: int = 10
):
    if not file.filename.endswith(".pdf"):
        raise HTTPException(
            status_code=400,
            detail="Only PDF files accepted"
        )

    with tempfile.NamedTemporaryFile(
        suffix=".pdf", delete=False
    ) as tmp:
        tmp.write(await file.read())
        tmp_path = tmp.name

    try:
        result = core_generate_from_pdf(
            pdf_path      = tmp_path,
            subject       = subject,
            chapter       = chapter,
            topic         = topic,
            exam_type     = exam_type,
            difficulty    = difficulty,
            num_questions = num_questions
        )
        return QuizResponse(**result)
    except Exception as e:
        raise HTTPException(
            status_code=500,
            detail=f"PDF processing failed: {str(e)}"
        )
    finally:
        os.unlink(tmp_path)


@fastapi_app.post("/score")
async def api_score_quiz(submission: dict):
    return core_score_quiz(
        answers     = submission.get("answers", {}),
        correct_map = submission.get("correct_answers", {}),
        exam_type   = submission.get("exam_type", "NEET")
    )
    
@fastapi_app.get("/test-models")
async def test_models():
    """Visit /test-models to check which free models work"""
    from openai import OpenAI

    api_key = os.getenv("OPENROUTER_API_KEY")
    client  = OpenAI(
        api_key  = api_key,
        base_url = "https://openrouter.ai/api/v1"
    )

    models_to_test = [
        "nvidia/nemotron-3-super-120b-a12b:free",
        "qwen/qwen3-next-80b-a3b-instruct:free",
        "google/gemma-4-31b-it:free",
        "openai/gpt-oss-120b:free",
    ]

    results = {}
    for model_name in models_to_test:
        try:
            res = client.chat.completions.create(
                model    = model_name,
                messages = [{"role": "user", "content": "Say OK"}],
                max_tokens = 5,
                extra_headers = {
                    "HTTP-Referer": "https://huggingface.co",
                    "X-Title"     : "NEET-JEE-Quiz-Generator"
                }
            )
            reply = res.choices[0].message.content.strip()
            results[model_name] = f"✅ Works → {reply}"
        except Exception as e:
            err = str(e)
            if "429" in err:
                results[model_name] = "⏳ Rate limited (try later)"
            elif "402" in err:
                results[model_name] = "💳 Requires credits"
            else:
                results[model_name] = f"❌ {err[:120]}"

    return results

# ═════════════════════════════════════════════════════════════
# GRADIO UI  (calls core_ functions directly — zero HTTP)
# ═════════════════════════════════════════════════════════════

def gradio_generate_text(
    subject, chapter, exam_type,
    difficulty, num_questions, notes_text
):
    """Gradio handler — calls core function directly, no requests lib."""
    if not notes_text or len(notes_text.strip()) < 100:
        return (
            "",
            "❌ Please paste at least a few paragraphs of notes",
            "", ""
        )
    try:
        result = core_generate_from_text(
            text          = notes_text,
            subject       = subject,
            chapter       = chapter,
            topic         = f"{subject} {chapter}",
            exam_type     = exam_type,
            difficulty    = difficulty,
            num_questions = int(num_questions)
        )
        return (
            json.dumps(result, indent=2),
            f" Generated {result['total_questions']} questions "
            f"in {result['generated_in_seconds']}s",
            result["session_id"],
            str(result["time_limit_minutes"])
        )
    except ValueError as e:
        return ("", f"❌ Input error: {e}", "", "")
    except Exception as e:
        return ("", f"❌ Error: {e}", "", "")


def gradio_generate_pdf(
    subject, chapter, exam_type,
    difficulty, num_questions, pdf_path
):
    """Gradio PDF handler — also direct call."""
    if pdf_path is None:
        return ("", "❌ Please upload a PDF file", "", "")
    try:
        result = core_generate_from_pdf(
            pdf_path      = pdf_path,
            subject       = subject,
            chapter       = chapter,
            topic = f"{subject} {chapter} key concepts, applications, exam-level MCQs",
            exam_type     = exam_type,
            difficulty    = difficulty,
            num_questions = int(num_questions)
        )
        return (
            json.dumps(result, indent=2),
            f" Generated {result['total_questions']} questions "
            f"in {result['generated_in_seconds']}s",
            result["session_id"],
            str(result["time_limit_minutes"])
        )
    except Exception as e:
        return ("", f"❌ Error: {e}", "", "")


def render_quiz(quiz_json_str: str) -> str:
    """Convert raw JSON output into readable markdown."""
    if not quiz_json_str or not quiz_json_str.strip():
        return ""
    try:
        data      = json.loads(quiz_json_str)
        quiz      = data.get("quiz", {})
        questions = quiz.get("questions", [])

        if not questions:
            return " No questions found in response"

        lines = []
        lines.append(f"#  {quiz.get('exam_type','Quiz')} Mock Test")
        lines.append(
            f"**Subject:** {quiz.get('subject','')} | "
            f"**Chapter:** {quiz.get('chapter','')} | "
            f"**Difficulty:** {quiz.get('difficulty','')}"
        )
        lines.append(
            f"**Questions:** {quiz.get('total_questions',0)} | "
            f"**Time:** {quiz.get('time_limit_minutes',0)} min | "
            f"**Marking:** +{questions[0].get('marks',4)} / "
            f"-{questions[0].get('negative_marks',1)}"
        )
        lines.append("\n---\n")

        for i, q in enumerate(questions, 1):
            opts = q.get("options", {})
            ans  = q.get("answer", "")
            lines.append(
                f"### Q{i}. `[{q.get('difficulty','')}]` "
                f"{q.get('question','')}"
            )
            for key in ["A", "B", "C", "D"]:
                marker = "✅" if key == ans else "⬜"
                lines.append(
                    f"{marker} **({key})** {opts.get(key,'')}"
                )
            lines.append(
                f"\n>  **Answer: {ans}** — "
                f"{q.get('explanation','')}"
            )
            lines.append("\n---\n")

        return "\n".join(lines)

    except json.JSONDecodeError:
        return " Could not parse quiz JSON"
    except Exception as e:
        return f" Render error: {str(e)}"


def build_gradio_blocks() -> gr.Blocks:
    with gr.Blocks(
        title="NEET/JEE Quiz Generator",
        theme=gr.themes.Soft()
    ) as demo:

        gr.Markdown("""
        # 📚 NEET / JEE Mock Quiz Generator
        Paste your study notes or upload a PDF → get exam-ready MCQs instantly
        """)

        with gr.Tabs():

            # ── Tab 1: Text input ────────────────────────────
            with gr.TabItem(" From Text Notes"):
                with gr.Row():
                    with gr.Column(scale=1, min_width=260):
                        gr.Markdown("###  Settings")
                        t_subject = gr.Dropdown(
                            choices=["Biology", "Physics",
                                     "Chemistry", "Mathematics"],
                            value="Biology", label="Subject"
                        )
                        t_chapter = gr.Textbox(
                            label="Chapter Name",
                            placeholder="e.g. Photosynthesis in Higher Plants"
                        )
                        t_exam = gr.Dropdown(
                            choices=["NEET", "JEE_MAINS", "JEE_ADVANCED"],
                            value="NEET", label="Exam Type"
                        )
                        t_diff = gr.Dropdown(
                            choices=["Easy", "Medium", "Hard", "Mixed"],
                            value="Medium", label="Difficulty"
                        )
                        t_num = gr.Slider(
                            minimum=1, maximum=30,
                            value=5, step=1,
                            label="Number of Questions"
                        )
                        t_btn = gr.Button(
                            " Generate Quiz",
                            variant="primary", size="lg"
                        )

                    with gr.Column(scale=2):
                        gr.Markdown("###  Paste Your Notes")
                        t_notes = gr.Textbox(
                            label="Study Notes",
                            placeholder="Paste your notes here...",
                            lines=18
                        )

                with gr.Row():
                    t_status  = gr.Textbox(
                        label="Status", interactive=False, scale=3
                    )
                    t_session = gr.Textbox(
                        label="Session ID", interactive=False, scale=1
                    )
                    t_time    = gr.Textbox(
                        label="Time Limit (min)", interactive=False, scale=1
                    )

                with gr.Tabs():
                    with gr.TabItem(" Rendered Quiz"):
                        t_rendered = gr.Markdown()
                    with gr.TabItem("🔧 Raw JSON"):
                        t_json = gr.Code(
                            language="json",
                            label="JSON Output",
                            lines=20
                        )

            # ── Tab 2: PDF input ─────────────────────────────
            with gr.TabItem(" From PDF"):
                with gr.Row():
                    with gr.Column(scale=1, min_width=260):
                        gr.Markdown("###  Settings")
                        p_subject = gr.Dropdown(
                            choices=["Biology", "Physics",
                                     "Chemistry", "Mathematics"],
                            value="Biology", label="Subject"
                        )
                        p_chapter = gr.Textbox(
                            label="Chapter Name",
                            placeholder="e.g. Laws of Motion"
                        )
                        p_exam = gr.Dropdown(
                            choices=["NEET", "JEE_MAINS", "JEE_ADVANCED"],
                            value="NEET", label="Exam Type"
                        )
                        p_diff = gr.Dropdown(
                            choices=["Easy", "Medium", "Hard", "Mixed"],
                            value="Medium", label="Difficulty"
                        )
                        p_num = gr.Slider(
                            minimum=1, maximum=30,
                            value=5, step=1,
                            label="Number of Questions"
                        )
                        p_btn = gr.Button(
                            "🚀 Generate from PDF",
                            variant="primary", size="lg"
                        )

                    with gr.Column(scale=2):
                        gr.Markdown("###  Upload PDF")
                        p_file = gr.File(
                            label="Upload PDF Notes",
                            file_types=[".pdf"],
                            type="filepath"   # gives us the path string
                        )
                        gr.Markdown("""
                        **Tips:** Single chapter PDFs work best.
                        Text-based PDFs only (not scanned images).
                        """)

                with gr.Row():
                    p_status  = gr.Textbox(
                        label="Status", interactive=False, scale=3
                    )
                    p_session = gr.Textbox(
                        label="Session ID", interactive=False, scale=1
                    )
                    p_time    = gr.Textbox(
                        label="Time Limit (min)", interactive=False, scale=1
                    )

                with gr.Tabs():
                    with gr.TabItem(" Rendered Quiz"):
                        p_rendered = gr.Markdown()
                    with gr.TabItem(" Raw JSON"):
                        p_json = gr.Code(
                            language="json",
                            label="JSON Output",
                            lines=20
                        )

            # ── Tab 3: API reference ─────────────────────────
            with gr.TabItem(" API Reference"):
                gr.Markdown("""
                ## Using the REST API from your webapp

                ### Base URL
                ```
                https://YOUR-USERNAME-YOUR-SPACE-NAME.hf.space
                ```

                ### Health check
                ```
                GET /health
                ```

                ### Generate from text
                ```
                POST /generate/from-text
                Content-Type: application/json

                {
                  "text": "your notes...",
                  "subject": "Biology",
                  "chapter": "Photosynthesis",
                  "topic": "Light reactions Calvin cycle",
                  "exam_type": "NEET",
                  "difficulty": "Medium",
                  "num_questions": 10
                }
                ```

                ### Generate from PDF
                ```
                POST /generate/from-pdf
                Content-Type: multipart/form-data

                file: <pdf>
                subject: Biology
                chapter: Photosynthesis
                exam_type: NEET
                difficulty: Medium
                num_questions: 10
                ```

                ### Interactive Swagger docs
                ```
                GET /docs
                ```
                """)

        gr.Markdown(
            "---\nBuilt with 🤗 Gradio + FastAPI + Gemini + FAISS"
        )

        # ── Wire up buttons ──────────────────────────────────
        t_btn.click(
            fn=gradio_generate_text,
            inputs=[t_subject, t_chapter, t_exam,
                    t_diff, t_num, t_notes],
            outputs=[t_json, t_status, t_session, t_time]
        ).then(
            fn=render_quiz,
            inputs=[t_json],
            outputs=[t_rendered]
        )

        p_btn.click(
            fn=gradio_generate_pdf,
            inputs=[p_subject, p_chapter, p_exam,
                    p_diff, p_num, p_file],
            outputs=[p_json, p_status, p_session, p_time]
        ).then(
            fn=render_quiz,
            inputs=[p_json],
            outputs=[p_rendered]
        )

    return demo


# ═════════════════════════════════════════════════════════════
# MOUNT GRADIO INTO FASTAPI  — single process, single port
# ═════════════════════════════════════════════════════════════
gradio_blocks = build_gradio_blocks()

app = gr.mount_gradio_app(
    app    = fastapi_app,
    blocks = gradio_blocks,
    path   = "/"
)