main.py

from __future__ import annotations

import sys
from pathlib import Path

# Make the prototype directory importable regardless of cwd
sys.path.insert(0, str(Path(__file__).parent))

from fastapi import FastAPI, HTTPException, Request
from fastapi.responses import FileResponse, HTMLResponse, RedirectResponse
from fastapi.staticfiles import StaticFiles
from fastapi.templating import Jinja2Templates
from pydantic import BaseModel

from config import PAPERS, TOP_N_RESULTS
from paper_loader import load_paper, get_paper_title
from renderer import render_paper
from retrieval import active_scorer

app = FastAPI(title="RAG Prototype")

_here = Path(__file__).parent
app.mount("/static", StaticFiles(directory=_here / "static"), name="static")
templates = Jinja2Templates(directory=_here / "templates")


# ---------------------------------------------------------------------------
# Warm up: pre-load paper data on startup so first request is fast
# ---------------------------------------------------------------------------

@app.on_event("startup")
async def startup():
    for key in PAPERS:
        load_paper(key)


# ---------------------------------------------------------------------------
# Routes
# ---------------------------------------------------------------------------

@app.get("/", response_class=RedirectResponse)
async def index():
    first_key = next(iter(PAPERS))
    return RedirectResponse(f"/papers/{first_key}")


@app.get("/papers/{paper_key}", response_class=HTMLResponse)
async def paper_view(request: Request, paper_key: str):
    if paper_key not in PAPERS:
        raise HTTPException(status_code=404, detail="Paper not found")

    sections, chunks, block_to_chunk = load_paper(paper_key)
    paper_html = render_paper(paper_key, sections, block_to_chunk)
    title = get_paper_title(paper_key)

    chunk_index = {
        c.chunk_id: {
            "chunk_id": c.chunk_id,
            "section_title": c.section_title,
            "section_path": c.section_path,
            "page_span": c.page_span,
            "chunk_type": c.chunk_type,
            "source_block_ids": c.source_block_ids,
            "text": c.text,
        }
        for c in chunks
    }

    return templates.TemplateResponse(request, "paper.html", {
        "paper_key": paper_key,
        "paper_title": title,
        "paper_html": paper_html,
        "chunk_count": len(chunks),
        "top_n": TOP_N_RESULTS,
        "chunk_index": chunk_index,
    })


# ---------------------------------------------------------------------------
# Scoring API
# ---------------------------------------------------------------------------

class ScoreRequest(BaseModel):
    query: str
    top_n: int = TOP_N_RESULTS


@app.post("/api/{paper_key}/score")
async def score(paper_key: str, req: ScoreRequest):
    if paper_key not in PAPERS:
        raise HTTPException(status_code=404, detail="Paper not found")
    if not req.query.strip():
        raise HTTPException(status_code=422, detail="Query must not be empty")

    _, chunks, _ = load_paper(paper_key)
    results = active_scorer.score(req.query, chunks)
    top = results[: req.top_n]

    return {"results": [r.to_dict() for r in top], "total_chunks": len(chunks)}


# ---------------------------------------------------------------------------
# Image serving
# ---------------------------------------------------------------------------

@app.get("/images/{paper_key}/{filename}")
async def serve_image(paper_key: str, filename: str):
    if paper_key not in PAPERS:
        raise HTTPException(status_code=404, detail="Paper not found")
    images_dir = PAPERS[paper_key]["images_dir"] / "images"
    img_path = images_dir / filename
    if not img_path.exists():
        raise HTTPException(status_code=404, detail="Image not found")
    return FileResponse(img_path)