main.py

import json
import shutil
import tempfile
import time
import logging
from contextlib import asynccontextmanager
from pathlib import Path

from dotenv import load_dotenv
from fastapi import FastAPI, File, HTTPException, UploadFile
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import StreamingResponse
from pydantic import BaseModel, Field

from ingestion.embedder import Embedder
from ingestion.pipeline import IngestionPipeline
from retrieval.index import VectorIndex
from retrieval.searcher import search
from generation.generator import Generator

load_dotenv()

logging.basicConfig(level=logging.INFO, format="%(levelname)s | %(name)s | %(message)s")
logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Application state (populated in lifespan, shared across requests)
# ---------------------------------------------------------------------------

class AppState:
    embedder: Embedder
    index: VectorIndex
    pipeline: IngestionPipeline
    generator: Generator

state = AppState()


@asynccontextmanager
async def lifespan(app: FastAPI):
    logger.info("Loading embedder model...")
    state.embedder = Embedder()

    logger.info("Initialising FAISS index (dim=%d)...", state.embedder.dimension)
    state.index = VectorIndex(dimension=state.embedder.dimension)

    logger.info("Building ingestion pipeline...")
    state.pipeline = IngestionPipeline(
        embedder=state.embedder,
        index=state.index,
        strategy="recursive_character",
        chunk_size=500,
        overlap=50,
    )

    logger.info("Initialising Gemini generator...")
    state.generator = Generator()

    logger.info("Startup complete — ready to serve.")
    yield
    logger.info("Shutting down.")


# ---------------------------------------------------------------------------
# App
# ---------------------------------------------------------------------------

app = FastAPI(
    title="RAG Document Q&A",
    version="1.0.0",
    description="Upload PDFs, ask questions, get grounded answers with citations.",
    lifespan=lifespan,
)

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


# ---------------------------------------------------------------------------
# Request / response schemas
# ---------------------------------------------------------------------------

class QueryRequest(BaseModel):
    question: str = Field(..., min_length=1)
    top_k: int = Field(default=5, ge=1, le=20)


class SourceInfo(BaseModel):
    chunk_id: str
    source: str
    page_num: int
    score: float


class QueryResponse(BaseModel):
    question: str
    answer: str
    sources: list[SourceInfo]
    duration_ms: float
    confidence_score: float
    confidence_level: str  # "high" | "medium" | "low"


class IngestResponse(BaseModel):
    file: str
    pages: int
    chunks: int
    chunk_ids: list[str]
    duration_ms: float


class StatsResponse(BaseModel):
    index_size: int
    total_chunks_ingested: int
    embedding_model: str
    embedding_dimension: int


class HealthResponse(BaseModel):
    status: str


# ---------------------------------------------------------------------------
# Endpoints
# ---------------------------------------------------------------------------

@app.get("/health", response_model=HealthResponse, tags=["system"])
def health():
    return HealthResponse(status="ok")


@app.get("/stats", response_model=StatsResponse, tags=["system"])
def stats():
    return StatsResponse(
        index_size=state.index.size,
        total_chunks_ingested=state.pipeline.chunk_count,
        embedding_model="all-MiniLM-L6-v2",
        embedding_dimension=state.embedder.dimension,
    )


@app.post("/ingest", response_model=IngestResponse, tags=["ingestion"])
def ingest(file: UploadFile = File(...)):
    """Upload a PDF and add its content to the vector index.

    The file is written to a temp path, processed by the ingestion pipeline
    (extract → chunk → embed → index), then deleted.  Returns the number of
    chunks added and wall-clock timing.
    """
    if not (file.filename or "").lower().endswith(".pdf"):
        raise HTTPException(status_code=400, detail="Only PDF files are accepted.")

    t0 = time.perf_counter()

    tmp_path: Path | None = None
    try:
        tmp_dir = Path(tempfile.mkdtemp())
        tmp_path = tmp_dir / file.filename
        with tmp_path.open("wb") as f:
            shutil.copyfileobj(file.file, f)

        result = state.pipeline.ingest_pdf(tmp_path)
    finally:
        file.file.close()
        if tmp_path and tmp_path.exists():
            shutil.rmtree(tmp_path.parent, ignore_errors=True)

    if result.error:
        raise HTTPException(status_code=422, detail=result.error)

    duration_ms = round((time.perf_counter() - t0) * 1000, 2)

    return IngestResponse(
        file=result.file,
        pages=result.pages,
        chunks=result.chunks,
        chunk_ids=result.chunk_ids,
        duration_ms=duration_ms,
    )


