src/api_helpers_v1.py

"""V1 helper API endpoints — browser-memory storage mode.

Inference endpoints accept the image directly in the request body (base64 JSON).
No session state lookup needed.

On ZeroGPU Spaces, inference is routed through @spaces.GPU-decorated functions
from gradio_endpoints so a real GPU is allocated for each call.
"""

from __future__ import annotations

import asyncio
import base64
import json
import logging
import os
import re
from collections.abc import AsyncGenerator, Generator
from io import BytesIO
from typing import Any

from fastapi import HTTPException, Request
from fastapi.responses import StreamingResponse
from PIL import Image as PILImage

logger = logging.getLogger(__name__)


def _is_zerogpu() -> bool:
    """Detect HuggingFace ZeroGPU Spaces."""
    return bool(os.environ.get("SPACE_ID")) and bool(os.environ.get("ZERO_GPU"))


def _decode_request_image(image_b64: str) -> PILImage.Image:
    """Decode a base64-encoded image from the request body."""
    img_bytes = base64.b64decode(image_b64)
    return PILImage.open(BytesIO(img_bytes)).convert("RGB")


def _sse_event(data: dict[str, str]) -> str:
    """Format a dict as an SSE data line."""
    return f"data: {json.dumps(data)}\n\n"


async def _sse_stream_v1(
    task_label: str,
    token_generator: Generator[str, None, None],
    post_process: Any | None = None,
) -> AsyncGenerator[str, None]:
    """Wrap a blocking token generator as SSE events (no session check needed).

    The generator yields token deltas. Each delta is sent as a chunk event.
    """
    yield _sse_event({"type": "status", "message": task_label})

    full_text = ""
    try:
        for delta in token_generator:
            if delta:
                full_text += delta
                yield _sse_event({"type": "chunk", "content": delta})
            await asyncio.sleep(0)

        if post_process is not None:
            processed = post_process(full_text)
            if processed != full_text:
                yield _sse_event({"type": "replace", "content": processed})

        yield _sse_event({"type": "done"})
    except Exception as e:  # noqa: BLE001
        logger.error("V1 streaming error: %s", e, exc_info=True)
        yield _sse_event({"type": "error", "message": str(e)})


_GPU_MAX_RETRIES = 3
_GPU_BASE_DELAY = 2.0  # seconds


async def _sse_gpu_call(
    task_label: str,
    gpu_fn: Any,
    image_b64: str,
    post_process: Any | None = None,
) -> AsyncGenerator[str, None]:
    """Call a @spaces.GPU function in a thread and yield the result as SSE.

    Retries up to _GPU_MAX_RETRIES times on 429 / rate-limit errors with
    exponential backoff.
    """
    yield _sse_event({"type": "status", "message": task_label})

    last_exc: Exception | None = None
    loop = asyncio.get_event_loop()

    for attempt in range(_GPU_MAX_RETRIES):
        try:
            result = await loop.run_in_executor(None, gpu_fn, image_b64)

            if post_process is not None:
                result = post_process(result)

            yield _sse_event({"type": "replace", "content": result})
            yield _sse_event({"type": "done"})
            return
        except Exception as e:  # noqa: BLE001
            exc_str = str(e).lower()
            if "429" in exc_str or "too many requests" in exc_str or "queue" in exc_str or ("exceeded" in exc_str and "gpu quota" in exc_str):
                last_exc = e
                delay = _GPU_BASE_DELAY * (2**attempt)
                logger.warning(
                    "ZeroGPU rate-limited (attempt %d/%d), retrying in %.1fs: %s",
                    attempt + 1,
                    _GPU_MAX_RETRIES,
                    delay,
                    e,
                )
                yield _sse_event({"type": "status", "message": "Waiting for GPU to become available..."})
                await asyncio.sleep(delay)
            else:
                logger.error("V1 GPU inference error: %s", e, exc_info=True)
                yield _sse_event({"type": "error", "message": str(e)})
                return

    logger.error("V1 GPU inference failed after %d retries: %s", _GPU_MAX_RETRIES, last_exc)
    yield _sse_event({"type": "error", "message": f"GPU rate-limited after {_GPU_MAX_RETRIES} retries: {last_exc}"})


_STREAM_HEADERS = {
    "Cache-Control": "no-cache",
    "Connection": "keep-alive",
    "X-Accel-Buffering": "no",
}


async def _get_image_from_body(request: Request) -> PILImage.Image:
    """Extract and decode the image from a JSON request body."""
    image_b64 = await _get_image_b64_from_body(request)
    try:
        return _decode_request_image(image_b64)
    except Exception as e:  # noqa: BLE001
        raise HTTPException(status_code=400, detail=f"Invalid image data: {e}") from e


