backend/app/api/visuals.py

"""
Visuals API — Serve Page Images
================================
GET /api/visuals/{doc_id}/{page_number}
Returns the rendered PDF page image for display in the UI.
"""

import logging
from pathlib import Path
from fastapi import APIRouter, HTTPException
from fastapi.responses import FileResponse

logger = logging.getLogger(__name__)
router = APIRouter()

BASE_DIR  = Path(__file__).parent.parent.parent
DATA_DIR  = BASE_DIR / "data"
PAGES_DIR = DATA_DIR / "pages"


def _page_image_candidates(doc_id: str, page_number: int) -> list[Path]:
    pages_dir = PAGES_DIR / doc_id / "pages"
    page_stem = f"page_{page_number:04d}"
    return [
        pages_dir / f"{page_stem}.png",
        pages_dir / f"{page_stem}_colpali_index.png",
        pages_dir / f"{page_stem}_colpali.png",
    ]


@router.get("/{doc_id}/{page_number}")
async def get_page_image(doc_id: str, page_number: int):
    """
    Serve the rendered full-resolution page image for a given document page.
    Used by the RightEvidencePanel PDF viewer.
    """
    img_path = next(
        (candidate for candidate in _page_image_candidates(doc_id, page_number) if candidate.exists()),
        None,
    )

    if not img_path:
        raise HTTPException(
            status_code=404,
            detail=f"Page image not found: doc={doc_id}, page={page_number}. "
                   "Run the ingestion script first.",
        )

    return FileResponse(
        path=str(img_path),
        media_type="image/png",
        headers={"Cache-Control": "public, max-age=3600"},
    )


@router.get("/{doc_id}/{page_number}/colpali")
async def get_colpali_image(doc_id: str, page_number: int):
    """Serve the image used for ColPali indexing (for debugging)."""
    pages_dir = PAGES_DIR / doc_id / "pages"
    page_stem = f"page_{page_number:04d}"
    img_path = next(
        (
            candidate
            for candidate in [
                pages_dir / f"{page_stem}_colpali_index.png",
                pages_dir / f"{page_stem}_colpali.png",
            ]
            if candidate.exists()
        ),
        None,
    )

    if not img_path:
        raise HTTPException(status_code=404, detail="ColPali image not found")

    return FileResponse(str(img_path), media_type="image/png")


@router.get("/{doc_id}/index")
async def get_visual_index(doc_id: str):
    """Return list of all indexed pages with metadata for a document."""
    index_path = BASE_DIR / "data" / "colpali_index" / doc_id / "colpali_index.json"

    if not index_path.exists():
        raise HTTPException(status_code=404, detail=f"No ColPali index for doc '{doc_id}'")

    import json
    with open(index_path) as f:
        records = json.load(f)

    return {
        "doc_id": doc_id,
        "total_pages": len(records),
        "pages": [
            {
                "page_number": r["page_number"],
                "has_embedding": bool(r.get("colpali_embedding_path")),
                "image_url": f"/api/visuals/{doc_id}/{r['page_number']}",
            }
            for r in records
        ],
    }