src/app.py

"""Granite Vision Document Intelligence Demo.

Upload a PDF or image to explore Granite-Vision-4.1-4B capabilities including
Chart2CSV, Chart2Code, Chart2Summary, Table Extraction, and Image Q&A.
"""

from __future__ import annotations

# Import spaces first on ZeroGPU — it must be imported before torch to set up
# CUDA emulation correctly.
import os

_GRADIO_MODE = bool(os.environ.get("SPACE_ID"))

if _GRADIO_MODE:
    try:
        import spaces  # noqa: F401, E402
    except ImportError:
        pass

import logging
import uuid
from collections.abc import AsyncGenerator
from contextlib import asynccontextmanager
from pathlib import Path
from typing import Any

logging.basicConfig(
    level=os.environ.get("LOG_LEVEL", "INFO").upper(),
    format="%(asctime)s %(levelname)s %(name)s: %(message)s",
)

from dotenv import load_dotenv

load_dotenv()
load_dotenv(Path(__file__).resolve().parent / ".env", override=False)

from storage import init_storage
init_storage()

if _GRADIO_MODE:
    # Monkey-patch gradio_client to handle bool JSON Schema values.
    # gradio 5.x emits additionalProperties: false/true (valid JSON Schema)
    # but gradio_client 1.5.x does not guard against bool in get_type(),
    # causing TypeError on every request to the /info endpoint.
    try:
        import gradio_client.utils as _gcu

        _orig_get_type = _gcu.get_type
        _orig_j2p = _gcu._json_schema_to_python_type

        def _patched_get_type(schema):  # noqa: ANN001, ANN202
            if not isinstance(schema, dict):
                return "unknown"
            return _orig_get_type(schema)

        def _patched_j2p(schema, defs=None):  # noqa: ANN001, ANN202
            if not isinstance(schema, dict):
                return "any" if schema else "unknown"
            return _orig_j2p(schema, defs)

        _gcu.get_type = _patched_get_type
        _gcu._json_schema_to_python_type = _patched_j2p
    except Exception:  # noqa: BLE001
        pass

    import gradio as gr
    from gradio import Server

from fastapi import FastAPI, UploadFile, File, HTTPException
from fastapi.responses import JSONResponse
from fastapi.middleware.cors import CORSMiddleware
from starlette.middleware.base import BaseHTTPMiddleware
from starlette.responses import Response
from PIL import Image

from crops import extract_figures
from document_parser import parse_document
from infer_chart2csv import extract_csv, extract_csv_stream
from infer_vision_qa import answer_question, answer_question_stream
from pdf_io import load_pdf_pages
from storage import load_parse_cache, resolve_for_gradio, save_parse_cache, save_image, use_disk_images
from ui_state import create_initial_state, hash_bytes, page_cache, parse_cache

if _GRADIO_MODE:
    from themes.research_monochrome import theme

TITLE = "Granite Vision: Document Intelligence"
DESCRIPTION = (
    "Upload a PDF or image to explore Granite-Vision-4.1-4B's document intelligence capabilities — "
    "including Chart2Summary, Chart2CSV, Chart2Code, Table Extraction, and Image Description — "
    "with automatic Docling-powered parsing for PDFs and direct inference on uploaded images."
)

IMAGE_EXTENSIONS = {".jpg", ".jpeg", ".jfif", ".png", ".bmp", ".dib", ".gif", ".tif", ".tiff", ".webp"}
OFFICE_EXTENSIONS = {".docx", ".xlsx", ".pptx"}

css_file_path = Path(Path(__file__).parent / "app.css")
head_file_path = Path(Path(__file__).parent / "app_head.html")

# In-memory session storage for API requests
session_states: dict[str, dict[str, Any]] = {}


def _is_image_file(file_path: str) -> bool:
    """Check whether a file path points to a supported image format."""
    ext = os.path.splitext(file_path)[1].lower()
    return ext in IMAGE_EXTENSIONS


def _is_office_file(file_path: str) -> bool:
    """Check whether a file path points to a supported Office format (DOCX/XLSX/PPTX)."""
    ext = os.path.splitext(file_path)[1].lower()
    return ext in OFFICE_EXTENSIONS


