# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0

"""Figure captioning for RAG ingestion.

Two captioning sources, tried in order:

1. The user's currently-loaded chat VLM — when ``vlm_url`` / ``vlm_model``
   are provided by the parent process (it probes ``llama_cpp.is_vision``
   at enqueue time). Captions go through that model's OpenAI-compatible
   ``/v1/chat/completions`` endpoint as base64 ``image_url``.

2. A helper llama-server fallback that loads the pre-cached
   ``unsloth/gemma-4-E2B-it-GGUF`` (gemma-3n family, multimodal) with
   its mmproj for vision. Spawned for the lifetime of a
   ``caption_images`` call, unloaded before return so no llama-server
   process leaks past ingestion.

Defensive: any per-image failure returns an empty string; total
captioner unavailability (no chat VLM + helper load failure) returns
empty strings for every image. The caller (``_stream_image_chunks``)
falls back to the parser's page-text caption in that case.
"""

from __future__ import annotations

import base64
from io import BytesIO
from typing import Any, Optional

from loggers import get_logger

logger = get_logger(__name__)

_PROMPT = (
    "This image is a region cropped from a PDF page that contains a "
    "single figure (schematic, chart, diagram, table, photo, or "
    "their combination). Describe the figure's structure and content "
    "in <=80 words. Focus on factual visible content: axes, labels, "
    "arrow labels, box labels, legends, visible text in the figure, "
    "and what entities are connected to what. Do not speculate beyond "
    "what is visible and do not describe the page header/footer or "
    "body paragraphs."
)
_MAX_NEW_TOKENS = 200
# Downscale large images to keep the base64 payload manageable; prefill cost
# scales with tile count not pixels, but huge inputs still bloat the JSON body.
# 1600 px on the long side matches PR #5351's chat-composer extractor.
_MAX_IMAGE_SIZE = 1600
_REQUEST_TIMEOUT_SECONDS = 120.0

# Helper VLM (used when no vision-capable chat model is loaded). Matches the
# model pre-cached by precache_helper_gguf() at startup to avoid a fresh download.
_HELPER_REPO = "unsloth/gemma-4-E2B-it-GGUF"
_HELPER_VARIANT = "UD-Q4_K_XL"
_HELPER_MODEL_NAME = "helper"


def _image_to_data_url(blob: bytes) -> str:
    from PIL import Image

    img = Image.open(BytesIO(blob)).convert("RGB")
    if max(img.size) > _MAX_IMAGE_SIZE:
        img.thumbnail((_MAX_IMAGE_SIZE, _MAX_IMAGE_SIZE))
    buf = BytesIO()
    img.save(buf, format = "JPEG", quality = 88)
    encoded = base64.b64encode(buf.getvalue()).decode("ascii")
    return f"data:image/jpeg;base64,{encoded}"


def _load_helper_vlm() -> Optional[tuple[Any, str, str]]:
    """Spawn a private LlamaCppBackend with the helper VLM + mmproj.

    Returns ``(backend, base_url, model_name)`` on success, ``None`` on
    failure. The caller is responsible for unloading the backend when
    done (so the helper doesn't outlive the ingestion subprocess).
    """
    try:
        from core.inference.llama_cpp import LlamaCppBackend

        # kill_orphans=False is critical: the global singleton already runs the
        # user's chat-model llama-server. Killing "orphans" here would reap that
        # healthy process — the orphan-killer can't tell two backends apart by PID.
        backend = LlamaCppBackend(kill_orphans = False)
        logger.info(
            "RAG captioner: loading helper VLM as fallback",
            repo = _HELPER_REPO,
            variant = _HELPER_VARIANT,
        )
        ok = backend.load_model(
            hf_repo = _HELPER_REPO,
            hf_variant = _HELPER_VARIANT,
            model_identifier = f"rag-captioner:{_HELPER_REPO}:{_HELPER_VARIANT}",
            is_vision = True,
            n_ctx = 4096,
            n_gpu_layers = -1,
        )
        if not ok:
            logger.warning("RAG captioner: helper VLM failed to start")
            return None
        return backend, backend.base_url, _HELPER_MODEL_NAME
    except Exception as exc:  # noqa: BLE001
        logger.warning("RAG captioner: helper VLM load raised", error = str(exc))
        return None


def _post_one(client: Any, endpoint: str, model: str, blob: bytes) -> str:
    """POST one image to the OpenAI-compatible endpoint, return caption."""
    data_url = _image_to_data_url(blob)
    payload = {
        "model": model,
        "messages": [
            {
                "role": "user",
                "content": [
                    {"type": "text", "text": _PROMPT},
                    {
                        "type": "image_url",
                        "image_url": {"url": data_url},
                    },
                ],
            }
        ],
        "max_tokens": _MAX_NEW_TOKENS,
        "temperature": 0.0,
        # Reasoning models (gemma-4, qwen3-thinking, etc.) spend the whole budget on
        # <thinking> and emit empty visible content — useless for a short caption.
        # Disable thinking for this request only; the user's chat sessions stay unaffected.
        "chat_template_kwargs": {"enable_thinking": False},
    }
    response = client.post(endpoint, json = payload)
    response.raise_for_status()
    data = response.json()
    content = data.get("choices", [{}])[0].get("message", {}).get("content", "")
    return content.strip() if isinstance(content, str) else ""


def caption_images(
    image_bytes_list: list[bytes],
    *,
    vlm_url: Optional[str] = None,
    vlm_model: Optional[str] = None,
) -> list[str]:
    """Generate one short caption per image; same-length output.

    Tries the loaded chat VLM first (``vlm_url`` + ``vlm_model``). If
    those are missing, spawns the helper VLM, captions, and unloads it
    before returning. On any failure returns ``""`` for the affected
    image. Never raises.
    """
    logger.info(
        "caption_images: invoked",
        n_images = len(image_bytes_list),
        vlm_url = vlm_url,
        vlm_model = vlm_model,
    )
    if not image_bytes_list:
        return []

    import httpx

    helper_backend: Optional[Any] = None
    try:
        # Resolve endpoint + model: chat VLM if available, else helper.
        if vlm_url and vlm_model:
            endpoint = f"{vlm_url.rstrip('/')}/v1/chat/completions"
            model_name = vlm_model
        else:
            loaded = _load_helper_vlm()
            if loaded is None:
                logger.warning(
                    "caption_images: helper load failed, returning empty captions"
                )
                return ["" for _ in image_bytes_list]
            helper_backend, helper_base_url, helper_model_name = loaded
            endpoint = f"{helper_base_url.rstrip('/')}/v1/chat/completions"
            model_name = helper_model_name

        out: list[str] = []
        with httpx.Client(timeout = _REQUEST_TIMEOUT_SECONDS) as client:
            for idx, blob in enumerate(image_bytes_list):
                try:
                    out.append(_post_one(client, endpoint, model_name, blob))
                except Exception as exc:  # noqa: BLE001
                    logger.warning(
                        "caption_images: per-image request failed",
                        idx = idx,
                        endpoint = endpoint,
                        error = str(exc),
                    )
                    out.append("")
        non_empty = sum(1 for c in out if c.strip())
        logger.info(
            "caption_images: complete",
            total = len(out),
            non_empty = non_empty,
        )
        return out
    finally:
        # Tear down the helper if we spawned one. The parent-provided chat VLM is
        # left alone — it's not ours to manage.
        if helper_backend is not None:
            try:
                helper_backend.unload_model()
                logger.info("RAG captioner: helper VLM unloaded")
            except Exception as exc:  # noqa: BLE001
                logger.warning("RAG captioner: helper unload failed", error = str(exc))