async def _get_image_b64_from_body(request: Request) -> str:
    """Extract the raw base64 image string from a JSON request body."""
    try:
        body = await request.json()
    except Exception as e:  # noqa: BLE001
        raise HTTPException(status_code=400, detail=f"Invalid JSON body: {e}") from e

    image_b64 = body.get("image")
    if not image_b64:
        raise HTTPException(status_code=400, detail="Missing 'image' field in request body")

    return image_b64


def create_helper_routes_v1(app: Any) -> None:
    """Register v1 helper API routes (image in request body)."""

    zerogpu = _is_zerogpu()

    @app.post("/api/v1/helpers/chart2summary/stream")
    async def api_v1_chart2summary_stream(request: Request) -> StreamingResponse:
        if zerogpu:
            from gradio_endpoints import infer_chart2summary_sync

            image_b64 = await _get_image_b64_from_body(request)
            return StreamingResponse(
                _sse_gpu_call("Generating summary...", infer_chart2summary_sync, image_b64),
                media_type="text/event-stream",
                headers=_STREAM_HEADERS,
            )
        from infer_vision_qa import answer_question_stream

        image = await _get_image_from_body(request)
        gen = answer_question_stream(image, "<chart2summary>", [], None)
        return StreamingResponse(
            _sse_stream_v1("Generating summary...", gen),
            media_type="text/event-stream",
            headers=_STREAM_HEADERS,
        )

    @app.post("/api/v1/helpers/chart2csv/stream")
    async def api_v1_chart2csv_stream(request: Request) -> StreamingResponse:
        if zerogpu:
            from gradio_endpoints import infer_chart2csv_sync

            image_b64 = await _get_image_b64_from_body(request)
            return StreamingResponse(
                _sse_gpu_call("Extracting CSV...", infer_chart2csv_sync, image_b64),
                media_type="text/event-stream",
                headers=_STREAM_HEADERS,
            )
        from infer_chart2csv import extract_csv_stream

        image = await _get_image_from_body(request)
        gen = extract_csv_stream(image)
        return StreamingResponse(
            _sse_stream_v1("Extracting CSV...", gen),
            media_type="text/event-stream",
            headers=_STREAM_HEADERS,
        )

    @app.post("/api/v1/helpers/chart2code/stream")
    async def api_v1_chart2code_stream(request: Request) -> StreamingResponse:
        if zerogpu:
            from gradio_endpoints import infer_chart2code_sync

            image_b64 = await _get_image_b64_from_body(request)
            return StreamingResponse(
                _sse_gpu_call("Generating code...", infer_chart2code_sync, image_b64),
                media_type="text/event-stream",
                headers=_STREAM_HEADERS,
            )
        from app import PROMPT_TEXT_CODE
        from infer_vision_qa import answer_question_stream

        image = await _get_image_from_body(request)
        gen = answer_question_stream(image, PROMPT_TEXT_CODE, [], None)
        return StreamingResponse(
            _sse_stream_v1("Generating code...", gen),
            media_type="text/event-stream",
            headers=_STREAM_HEADERS,
        )

    @app.post("/api/v1/helpers/table-extract/stream")
    async def api_v1_table_extract_stream(request: Request) -> StreamingResponse:
        if zerogpu:
            from gradio_endpoints import infer_table_extract_sync

            image_b64 = await _get_image_b64_from_body(request)
            return StreamingResponse(
                _sse_gpu_call("Extracting table...", infer_table_extract_sync, image_b64),
                media_type="text/event-stream",
                headers=_STREAM_HEADERS,
            )
        from infer_vision_qa import answer_question_stream

        def _clean_table_html(text: str) -> str:
            text = re.sub(r"^```(?:html)?\s*", "", text.strip())
            text = re.sub(r"\s*```$", "", text.strip())
            text = re.sub(r"^\[\s*", "", text.strip())
            text = re.sub(r"\s*\]$", "", text.strip())
            return text

        image = await _get_image_from_body(request)
        gen = answer_question_stream(image, "<tables_html>", [], None)
        return StreamingResponse(
            _sse_stream_v1("Extracting table...", gen, post_process=_clean_table_html),
            media_type="text/event-stream",
            headers=_STREAM_HEADERS,
        )

    @app.post("/api/v1/helpers/describe-image/stream")
    async def api_v1_describe_image_stream(request: Request) -> StreamingResponse:
        if zerogpu:
            from gradio_endpoints import infer_describe_image_sync

            image_b64 = await _get_image_b64_from_body(request)
            return StreamingResponse(
                _sse_gpu_call("Describing image...", infer_describe_image_sync, image_b64),
                media_type="text/event-stream",
                headers=_STREAM_HEADERS,
            )
        from infer_vision_qa import answer_question_stream

        image = await _get_image_from_body(request)
        gen = answer_question_stream(image, "Describe this image in detail", [], None)
        return StreamingResponse(
            _sse_stream_v1("Describing image...", gen),
            media_type="text/event-stream",
            headers=_STREAM_HEADERS,
        )