def process_upload(file_path: str, session_state: dict[str, Any]) -> tuple:
    """Parse an uploaded PDF or load an image and extract figures.

    Args:
        file_path: Path to the uploaded file.
        session_state: Current Gradio session state dictionary.

    Returns:
        Tuple of (status, html_content, fig_status, fig_caption, fig_image, session_state).
    """
    max_pages = 20
    sid = str(uuid.uuid4())

    session_state["current_figure_index"] = 0
    session_state["conversation_history"] = []
    session_state["current_image_path"] = None

    if not file_path:
        return "Please upload a PDF, Office document, or image.", "No document loaded", "No figures", "", None, session_state

    try:
        with open(file_path, "rb") as f:
            file_bytes = f.read()

        file_hash = hash_bytes(file_bytes)
        session_state["uploaded_file_hash"] = file_hash
        if not use_disk_images():
            session_state["uploaded_file_bytes"] = file_bytes
        if _is_image_file(file_path):
            image = Image.open(file_path).convert("RGB")
            lazy = save_image(sid, "figures", 0, image)  # LazyImage or PIL Image
            figures_info = [{"image": lazy, "page": 0, "bbox": None, "caption": ""}]

            session_state["page_images"] = [lazy]  # LazyImage proxies in disk mode, PIL Images in memory mode
            if not use_disk_images():
                session_state["parsed_result"] = {}
            session_state["figures_info"] = figures_info  # fig["image"] is LazyImage or PIL Image
            session_state["selected_figure"] = figures_info[0]  # reference to figures_info entry

            return (
                "Image loaded successfully.\nNumber of figures: 1.",
                "Image uploaded directly (no document parsing needed)",
                "Figure 1 of 1 (Page 1)",
                "",
                image,
                session_state,
            )

        file_ext = os.path.splitext(file_path)[1].lower()
        is_office = _is_office_file(file_path)
        fmt_label = file_ext.lstrip(".").upper()
        status_lines = [f"{fmt_label} loaded successfully."]

        if is_office:
            page_images: list = []
            session_state["page_images"] = []
        else:
            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(file_bytes, max_pages=max_pages)
                if not use_disk_images():
                    page_cache[cache_key] = page_images
            session_state["page_images"] = [  # LazyImage proxies in disk mode, PIL Images in memory mode
                save_image(sid, "pages", i, img) for i, img in enumerate(page_images)
            ]
            status_lines.append(f"Number of pages rendered: {len(page_images)} (max {max_pages}).")

        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=sid)
            if parse_result is None:
                parse_result = parse_document(file_bytes, file_ext=file_ext)
                save_parse_cache(file_hash, parse_result, session_id=sid)
            if not use_disk_images():
                parse_cache[file_hash] = parse_result
        if not use_disk_images():
            session_state["parsed_result"] = parse_result
        status_lines.append("Document parsing done using Docling.")

        figures_info = extract_figures(page_images, parse_result.get("figures", []))
        for i, fig in enumerate(figures_info):
            fig["image"] = save_image(sid, "figures", i, fig["image"])  # LazyImage or PIL Image
        session_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:
            session_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 = resolve_for_gradio(figures_info[0]["image"])
        else:
            session_state["selected_figure"] = None
            fig_status = "No figures found"
            fig_caption = ""
            fig_image = None

        html_content = parse_result.get("html", "No content available")
        status = "\n".join(status_lines)

        return status, html_content, fig_status, fig_caption, fig_image, session_state

    except Exception as e:  # noqa: BLE001
        import traceback

        print(f"Error: {e}")
        traceback.print_exc()
        return f"Error: {e!s}", f"Error loading document: {e!s}", "Error", "", None, session_state


