src/api_helpers.py

"""API endpoints for document intelligence helper functions.

Provides JSON endpoints to call Chart2Summary, Chart2CSV, Chart2Code,
Table Extraction, and Image Q&A helpers.
"""

from __future__ import annotations

import asyncio
import base64
import io
import json
from collections.abc import AsyncGenerator
import logging
import queue
import re
from collections.abc import AsyncGenerator, Callable, Generator

from typing import Any

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

logger = logging.getLogger(__name__)

async def _stream_with_keepalive(sync_fn: Any, session_id: str = "") -> AsyncGenerator[str, None]:
    """Run *sync_fn* in a thread and yield SSE keepalive comments every 15 s.

    Yields SSE comment lines (``:\\n\\n``) while the function is running,
    then a ``data:`` event with the JSON result once it completes.
    On error, yields an ``event: error`` SSE event.
    """
    logger.debug(f"Calling {sync_fn.__name__} for session {session_id}")
    loop = asyncio.get_event_loop()
    task = loop.run_in_executor(None, sync_fn)

    while True:
        try:
            result = await asyncio.wait_for(asyncio.shield(task), timeout=15)
            logger.info(f"Successfully completed {sync_fn.__name__} for session {session_id}")
            yield f"data: {json.dumps(result)}\n\n"
            return
        except asyncio.TimeoutError:
            yield ":\n\n"  # SSE comment — keeps connection alive
        except Exception as exc:  # noqa: BLE001
            logger.error("Helper %s failed for session %s: %s", sync_fn.__name__, session_id, exc, exc_info=True)
            yield f"event: error\ndata: {json.dumps({'detail': str(exc)})}\n\n"
            return

async def _stream_with_progress(
    sync_fn: Callable[[Callable[[str], None]], Any],
    session_id: str = "",
) -> AsyncGenerator[str, None]:
    """Run *sync_fn(on_progress)* in a thread, yielding SSE progress events in real-time.

    The sync function receives an ``on_progress(phase)`` callback that it can
    call at any point.  Each call puts a message on a thread-safe queue which
    the async loop drains as ``event: progress`` SSE events.
    """
    progress_q: queue.Queue[str] = queue.Queue()
    loop = asyncio.get_event_loop()

    def _on_progress(phase: str) -> None:
        progress_q.put(phase)

    task = loop.run_in_executor(None, sync_fn, _on_progress)

    while True:
        while not progress_q.empty():
            phase = progress_q.get_nowait()
            yield f"event: progress\ndata: {json.dumps({'phase': phase})}\n\n"

        try:
            result = await asyncio.wait_for(asyncio.shield(task), timeout=1)
            while not progress_q.empty():
                phase = progress_q.get_nowait()
                yield f"event: progress\ndata: {json.dumps({'phase': phase})}\n\n"
            logger.info("Upload completed for session %s", session_id)
            yield f"data: {json.dumps(result)}\n\n"
            return
        except asyncio.TimeoutError:
            yield ":\n\n"
        except Exception as exc:  # noqa: BLE001
            logger.error("Upload failed for session %s: %s", session_id, exc, exc_info=True)
            yield f"event: error\ndata: {json.dumps({'detail': str(exc)})}\n\n"
            return