@app.post("/query", response_model=QueryResponse, tags=["query"])
def query(request: QueryRequest):
    """Ask a question against the indexed documents.

    Embeds the question, retrieves the top-k matching chunks from FAISS,
    and sends them to Gemini 1.5 Flash with a grounding prompt.  The model
    is instructed to cite sources inline using [Source N] notation.
    """
    if state.index.size == 0:
        raise HTTPException(
            status_code=400,
            detail="The index is empty. Upload at least one PDF via POST /ingest first.",
        )

    t0 = time.perf_counter()

    search_resp = search(request.question, state.embedder, state.index, k=request.top_k)
    answer = state.generator.generate_answer(
        request.question, search_resp.chunks, max_score=search_resp.max_score
    )

    duration_ms = round((time.perf_counter() - t0) * 1000, 2)

    sources = [
        SourceInfo(
            chunk_id=r.metadata.get("chunk_id", ""),
            source=r.metadata.get("source", ""),
            page_num=r.metadata.get("page_num", 0),
            score=round(r.score, 4),
        )
        for r in search_resp.chunks
    ]

    return QueryResponse(
        question=answer.question,
        answer=answer.answer,
        sources=sources,
        duration_ms=duration_ms,
        confidence_score=round(search_resp.max_score, 4),
        confidence_level=answer.confidence_level,
    )


@app.post("/query/stream", tags=["query"])
def query_stream(request: QueryRequest):
    """Stream an answer as Server-Sent Events (text/event-stream).

    Each SSE event carries a JSON payload: {"text": "<chunk>"}.
    The final event is {"done": true}.  On error, {"error": "<message>"} is
    sent and the stream closes.

    The existing POST /query endpoint is unaffected.
    """
    if state.index.size == 0:
        raise HTTPException(
            status_code=400,
            detail="The index is empty. Upload at least one PDF via POST /ingest first.",
        )

    search_resp = search(request.question, state.embedder, state.index, k=request.top_k)

    def event_generator():
        try:
            for chunk in state.generator.generate_answer_stream(
                request.question, search_resp.chunks, max_score=search_resp.max_score
            ):
                yield f"data: {json.dumps({'text': chunk})}\n\n"
        except Exception as exc:
            logger.error("Streaming generation error: %s", exc)
            yield f"data: {json.dumps({'error': str(exc)})}\n\n"
        yield f"data: {json.dumps({'done': True})}\n\n"

    return StreamingResponse(event_generator(), media_type="text/event-stream")


# ---------------------------------------------------------------------------
# Debug endpoints
# ---------------------------------------------------------------------------

@app.post("/debug/chunks", tags=["debug"])
def debug_chunks(request: QueryRequest):
    """Show retrieved chunks without generating an answer. For debugging."""
    if state.index.size == 0:
        raise HTTPException(status_code=400, detail="Index is empty.")

    search_resp = search(request.question, state.embedder, state.index, k=request.top_k)
    return {
        "question": request.question,
        "max_score": round(search_resp.max_score, 4),
        "expansion_used": search_resp.expansion_used,
        "chunks": [
            {
                "rank": i,
                "score": round(r.score, 4),
                "source": r.metadata.get("source", "?"),
                "page": r.metadata.get("page_num", "?"),
                "section": r.metadata.get("section_header", None),
                "text_preview": r.metadata.get("text", "")[:300],
            }
            for i, r in enumerate(search_resp.chunks, 1)
        ],
    }


# ---------------------------------------------------------------------------
# Dev entry point
# ---------------------------------------------------------------------------

if __name__ == "__main__":
    import uvicorn
    uvicorn.run("main:app", host="0.0.0.0", port=8000, reload=True)