def _get_figure_display(session_state: dict[str, Any]) -> tuple[str, str, Image.Image | None]:
    """Return the current figure's display info, caption, and image.

    Args:
        session_state: Current session state dictionary.

    Returns:
        Tuple of (fig_status, fig_caption, fig_image).
    """
    figures_info = session_state.get("figures_info", [])
    idx = session_state.get("current_figure_index", 0)

    if not figures_info:
        return "No figures found", "", None

    fig = figures_info[idx]
    fig_status = f"Figure {idx + 1} of {len(figures_info)} (Page {fig['page'] + 1})"
    fig_caption = fig.get("caption", "No caption")
    # fig["image"] is a LazyImage (disk mode) or PIL Image (memory mode).
    # resolve_for_gradio converts LazyImage to a Path that Gradio can postprocess.
    return fig_status, fig_caption, resolve_for_gradio(fig["image"])


def next_figure(session_state: dict[str, Any]) -> tuple:
    """Advance to the next figure.

    Args:
        session_state: Current session state dictionary.

    Returns:
        Tuple of (fig_status, fig_caption, fig_image, session_state).
    """
    figures_info = session_state.get("figures_info", [])

    if not figures_info:
        return "No figures found", "", None, session_state

    idx = (session_state.get("current_figure_index", 0) + 1) % len(figures_info)
    session_state["current_figure_index"] = idx
    session_state["selected_figure"] = figures_info[idx]
    session_state["conversation_history"] = []
    session_state["current_image_path"] = None

    fig_status, fig_caption, fig_image = _get_figure_display(session_state)
    return fig_status, fig_caption, fig_image, session_state


def prev_figure(session_state: dict[str, Any]) -> tuple:
    """Go back to the previous figure.

    Args:
        session_state: Current session state dictionary.

    Returns:
        Tuple of (fig_status, fig_caption, fig_image, session_state).
    """
    figures_info = session_state.get("figures_info", [])

    if not figures_info:
        return "No figures found", "", None, session_state

    idx = (session_state.get("current_figure_index", 0) - 1) % len(figures_info)
    session_state["current_figure_index"] = idx
    session_state["selected_figure"] = figures_info[idx]
    session_state["conversation_history"] = []
    session_state["current_image_path"] = None

    fig_status, fig_caption, fig_image = _get_figure_display(session_state)
    return fig_status, fig_caption, fig_image, session_state


def describe_image_helper(session_state: dict[str, Any]):  # noqa: ANN201
    """Generate a detailed description of the selected figure (streaming)."""
    selected_fig = session_state.get("selected_figure")
    if selected_fig is None:
        yield "No figure selected", session_state
        return
    try:
        image = selected_fig["image"]
        accumulated = ""
        for token in answer_question_stream(image, "Describe this image in detail", [], None):
            accumulated += token
            yield accumulated, session_state
    except Exception as e:  # noqa: BLE001
        yield f"Error: {e!s}", session_state


def load_current_figure(session_state: dict[str, Any]) -> tuple[str, str, Image.Image | None]:
    """Load the current figure into display components (called on tab select).

    Args:
        session_state: Current session state dictionary.

    Returns:
        Tuple of (fig_status, fig_caption, fig_image).
    """
    return _get_figure_display(session_state)


PROMPT_TEXT_CODE = (
    "Please take a look at this chart image and generate Python code that perfectly reconstructs this chart image."
)

PROMPT_TEXT_SUMMARY = "<chart2summary>"

PROMPT_TEXT_TABLE = "<tables_html>"


def extract_code_helper(session_state: dict[str, Any]):  # noqa: ANN201
    """Generate Python code to reconstruct the selected chart (streaming)."""
    selected_fig = session_state.get("selected_figure")
    if selected_fig is None:
        yield "No figure selected", session_state
        return
    try:
        image = selected_fig["image"]
        accumulated = ""
        for token in answer_question_stream(image, PROMPT_TEXT_CODE, [], None):
            accumulated += token
            yield accumulated, session_state
    except Exception as e:  # noqa: BLE001
        yield f"Error: {e!s}", session_state


def extract_summary_helper(session_state: dict[str, Any]):  # noqa: ANN201
    """Generate a text summary of the selected chart (streaming)."""
    selected_fig = session_state.get("selected_figure")
    if selected_fig is None:
        yield "No figure selected", session_state
        return
    try:
        image = selected_fig["image"]
        accumulated = ""
        for token in answer_question_stream(image, PROMPT_TEXT_SUMMARY, [], None):
            accumulated += token
            yield accumulated, session_state
    except Exception as e:  # noqa: BLE001
        yield f"Error: {e!s}", session_state