def create_helper_routes(app: Any, session_states: dict[str, dict[str, Any]]) -> None:
    """Register helper API routes on the FastAPI app.

    Args:
        app: FastAPI application instance.
        session_states: Global session states dictionary.
    """

    @app.post("/api/helpers/chart2summary")
    async def api_chart2summary(session_id: str = Query(...)) -> StreamingResponse:  # noqa: ANN001
        """Extract chart summary for current session figure.

        Args:
            session_id: Session ID to get figure from.

        Returns:
            JSON with result and session_state.
        """
        logger.info(f"API chart2summary request received for session_id: {session_id}")

        if session_id not in session_states:
            logger.error(f"Session not found: {session_id}")
            raise HTTPException(status_code=404, detail="Session not found")

        def _work() -> dict:
            logger.debug(f"Importing extract_summary_helper for session {session_id}")
            from app import extract_summary_helper

            logger.debug(f"Calling extract_summary_helper for session {session_id}")
            result, updated_state = extract_summary_helper(session_states[session_id])
            session_states[session_id] = updated_state
            return {"result": result}

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

    @app.post("/api/helpers/chart2csv")
    async def api_chart2csv(session_id: str = Query(...)) -> StreamingResponse:  # noqa: ANN001
        """Extract chart data as CSV for current session figure.

        Args:
            session_id: Session ID to get figure from.

        Returns:
            JSON with result and session_state.
        """
        logger.info(f"API chart2csv request received for session_id: {session_id}")

        if session_id not in session_states:
            logger.error(f"Session not found: {session_id}")
            raise HTTPException(status_code=404, detail="Session not found")

        def _work() -> dict:
            logger.debug(f"Importing extract_csv_helper for session {session_id}")
            from app import extract_csv_helper

            logger.debug(f"Calling extract_csv_helper for session {session_id}")
            result, updated_state = extract_csv_helper(session_states[session_id])
            session_states[session_id] = updated_state
            return {"result": result}

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

    @app.post("/api/helpers/chart2code")
    async def api_chart2code(session_id: str = Query(...)) -> StreamingResponse:  # noqa: ANN001
        """Extract chart generation code for current session figure.

        Args:
            session_id: Session ID to get figure from.

        Returns:
            JSON with result and session_state.
        """
        logger.info(f"API chart2code request received for session_id: {session_id}")

        if session_id not in session_states:
            logger.error(f"Session not found: {session_id}")
            raise HTTPException(status_code=404, detail="Session not found")

        def _work() -> dict:
            logger.debug(f"Importing extract_code_helper for session {session_id}")
            from app import extract_code_helper

            logger.debug(f"Calling extract_code_helper for session {session_id}")
            result, updated_state = extract_code_helper(session_states[session_id])
            session_states[session_id] = updated_state
            return {"result": result}

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

    @app.post("/api/helpers/table-extract")
    async def api_table_extract(session_id: str = Query(...)) -> StreamingResponse:  # noqa: ANN001
        """Extract table data for current session figure.

        Args:
            session_id: Session ID to get figure from.

        Returns:
            JSON with result and session_state.
        """
        logger.info(f"API table-extract request received for session_id: {session_id}")

        if session_id not in session_states:
            logger.error(f"Session not found: {session_id}")
            raise HTTPException(status_code=404, detail="Session not found")

        def _work() -> dict:
            logger.debug(f"Importing extract_table_helper for session {session_id}")
            from app import extract_table_helper

            logger.debug(f"Calling extract_table_helper for session {session_id}")
            result, updated_state = extract_table_helper(session_states[session_id])
            session_states[session_id] = updated_state
            return {"result": result}

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



    @app.post("/api/helpers/describe-image")
    async def api_describe_image(session_id: str = Query(...)) -> StreamingResponse:  # noqa: ANN001
        """Generate a detailed description of the current session figure.

        Args:
            session_id: Session ID to get figure from.

        Returns:
            JSON with result and session_state.
        """
        logger.info(f"API describe-image request received for session_id: {session_id}")

        if session_id not in session_states:
            logger.error(f"Session not found: {session_id}")
            raise HTTPException(status_code=404, detail="Session not found")

        def _work() -> dict:
            logger.debug(f"Importing describe_image_helper for session {session_id}")
            from app import describe_image_helper

            logger.debug(f"Calling describe_image_helper for session {session_id}")
            result, updated_state = describe_image_helper(session_states[session_id])
            session_states[session_id] = updated_state
            return {"result": result}

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

    # ------------------------------------------------------------------
    # SSE streaming endpoints (token-by-token)
    # ------------------------------------------------------------------

    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(
        session_id: str,
        task_label: str,
        token_generator: Generator[str, None, None],
        post_process: Any | None = None,
    ) -> AsyncGenerator[str, None]:
        """Shared async generator that wraps a blocking token generator as SSE events.

        Args:
            session_id: Session to validate.
            task_label: Human-readable label for the status event.
            token_generator: Synchronous generator yielding token strings.
            post_process: Optional callable(full_text) -> str for cleanup.
        """
        if session_id not in session_states:
            yield _sse_event({"type": "error", "message": "Session not found"})
            return

        state = session_states[session_id]
        if state.get("selected_figure") is None:
            yield _sse_event({"type": "error", "message": "No figure selected"})
            return

        yield _sse_event({"type": "status", "message": task_label})

        full_text = ""
        try:
            for token in token_generator:
                full_text += token
                yield _sse_event({"type": "chunk", "content": token})
                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(f"Streaming error for session {session_id}: {e}", exc_info=True)
            yield _sse_event({"type": "error", "message": str(e)})

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

    def _get_selected_image(session_id: str) -> PILImage.Image:
        """Return the selected figure image or raise a 404."""
        if session_id not in session_states:
            raise HTTPException(status_code=404, detail="Session not found")
        fig = session_states[session_id].get("selected_figure")
        if not fig or fig.get("image") is None:
            raise HTTPException(status_code=404, detail="No figure selected")
        return fig["image"]

    @app.post("/api/helpers/chart2summary/stream")
    async def api_chart2summary_stream(session_id: str = Query(...)) -> StreamingResponse:
        from infer_vision_qa import answer_question_stream

        image = _get_selected_image(session_id)
        gen = answer_question_stream(image, "<chart2summary>", [], None)
        return StreamingResponse(
            _sse_stream(session_id, "Generating summary...", gen),
            media_type="text/event-stream",
            headers=_STREAM_HEADERS,
        )

    @app.post("/api/helpers/chart2csv/stream")
    async def api_chart2csv_stream(session_id: str = Query(...)) -> StreamingResponse:
        from infer_chart2csv import extract_csv_stream

        image = _get_selected_image(session_id)
        gen = extract_csv_stream(image)
        return StreamingResponse(
            _sse_stream(session_id, "Extracting CSV...", gen),
            media_type="text/event-stream",
            headers=_STREAM_HEADERS,
        )

    @app.post("/api/helpers/chart2code/stream")
    async def api_chart2code_stream(session_id: str = Query(...)) -> StreamingResponse:
        from app import PROMPT_TEXT_CODE
        from infer_vision_qa import answer_question_stream

        image = _get_selected_image(session_id)
        gen = answer_question_stream(image, PROMPT_TEXT_CODE, [], None)
        return StreamingResponse(
            _sse_stream(session_id, "Generating code...", gen),
            media_type="text/event-stream",
            headers=_STREAM_HEADERS,
        )

    @app.post("/api/helpers/table-extract/stream")
    async def api_table_extract_stream(session_id: str = Query(...)) -> StreamingResponse:
        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 = _get_selected_image(session_id)
        gen = answer_question_stream(image, "<tables_html>", [], None)
        return StreamingResponse(
            _sse_stream(session_id, "Extracting table...", gen, post_process=_clean_table_html),
            media_type="text/event-stream",
            headers=_STREAM_HEADERS,
        )

    @app.post("/api/helpers/describe-image/stream")
    async def api_describe_image_stream(session_id: str = Query(...)) -> StreamingResponse:
        from infer_vision_qa import answer_question_stream

        image = _get_selected_image(session_id)
        gen = answer_question_stream(image, "Describe this image in detail", [], None)
        return StreamingResponse(
            _sse_stream(session_id, "Describing image...", gen),
            media_type="text/event-stream",
            headers=_STREAM_HEADERS,
        )