src/api_routes.py

"""API routes for document intelligence endpoints.

Includes file upload, figure navigation, and helper function endpoints.
All endpoints work with session-based state management.
"""

import base64
import os
import uuid
from io import BytesIO
from pathlib import Path
from typing import Any

from fastapi import File, HTTPException, Query, UploadFile
from fastapi.responses import JSONResponse, StreamingResponse

from api_helpers import _stream_with_progress
from storage import get_temp_dir, load_parse_cache, save_image, save_parse_cache, save_upload, use_disk_images


def _image_to_b64(img: "Path | Image.Image | None") -> "str | None":
    """Encode a PIL Image or Path to a base64 PNG string for JSON responses."""
    if img is None:
        return None
    if isinstance(img, Path):
        return base64.b64encode(img.read_bytes()).decode("utf-8")
    buf = BytesIO()
    img.save(buf, format="PNG")
    buf.seek(0)
    return base64.b64encode(buf.getvalue()).decode("utf-8")


def create_document_routes(
    app: Any,
    session_states: dict[str, dict[str, Any]],
    process_upload: Any,
    next_figure: Any,
    prev_figure: Any,
) -> None:
    """Register document processing API routes on the FastAPI app.

    Args:
        app: FastAPI application instance.
        session_states: Global session states dictionary.
        process_upload: Function to process uploaded files.
        next_figure: Function to get next figure.
        prev_figure: Function to get previous figure.
    """

    @app.post("/api/upload")
    async def api_upload_file(file: UploadFile = File(...)) -> StreamingResponse:  # noqa: ANN001
        """API endpoint for file upload from Next.js frontend.

        Args:
            file: Uploaded file from form-data.

        Returns:
            SSE stream with progress events and final JSON result.
        """
        from collections.abc import Callable

        from PIL import Image

        from crops import extract_figures
        from document_parser import parse_document
        from pdf_io import load_pdf_pages
        from ui_state import create_initial_state, hash_bytes, page_cache, parse_cache

        file_bytes = await file.read()
        session_id = str(uuid.uuid4())
        session_states[session_id] = create_initial_state()

        suffix = Path(file.filename).suffix if file.filename else ".tmp"
        upload_path = save_upload(session_id, file_bytes, suffix)
        if upload_path is not None:
            temp_path = str(upload_path)
            _cleanup_temp = False
        else:
            temp_path = str(get_temp_dir() / f"{uuid.uuid4().hex}{suffix}")
            Path(temp_path).write_bytes(file_bytes)
            _cleanup_temp = True

        image_exts = {".jpg", ".jpeg", ".jfif", ".png", ".bmp", ".dib", ".gif", ".tif", ".tiff", ".webp"}
        office_exts = {".docx", ".xlsx", ".pptx"}
        max_pages = 20

        def _work(on_progress: Callable[[str], None]) -> dict:
            try:
                ext = Path(temp_path).suffix.lower()
                state = session_states[session_id]
                state["current_figure_index"] = 0
                state["conversation_history"] = []
                state["current_image_path"] = None

                with open(temp_path, "rb") as f:
                    raw = f.read()

                file_hash = hash_bytes(raw)
                state["uploaded_file_hash"] = file_hash
                if ext in image_exts:
                    on_progress("Loading image...")
                    image = Image.open(temp_path).convert("RGB")
                    lazy = save_image(session_id, "figures", 0, image)  # LazyImage or PIL Image
                    figures_info = [{"image": lazy, "page": 0, "bbox": None, "caption": ""}]
                    state["page_images"] = [lazy]  # LazyImage proxies in disk mode, PIL Images in memory mode
                    state["figures_info"] = figures_info  # fig["image"] is LazyImage or PIL Image
                    state["selected_figure"] = figures_info[0]  # reference to figures_info entry
                    session_states[session_id] = state

                    return {
                        "status": "Image loaded successfully.\nNumber of figures: 1.",
                        "html_content": "Image uploaded directly (no document parsing needed)",
                        "fig_status": "Figure 1 of 1 (Page 1)",
                        "fig_caption": "",
                        "fig_image": _image_to_b64(image),
                        "session_id": session_id,
                    }

                fmt_label = ext.lstrip(".").upper()
                status_lines = [f"{fmt_label} loaded successfully."]

                if ext in office_exts:
                    page_images: list = []
                    state["page_images"] = []
                else:
                    on_progress("Rendering PDF pages...")
                    cache_key = f"{file_hash}_{max_pages}"
                    if cache_key in page_cache:
                        page_images = page_cache[cache_key]
                    else:
                        page_images = load_pdf_pages(raw, max_pages=max_pages)
                        if not use_disk_images():
                            page_cache[cache_key] = page_images
                    state["page_images"] = [  # LazyImage proxies in disk mode, PIL Images in memory mode
                        save_image(session_id, "pages", i, img)
                        for i, img in enumerate(page_images)
                    ]
                    status_lines.append(f"Number of pages rendered: {len(page_images)} (max {max_pages}).")

                on_progress("Parsing document with Docling...")
                if not use_disk_images() and file_hash in parse_cache:
                    parse_result = parse_cache[file_hash]
                else:
                    parse_result = load_parse_cache(file_hash, session_id=session_id)
                    if parse_result is None:
                        parse_result = parse_document(raw, file_ext=ext, on_progress=on_progress)
                        save_parse_cache(file_hash, parse_result, session_id=session_id)
                    if not use_disk_images():
                        parse_cache[file_hash] = parse_result
                status_lines.append("Document parsing done using Docling.")

                on_progress("Extracting figures...")
                figures_info = extract_figures(page_images, parse_result.get("figures", []))
                for i, fig in enumerate(figures_info):
                    fig["image"] = save_image(session_id, "figures", i, fig["image"])  # LazyImage or PIL Image
                state["figures_info"] = figures_info  # fig["image"] is LazyImage or PIL Image
                status_lines.append(f"Number of figures extracted: {len(figures_info)}.")

                if figures_info:
                    state["selected_figure"] = figures_info[0]  # reference to figures_info entry
                    fig_status = f"Figure 1 of {len(figures_info)} (Page {figures_info[0]['page'] + 1})"
                    fig_caption = figures_info[0].get("caption", "No caption")
                    fig_image = figures_info[0]["image"]
                else:
                    state["selected_figure"] = None
                    fig_status = "No figures found"
                    fig_caption = ""
                    fig_image = None

                on_progress("Finalizing...")
                session_states[session_id] = state

                fig_image_data = _image_to_b64(fig_image)

                return {
                    "status": "\n".join(status_lines),
                    "html_content": parse_result.get("html", "No content available"),
                    "fig_status": fig_status,
                    "fig_caption": fig_caption,
                    "fig_image": fig_image_data,
                    "session_id": session_id,
                }
            finally:
                if _cleanup_temp:
                    os.remove(temp_path)

        return StreamingResponse(
            _stream_with_progress(_work, session_id),
            media_type="text/event-stream",
        )

    @app.get("/health")
    async def health_check() -> JSONResponse:  # noqa: ANN001
        """Health check endpoint to verify API is accessible.

        Returns:
            JSON response with status.
        """
        return JSONResponse({"status": "ok"})

    @app.post("/api/next_figure")
    async def api_next_figure(session_id: str = Query(...)) -> JSONResponse:  # noqa: ANN001
        """Get the next figure for a session.

        Args:
            session_id: The session ID to advance figures for.

        Returns:
            JSON response with figure info, caption, and image data.
        """
        if session_id not in session_states:
            raise HTTPException(status_code=404, detail="Session not found")

        try:
            fig_status, fig_caption, fig_image, updated_state = next_figure(session_states[session_id])
            session_states[session_id] = updated_state

            return JSONResponse(
                {
                    "fig_status": fig_status,
                    "fig_caption": fig_caption,
                    "fig_image": _image_to_b64(fig_image),
                    "fig_index": session_states[session_id].get("current_figure_index", 0),
                }
            )
        except Exception as e:  # noqa: BLE001
            import traceback

            traceback.print_exc()
            raise HTTPException(status_code=400, detail=str(e)) from e

    @app.post("/api/prev_figure")
    async def api_prev_figure(session_id: str = Query(...)) -> JSONResponse:  # noqa: ANN001
        """Get the previous figure for a session.

        Args:
            session_id: The session ID to go back figures for.

        Returns:
            JSON response with figure info, caption, and image data.
        """
        if session_id not in session_states:
            raise HTTPException(status_code=404, detail="Session not found")

        try:
            fig_status, fig_caption, fig_image, updated_state = prev_figure(session_states[session_id])
            session_states[session_id] = updated_state

            return JSONResponse(
                {
                    "fig_status": fig_status,
                    "fig_caption": fig_caption,
                    "fig_image": _image_to_b64(fig_image),
                    "fig_index": session_states[session_id].get("current_figure_index", 0),
                }
            )
        except Exception as e:  # noqa: BLE001
            import traceback

            traceback.print_exc()
            raise HTTPException(status_code=400, detail=str(e)) from e