def extract_table_helper(session_state: dict[str, Any]):  # noqa: ANN201
    """Extract tables as HTML from the selected figure (streaming)."""
    import re

    selected_fig = session_state.get("selected_figure")
    if selected_fig is None:
        yield "No figure selected", session_state
        return
    try:
        image = selected_fig["image"]
        result = ""
        for token in answer_question_stream(image, PROMPT_TEXT_TABLE, [], None):
            result += token
            yield result, session_state
        # Final cleanup pass on the complete output
        result = re.sub(r"^```(?:html)?\s*", "", result.strip())
        result = re.sub(r"\s*```$", "", result.strip())
        result = re.sub(r"^\[\s*", "", result.strip())
        result = re.sub(r"\s*\]$", "", result.strip())
        yield result, session_state
    except Exception as e:  # noqa: BLE001
        yield f"Error: {e!s}", session_state




def extract_csv_helper(session_state: dict[str, Any]):  # noqa: ANN201
    """Extract CSV data from the selected chart (streaming)."""
    selected_fig = session_state.get("selected_figure")
    if selected_fig is None:
        yield "No figure selected", session_state
        return
    try:
        image = selected_fig["image"]
        csv_text = ""
        for token in extract_csv_stream(image):
            csv_text += token
            yield csv_text, session_state
        session_state["last_csv"] = csv_text
    except Exception as e:  # noqa: BLE001
        yield f"Error: {e!s}", session_state


