src/api_routes_v1.py

"""V1 API routes — browser-memory storage mode.

Upload returns all figure images inline as base64; backend retains nothing
after the response is sent. No next_figure/prev_figure endpoints needed.
"""

import base64
import os
import uuid
from collections.abc import Callable
from io import BytesIO
from pathlib import Path
from typing import Any

from fastapi import File, UploadFile
from fastapi.responses import StreamingResponse

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


def _image_to_b64(img: "Path | Any | None") -> "str | None":
    """Encode a PIL Image or Path to a base64 PNG string."""
    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_v1(
    app: Any,
    session_states: dict[str, dict[str, Any]],
) -> None:
    """Register v1 document processing API routes."""

    @app.post("/api/v1/upload")
    async def api_v1_upload_file(file: UploadFile = File(...)) -> StreamingResponse:
        """V1 upload: returns all figures inline, then cleans up."""
        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")
                    figures = [
                        {
                            "image": _image_to_b64(image),
                            "page": 0,
                            "caption": "",
                            "status": "Figure 1 of 1 (Page 1)",
                        }
                    ]

                    return {
                        "status": "Image loaded successfully.\nNumber of figures: 1.",
                        "html_content": "Image uploaded directly (no document parsing needed)",
                        "figures": figures,
                        "session_id": session_id,
                        "v1": True,
                    }

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

                if ext in office_exts:
                    page_images: list = []
                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
                    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", []))
                status_lines.append(f"Number of figures extracted: {len(figures_info)}.")

                on_progress("Finalizing...")
                figures = []
                for i, fig in enumerate(figures_info):
                    fig_b64 = _image_to_b64(fig["image"])
                    page_num = fig.get("page", 0)
                    caption = fig.get("caption", "")
                    fig_status = f"Figure {i + 1} of {len(figures_info)} (Page {page_num + 1})"
                    figures.append({
                        "image": fig_b64,
                        "page": page_num,
                        "caption": caption,
                        "status": fig_status,
                    })

                return {
                    "status": "\n".join(status_lines),
                    "html_content": parse_result.get("html", "No content available"),
                    "figures": figures,
                    "session_id": session_id,
                    "v1": True,
                }
            finally:
                if _cleanup_temp:
                    try:
                        os.remove(temp_path)
                    except OSError:
                        pass
                cleanup_session_files(session_id)
                session_states.pop(session_id, None)

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