app.py

from __future__ import annotations

import json
from pathlib import Path

from fastapi import BackgroundTasks, FastAPI, File, HTTPException, Request, UploadFile
from fastapi.responses import HTMLResponse, StreamingResponse
from fastapi.staticfiles import StaticFiles
from fastapi.templating import Jinja2Templates
from pydantic import BaseModel

from pdf_to_text import ExtractionResult, extract_uploaded_pdf, stream_uploaded_pdf
from rag_pipeline import SmartNotesRAG


BASE_DIR = Path(__file__).resolve().parent
PROJECT_METADATA = {
    "name": "SmartNotes AI",
    "title": "SmartNotes PDF OCR",
    "version": "1.0.0",
    "description": "FastAPI app for PDF OCR, text extraction, and RAG-based question answering.",
    "runtime": "Python 3.12",
    "framework": "FastAPI",
    "default_port": 8000,
    "entrypoint": "app.py",
    "storage": "SQLite local fallback with optional PostgreSQL",
    "vector_store": "Local fallback with optional Qdrant",
    "features": [
        "PDF direct text extraction",
        "OCR fallback with OpenCV, Gemini Vision OCR, and TrOCR",
        "RAG indexing with parent and child chunks",
        "Hybrid retrieval with embeddings and BM25",
        "Reranked answers with citations",
    ],
}

app = FastAPI(
    title=PROJECT_METADATA["title"],
    description=PROJECT_METADATA["description"],
    version=PROJECT_METADATA["version"],
)
app.mount("/static", StaticFiles(directory=BASE_DIR / "static"), name="static")
templates = Jinja2Templates(directory=BASE_DIR / "templates")
rag = SmartNotesRAG()


class IndexTextRequest(BaseModel):
    file_name: str
    text: str


class QueryRequest(BaseModel):
    document_id: str
    query: str


class FeedbackRequest(BaseModel):
    document_id: str
    query: str
    rating: int
    comment: str = ""


@app.get("/", response_class=HTMLResponse)
async def home(request: Request):
    return templates.TemplateResponse(request, "index.html")


@app.get("/api/metadata")
async def project_metadata():
    return PROJECT_METADATA


@app.post("/api/pdf-to-text")
async def pdf_to_text(file: UploadFile = File(...)):
    if not _looks_like_pdf(file.filename, file.content_type):
        raise HTTPException(status_code=400, detail="Please upload a PDF file.")

    file_bytes = await file.read()
    if not file_bytes:
        raise HTTPException(status_code=400, detail="Uploaded PDF is empty.")

    try:
        result = extract_uploaded_pdf(file.filename, file_bytes)
    except RuntimeError as exc:
        raise HTTPException(status_code=500, detail=str(exc)) from exc
    except Exception as exc:
        raise HTTPException(status_code=500, detail=f"PDF processing failed: {exc}") from exc

    return _result_payload(result, original_file=file.filename)


@app.post("/api/pdf-to-text-stream")
async def pdf_to_text_stream(file: UploadFile = File(...)):
    if not _looks_like_pdf(file.filename, file.content_type):
        raise HTTPException(status_code=400, detail="Please upload a PDF file.")

    original_file = file.filename
    file_bytes = await file.read()
    if not file_bytes:
        raise HTTPException(status_code=400, detail="Uploaded PDF is empty.")

    def events():
        try:
            for event in stream_uploaded_pdf(original_file, file_bytes):
                yield json.dumps(event, ensure_ascii=True) + "\n"
        except Exception as exc:
            yield json.dumps(
                {
                    "type": "error",
                    "message": f"PDF processing failed: {exc}",
                },
                ensure_ascii=True,
            ) + "\n"

    return StreamingResponse(events(), media_type="application/x-ndjson")


@app.post("/api/index-text-stream")
async def index_text_stream(payload: IndexTextRequest):
    if not payload.text.strip():
        raise HTTPException(status_code=400, detail="Extracted text is empty.")

    def events():
        try:
            for event in rag.index_text_stream(payload.file_name, payload.text):
                yield json.dumps(event, ensure_ascii=True) + "\n"
        except Exception as exc:
            yield json.dumps(
                {
                    "type": "error",
                    "message": f"Indexing failed: {exc}",
                },
                ensure_ascii=True,
            ) + "\n"

    return StreamingResponse(events(), media_type="application/x-ndjson")


@app.post("/api/index-text-background")
async def index_text_background(payload: IndexTextRequest, background_tasks: BackgroundTasks):
    if not payload.text.strip():
        raise HTTPException(status_code=400, detail="Extracted text is empty.")

    def run_indexing() -> None:
        for _ in rag.index_text_stream(payload.file_name, payload.text):
            pass

    background_tasks.add_task(run_indexing)
    return {"status": "queued"}


@app.post("/api/query")
async def query_document(payload: QueryRequest):
    return rag.query(payload.document_id, payload.query)


@app.post("/api/feedback")
async def save_feedback(payload: FeedbackRequest):
    rating = max(1, min(payload.rating, 5))
    return rag.save_feedback(payload.document_id, payload.query, rating, payload.comment)


@app.get("/api/documents/{document_id}")
async def get_document(document_id: str):
    document = rag.repository.get_document(document_id)
    if not document:
        raise HTTPException(status_code=404, detail="Document not found.")
    document.pop("cleaned_text", None)
    return document


@app.get("/health")
async def health():
    return {"status": "ok"}


def _result_payload(result: ExtractionResult, original_file: str) -> dict:
    return {
        "file": original_file,
        "text": result.text,
        "route": result.route,
        "page_count": result.page_count,
        "direct_text_found": result.direct_text_found,
        "warnings": result.warnings,
        "pages": [
            {
                "page_number": page.page_number,
                "engine": page.engine,
                "confidence": page.confidence,
                "text_length": len(page.text),
            }
            for page in result.pages
        ],
    }


def _looks_like_pdf(file_name: str | None, content_type: str | None) -> bool:
    if content_type in {"application/pdf", "application/x-pdf"}:
        return True
    if file_name and Path(file_name).suffix.lower() == ".pdf":
        return True
    return False