if _GRADIO_MODE:
    demo = gr.Blocks(
        title=TITLE,
        theme=theme,
        css_paths=css_file_path,
        head_paths=head_file_path,
        fill_height=True,
    )
    demo.queue()

    with demo:
        gr.Markdown(f"# {TITLE}")
        gr.Markdown(DESCRIPTION)

        session_state = gr.State(create_initial_state())

        with gr.Tabs():
            # TAB 1: UPLOAD & PARSE
            with gr.Tab("Parse & Extract"):
                with gr.Row():
                    file_path = gr.File(
                        label="Upload PDF, Office Document, or Image",
                        file_types=[".pdf", ".docx", ".xlsx", ".pptx", ".jpg", ".jpeg", ".jfif", ".png", ".bmp", ".dib", ".gif", ".tif", ".tiff", ".webp"],
                        scale=4,
                    )
                    load_btn = gr.Button("Load Document", variant="primary", scale=1)

                status = gr.Textbox(label="Status", interactive=False, lines=2)

                with gr.Row():
                    with gr.Column(scale=1):
                        html_view = gr.Textbox(
                            label="Parsed Document (Docling)",
                            value="Upload a PDF to see parsed content",
                            lines=35,
                            interactive=False,
                        )

                    with gr.Column(scale=1):
                        gr.Markdown("### Extracted Figures")
                        fig_info = gr.Textbox(label="Figure Info", interactive=False)
                        fig_caption = gr.Textbox(label="Caption", interactive=False)
                        fig_image = gr.Image(label="Figure", type="pil", elem_classes=["figure-image"])

                        with gr.Row():
                            prev_btn = gr.Button("Previous", scale=1)
                            next_btn = gr.Button("Next", scale=1)

                load_btn.click(
                    process_upload,
                    inputs=[file_path, session_state],
                    outputs=[status, html_view, fig_info, fig_caption, fig_image, session_state],
                )
                next_btn.click(
                    next_figure,
                    inputs=[session_state],
                    outputs=[fig_info, fig_caption, fig_image, session_state],
                )
                prev_btn.click(
                    prev_figure,
                    inputs=[session_state],
                    outputs=[fig_info, fig_caption, fig_image, session_state],
                )

            # TAB 2: CHART2SUMMARY
            with gr.Tab("Chart2Summary") as summary_tab:
                gr.Markdown("Generate a text summary of the selected chart")

                with gr.Row():
                    with gr.Column(scale=1):
                        gr.Markdown("### Figure")
                        summary_fig_info = gr.Textbox(label="Figure Info", interactive=False)
                        summary_fig_caption = gr.Textbox(label="Caption", interactive=False)
                        summary_fig_image = gr.Image(label="Figure", type="pil", elem_classes=["figure-image"])

                        with gr.Row():
                            summary_prev_btn = gr.Button("Previous", scale=1)
                            summary_next_btn = gr.Button("Next", scale=1)

                    with gr.Column(scale=1):
                        gr.Markdown("### Summary")
                        summary_btn = gr.Button("Generate Summary", variant="primary")
                        summary_out = gr.Textbox(label="Chart Summary", lines=20, interactive=False)

                summary_prev_btn.click(prev_figure, inputs=[session_state], outputs=[summary_fig_info, summary_fig_caption, summary_fig_image, session_state])
                summary_next_btn.click(next_figure, inputs=[session_state], outputs=[summary_fig_info, summary_fig_caption, summary_fig_image, session_state])
                summary_btn.click(extract_summary_helper, inputs=[session_state], outputs=[summary_out, session_state])
                summary_tab.select(load_current_figure, inputs=[session_state], outputs=[summary_fig_info, summary_fig_caption, summary_fig_image])

            # TAB 3: CHART2CSV
            with gr.Tab("Chart2CSV") as csv_tab:
                gr.Markdown("Extract CSV data from the selected chart")

                with gr.Row():
                    with gr.Column(scale=1):
                        gr.Markdown("### Figure")
                        csv_fig_info = gr.Textbox(label="Figure Info", interactive=False)
                        csv_fig_caption = gr.Textbox(label="Caption", interactive=False)
                        csv_fig_image = gr.Image(label="Figure", type="pil", elem_classes=["figure-image"])

                        with gr.Row():
                            csv_prev_btn = gr.Button("Previous", scale=1)
                            csv_next_btn = gr.Button("Next", scale=1)

                    with gr.Column(scale=1):
                        gr.Markdown("### CSV Extraction")
                        extract_btn = gr.Button("Extract CSV", variant="primary")
                        csv_out = gr.Textbox(label="CSV", lines=20, interactive=False)

                csv_prev_btn.click(prev_figure, inputs=[session_state], outputs=[csv_fig_info, csv_fig_caption, csv_fig_image, session_state])
                csv_next_btn.click(next_figure, inputs=[session_state], outputs=[csv_fig_info, csv_fig_caption, csv_fig_image, session_state])
                extract_btn.click(extract_csv_helper, inputs=[session_state], outputs=[csv_out, session_state])
                csv_tab.select(load_current_figure, inputs=[session_state], outputs=[csv_fig_info, csv_fig_caption, csv_fig_image])

            # TAB 4: CHART2CODE
            with gr.Tab("Chart2Code") as code_tab:
                gr.Markdown("Generate Python code to reconstruct the selected chart")

                with gr.Row():
                    with gr.Column(scale=1):
                        gr.Markdown("### Figure")
                        code_fig_info = gr.Textbox(label="Figure Info", interactive=False)
                        code_fig_caption = gr.Textbox(label="Caption", interactive=False)
                        code_fig_image = gr.Image(label="Figure", type="pil", elem_classes=["figure-image"])

                        with gr.Row():
                            code_prev_btn = gr.Button("Previous", scale=1)
                            code_next_btn = gr.Button("Next", scale=1)

                    with gr.Column(scale=1):
                        gr.Markdown("### Generated Code")
                        code_btn = gr.Button("Generate Code", variant="primary")
                        code_out = gr.Textbox(label="Python Code", lines=20, interactive=False)

                code_prev_btn.click(prev_figure, inputs=[session_state], outputs=[code_fig_info, code_fig_caption, code_fig_image, session_state])
                code_next_btn.click(next_figure, inputs=[session_state], outputs=[code_fig_info, code_fig_caption, code_fig_image, session_state])
                code_btn.click(extract_code_helper, inputs=[session_state], outputs=[code_out, session_state])
                code_tab.select(load_current_figure, inputs=[session_state], outputs=[code_fig_info, code_fig_caption, code_fig_image])

            # TAB 5: TABLE EXTRACTION
            with gr.Tab("Table Extraction") as table_tab:
                gr.Markdown("Extract table data as HTML from the selected figure")

                with gr.Row():
                    with gr.Column(scale=1):
                        gr.Markdown("### Figure")
                        table_fig_info = gr.Textbox(label="Figure Info", interactive=False)
                        table_fig_caption = gr.Textbox(label="Caption", interactive=False)
                        table_fig_image = gr.Image(label="Figure", type="pil", elem_classes=["figure-image"])

                        with gr.Row():
                            table_prev_btn = gr.Button("Previous", scale=1)
                            table_next_btn = gr.Button("Next", scale=1)

                    with gr.Column(scale=1):
                        gr.Markdown("### Table Extraction")
                        table_btn = gr.Button("Extract Table", variant="primary")
                        table_out = gr.HTML(value="<p>Upload a document and click Extract Table to see results here</p>")

                table_prev_btn.click(prev_figure, inputs=[session_state], outputs=[table_fig_info, table_fig_caption, table_fig_image, session_state])
                table_next_btn.click(next_figure, inputs=[session_state], outputs=[table_fig_info, table_fig_caption, table_fig_image, session_state])
                table_btn.click(extract_table_helper, inputs=[session_state], outputs=[table_out, session_state])
                table_tab.select(load_current_figure, inputs=[session_state], outputs=[table_fig_info, table_fig_caption, table_fig_image])

            # TAB 6: IMAGE DESCRIPTION
            with gr.Tab("Image Description") as qa_tab:
                gr.Markdown("Get a detailed description of the selected figure")

                with gr.Row():
                    with gr.Column(scale=1):
                        gr.Markdown("### Figure")
                        qa_fig_info = gr.Textbox(label="Figure Info", interactive=False)
                        qa_fig_caption = gr.Textbox(label="Caption", interactive=False)
                        qa_fig_image = gr.Image(label="Figure", type="pil", elem_classes=["figure-image"])

                        with gr.Row():
                            qa_prev_btn = gr.Button("Previous", scale=1)
                            qa_next_btn = gr.Button("Next", scale=1)

                    with gr.Column(scale=1):
                        gr.Markdown("### Description")
                        describe_btn = gr.Button("Describe Image", variant="primary")
                        answer = gr.Textbox(label="Description", lines=20, interactive=False)

                qa_prev_btn.click(prev_figure, inputs=[session_state], outputs=[qa_fig_info, qa_fig_caption, qa_fig_image, session_state])
                qa_next_btn.click(next_figure, inputs=[session_state], outputs=[qa_fig_info, qa_fig_caption, qa_fig_image, session_state])
                describe_btn.click(describe_image_helper, inputs=[session_state], outputs=[answer, session_state])
                qa_tab.select(load_current_figure, inputs=[session_state], outputs=[qa_fig_info, qa_fig_caption, qa_fig_image])

        # Register inference endpoints inside the Blocks context so they are
        # discoverable by @gradio/client via /gradio_api/info
        from gradio_endpoints import ALL_ENDPOINTS as _gradio_endpoints

        for _api_name, _fn in _gradio_endpoints.items():
            gr.api(_fn, api_name=_api_name)


