app.py

import re
import uuid
import asyncio
import logging
import tempfile
import os

from fastapi import FastAPI, UploadFile, File, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from typing import Optional

from src.chatbot.engine import call_llm
from src.chatbot.prompts import AI_MODES, DEFAULT_PROMPT
from src.chatbot.document_hub import process_document, list_documents, delete_document
from src.chatbot.rag_pipeline import index_document, retrieve_context, build_rag_prompt

# ---------------------------------------------------------------------------
# Logging
# ---------------------------------------------------------------------------
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger("ylf-api")

# ---------------------------------------------------------------------------
# App
# ---------------------------------------------------------------------------
app = FastAPI(
    title="YLF AI Platform",
    description="AI Tutoring API — RAG, multiple explanation modes, exam generation.",
    version="1.0.0"
)

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)

VALID_MODES = list(AI_MODES.keys())

# ---------------------------------------------------------------------------
# Health check
# ---------------------------------------------------------------------------

@app.get("/")
def home():
    return {"message": "YLF API is running 🚀"}


# ---------------------------------------------------------------------------
# Document endpoints
# ---------------------------------------------------------------------------

@app.post("/upload")
async def upload_pdf(file: UploadFile = File(...)):
    """
    Upload a PDF → extract → chunk (with page metadata) → index for RAG.
    Returns doc_id to scope /chat requests to this document.
    """
    if not file.filename.endswith(".pdf"):
        raise HTTPException(status_code=400, detail="Only PDF files are supported.")

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

    try:
        result = process_document(tmp_path, filename=file.filename)
        if result["status"] == "error":
            raise HTTPException(status_code=422, detail=result["message"])

        index_result = index_document(result["doc_id"])
        logger.info(
            f"[Upload] {file.filename} → doc_id={result['doc_id']} "
            f"| {index_result['chunks_indexed']} chunks indexed"
        )

        return {
            "doc_id": result["doc_id"],
            "filename": result["filename"],
            "num_chunks": result["num_chunks"],
        }
    finally:
        os.unlink(tmp_path)


@app.get("/documents")
def get_documents():
    """List all currently indexed documents."""
    return {"documents": list_documents()}


@app.delete("/documents/{doc_id}")
def remove_document(doc_id: str):
    """Remove a document from the index."""
    removed = delete_document(doc_id)
    if not removed:
        raise HTTPException(status_code=404, detail=f"Document '{doc_id}' not found.")
    return {"status": "deleted", "doc_id": doc_id}


# ---------------------------------------------------------------------------
# Chat endpoint
# ---------------------------------------------------------------------------

class ChatRequest(BaseModel):
    message: str
    mode: str = "socratic"
    session_id: Optional[str] = None
    doc_id: Optional[str] = None    # restrict RAG to one document


@app.post("/chat")
async def chat_endpoint(request: ChatRequest):
    """
    Main chat endpoint with smart routing, fallback, and RAG support.

    Supported modes: socratic, explain_beginner, explain_intermediate,
                     explain_expert, question_gen, chunking_helper

    Response shape (Day 5 spec):
    {
        "answer":     str,
        "questions":  list[str],   # populated when mode == "question_gen"
        "source":     str | null,  # e.g. "page 4" from RAG retrieval
        "session_id": str,
        "mode":       str,
        "status":     str          # "success" | "temporary_failure"
    }
    """
    mode = request.mode.lower().strip()
    if mode not in VALID_MODES:
        raise HTTPException(
            status_code=400,
            detail=f"Invalid mode '{mode}'. Valid modes: {VALID_MODES}"
        )

    # Consistent session ID format with engine logs
    session_id = request.session_id or f"user_{uuid.uuid4().hex[:8]}"
    logger.info(f"Session: {session_id} | Mode: {mode} | Msg: {request.message[:80]}")

    # ------------------------------------------------------------------
    # RAG: retrieve context if any documents are indexed
    # ------------------------------------------------------------------
    source = None
    rag_context_str = None

    if list_documents():
        rag_result = retrieve_context(request.message, doc_id=request.doc_id)
        source = rag_result["source"]
        rag_context_str = rag_result["context_str"]

    # ------------------------------------------------------------------
    # Build system prompt — inject RAG context if available
    # ------------------------------------------------------------------
    base_prompt = AI_MODES.get(mode, DEFAULT_PROMPT)

    if rag_context_str and "No relevant context" not in rag_context_str:
        augmented_prompt = build_rag_prompt(request.message, rag_context_str, base_prompt)
        _rag_key = f"__rag_{session_id}"
        AI_MODES[_rag_key] = augmented_prompt
        try:
            # Use to_thread — engine performs blocking HTTP requests
            answer = await asyncio.to_thread(
                call_llm,
                user_query=request.message,
                mode=_rag_key,
                session_id=session_id
            )
        finally:
            AI_MODES.pop(_rag_key, None)    # always clean up
    else:
        answer = await asyncio.to_thread(
            call_llm,
            user_query=request.message,
            mode=mode,
            session_id=session_id
        )

    # ------------------------------------------------------------------
    # Handle total engine failure (all models + retries exhausted)
    # ------------------------------------------------------------------
    if "All models failed" in answer:
        logger.error(f"Critical Engine Failure: {answer}")
        return {
            "answer": "Sorry, our AI engines are currently under heavy load. Please try again in a minute.",
            "questions": [],
            "source": source,
            "session_id": session_id,
            "mode": mode,
            "status": "temporary_failure"
        }

    # ------------------------------------------------------------------
    # Parse questions list when mode is question_gen
    # ------------------------------------------------------------------
    questions = []
    if mode == "question_gen":
        questions = [
            line.strip()
            for line in answer.split("\n")
            if re.match(r'^[\d\-\*\•]', line.strip()) and len(line.strip()) > 5
        ]

    return {
        "answer": answer,
        "questions": questions,
        "source": source,
        "session_id": session_id,
        "mode": mode,
        "status": "success"
    }