def _verify_offline_models() -> None:
    """Check that required models are cached locally when OFFLINE_MODE is on."""
    from huggingface_hub import try_to_load_from_cache

    from model_loader import get_model_name, get_mlx_model_name, use_mlx_mode

    missing = []
    model_name = get_model_name()
    if try_to_load_from_cache(model_name, "config.json") is None:
        missing.append(model_name)

    if use_mlx_mode():
        mlx_name = get_mlx_model_name()
        if try_to_load_from_cache(mlx_name, "config.json") is None:
            missing.append(mlx_name)

    if missing:
        raise SystemExit(
            "OFFLINE_MODE is enabled but these models are not cached:\n"
            + "\n".join(f"  - {m}" for m in missing)
            + "\nRun while online: bash scripts/preload_offline.sh"
        )


@asynccontextmanager
async def _lifespan(app: Any) -> AsyncGenerator[None]:
    from model_loader import load_model, load_mlx_model, use_api_mode, use_mlx_mode

    if os.environ.get("OFFLINE_MODE", "").lower() in ("1", "true"):
        os.environ.setdefault("HF_HUB_OFFLINE", "1")
        if not use_api_mode():
            _verify_offline_models()

    if not use_api_mode():
        if use_mlx_mode():
            load_mlx_model()
        else:
            load_model()

    try:
        from document_parser import get_converter
        get_converter()
    except Exception:  # noqa: BLE001
        pass

    yield


if _GRADIO_MODE:
    app = Server(lifespan=_lifespan)
else:
    app = FastAPI(lifespan=_lifespan)


# Add CORS middleware
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)


if _GRADIO_MODE:
    class HeadRequestMiddleware(BaseHTTPMiddleware):
        """Return a plain 200 for HEAD requests to root.

        Gradio's template renderer crashes on HEAD / because it tries to render
        the Jinja template without a populated config. This intercepts HEAD
        requests before they reach Gradio's route handler.
        """

        async def dispatch(self, request, call_next):  # noqa: ANN001, ANN201
            if request.method == "HEAD" and request.url.path == "/":
                return Response(status_code=200)
            return await call_next(request)

    app.add_middleware(HeadRequestMiddleware)


class FrameAncestorsMiddleware(BaseHTTPMiddleware):
    """Restrict who may embed this app in an iframe.

    The frontend opts into the ZeroGPU token handshake only when framed by
    huggingface.co; this header is the primary control that prevents a hostile
    site from embedding the Space (and being asked for / supplying a visitor
    token) in the first place. Allowing only the Hugging Face origins keeps the
    legitimate Space embedding (``*.hf.space`` inside ``huggingface.co``) working.
    """

    async def dispatch(self, request, call_next):  # noqa: ANN001, ANN201
        response = await call_next(request)
        response.headers["Content-Security-Policy"] = (
            "frame-ancestors 'self' https://huggingface.co https://*.huggingface.co"
        )
        return response


app.add_middleware(FrameAncestorsMiddleware)




# Register API routes
from api_routes import create_document_routes
from api_helpers import create_helper_routes

create_document_routes(app, session_states, process_upload, next_figure, prev_figure)
create_helper_routes(app, session_states)


@app.get("/api/config")
async def api_config() -> JSONResponse:
    """Return runtime configuration flags for the frontend."""
    config: dict[str, Any] = {"v1": _v1 or BUILD_DIR.exists()}
    return JSONResponse(config)


from storage import v1_mode

_v1 = v1_mode()
logging.getLogger(__name__).info("BROWSER_MEMORY_MODE: %s", _v1)

if _GRADIO_MODE:
    from gradio_endpoints import register_gradio_api_endpoints
    register_gradio_api_endpoints(app)

# Serve static React build if available, otherwise mount Gradio UI
_app_dir = Path(__file__).resolve().parent
_env_build = os.environ.get("STATIC_BUILD_DIR")
if _env_build:
    _env_path = Path(_env_build)
    BUILD_DIR = _env_path if _env_path.is_absolute() else _app_dir / _env_path
else:
    # Check both: project-root/build (src/ layout) and app-dir/build (flat layout)
    _project_build = _app_dir.parent / "build"
    _flat_build = _app_dir / "build"
    BUILD_DIR = _flat_build if _flat_build.exists() else _project_build

# Enable V1 (browser-memory) routes when explicitly set or when serving a
# static React build (the React frontend requires V1 mode).
if _v1 or BUILD_DIR.exists():
    from api_helpers_v1 import create_helper_routes_v1
    from api_routes_v1 import create_document_routes_v1

    create_document_routes_v1(app, session_states)
    create_helper_routes_v1(app)
    if not _v1:
        logging.getLogger(__name__).info("V1 routes auto-enabled (static build detected)")

if BUILD_DIR.exists():
    from fastapi.staticfiles import StaticFiles
    from fastapi.responses import FileResponse

    @app.get("/")
    async def serve_index():
        if os.environ.get("NEXT_PUBLIC_VISION_ONLY", "").lower() in ("1", "true"):
            from fastapi.responses import RedirectResponse
            return RedirectResponse("/vision-demo")
        return FileResponse(BUILD_DIR / "index.html")

    @app.api_route("/granite-vision", methods=["GET", "HEAD"])
    async def serve_granite_vision():
        if os.environ.get("NEXT_PUBLIC_VISION_ONLY", "").lower() in ("1", "true"):
            from fastapi.responses import RedirectResponse
            return RedirectResponse("/vision-demo")
        return FileResponse(BUILD_DIR / "granite-vision.html")

    @app.get("/vision-demo")
    async def serve_vision_demo():
        return FileResponse(BUILD_DIR / "vision-demo.html")

    app.mount("/_next", StaticFiles(directory=BUILD_DIR / "_next"), name="next_static")
    app.mount("/assets", StaticFiles(directory=BUILD_DIR / "assets"), name="assets")

    if _GRADIO_MODE:
        # Pre-register Gradio's internal /gradio_api/* routes BEFORE adding the SPA
        # catch-all below. Otherwise the catch-all matches /gradio_api/startup-events
        # (which Server.launch() probes during startup) and causes a 404.
        from gradio.blocks import Blocks as _Blocks
        from gradio.events import api as _gr_api
        from gradio.routes import App as _GrApp

        with _Blocks() as _internal_blocks:
            for _fn, _api_kwargs in app._deferred_apis:
                _gr_api(fn=_fn, **_api_kwargs)
        _internal_blocks.config = _internal_blocks.get_config_file()
        _internal_blocks.validate_queue_settings()
        _GrApp.create_app(_internal_blocks, app=app)

    _RESERVED_PREFIXES = ("api/", "gradio_api/", "openapi.json", "docs", "redoc")

    @app.get("/{full_path:path}")
    async def serve_spa(full_path: str):
        if full_path.startswith(_RESERVED_PREFIXES) or full_path in ("api", "gradio_api"):
            raise HTTPException(status_code=404)
        candidate = BUILD_DIR / full_path
        if candidate.is_file():
            return FileResponse(candidate)
        html_candidate = BUILD_DIR / f"{full_path}.html"
        if html_candidate.is_file():
            return FileResponse(html_candidate)
        return FileResponse(BUILD_DIR / "index.html")

    logging.getLogger(__name__).info("Serving static frontend from %s", BUILD_DIR)
elif _GRADIO_MODE:
    gr.mount_gradio_app(app, demo, path="/")
    logging.getLogger(__name__).info("Serving Gradio UI (no static build at %s)", BUILD_DIR)
else:
    from fastapi.responses import HTMLResponse

    @app.get("/")
    async def local_mode_root():
        return HTMLResponse("""<!DOCTYPE html>
<html><head><meta charset="utf-8"><title>Granite Vision</title>
<style>body{font-family:system-ui,sans-serif;display:flex;justify-content:center;align-items:center;min-height:100vh;margin:0;background:#f5f5f5}
.card{background:#fff;border-radius:8px;padding:2.5rem;max-width:480px;box-shadow:0 2px 8px rgba(0,0,0,.08);text-align:center}
h1{margin:0 0 .5rem;font-size:1.5rem}p{color:#555;line-height:1.5}
a{color:#0066cc;text-decoration:none}a:hover{text-decoration:underline}</style></head>
<body><div class="card">
<h1>Granite Vision API</h1>
<p>The API server is running, but no frontend build was found.</p>
<p>To explore the available endpoints, visit the <a href="/docs">API docs</a>.</p>
<p style="margin-top:1.5rem;font-size:.85rem;color:#888">To serve the full UI, build the frontend and restart the server.</p>
</div></body></html>""")

    logging.getLogger(__name__).info("Local mode: serving API only (no static build at %s)", BUILD_DIR)


if __name__ == "__main__":
    if _GRADIO_MODE:
        app.launch(server_name="0.0.0.0", server_port=7860)
    else:
        import uvicorn
        uvicorn.run(app, host="0.0.0.0", port=7860)