"""Which models can see, and what happens when the one doing the job cannot.

A job's model is chosen for the job, not for the pictures that turn up in it.
The extraction and tutor models configured on this deployment report
`supports_vision: false`, and handing one an `image_url` part does not produce a
worse answer — it ends the request with a proxy error. So the picture goes to a
model that can see, the one an administrator sets for the `tool` job, and its
description is folded into the prompt as text, which a model that cannot see
reads perfectly well.

Capability is asked of the proxy rather than inferred from the model id: the
catalogue here runs to several hundred entries and changes without us.
`/model/info` carries `supports_vision` per deployment and settles most of it
outright — every OpenAI and Anthropic route, and a flat `false` on the DeepSeek
one. Where the field is absent the proxy genuinely does not know, and guessing
"no" would send work to the tool model that never needed to go there, so the
model itself is asked with an eight-pixel image, once.

Both answers are cached, because the question is about the model and not about
the request: the catalogue in-process for a few minutes, the probe in Redis for
a week. Descriptions are cached too — the tutor re-sends the same figure on
every turn of a conversation, and paying a second model call for each of them
was the first thing that made this feature look slow.

Nothing here degrades quietly. If the job's model cannot see and no tool model
is configured, the caller gets an error naming the setting to change; a
text-only answer about an image nobody looked at is worse than no answer.
"""
import base64
import hashlib
import io
import logging
import mimetypes
import re
import time
from dataclasses import dataclass, field

from app.config import settings
from app.services.ai_service import chat, get_configured_model

logger = logging.getLogger(__name__)

#: How long the proxy's catalogue is trusted. Models are added to it by hand,
#: so minutes are plenty and a stale "cannot see" is only ever a slower answer.
CATALOGUE_TTL = 600

#: A probe verdict is a fact about the model, not about today, so it is kept
#: long enough to be worth having and short enough to survive a proxy rewiring.
PROBE_TTL = 7 * 86400

DESCRIPTION_TTL = 30 * 86400

#: Beyond this the base64 payload costs more than the detail is worth — a stem
#: image here is a 2–4 MB scanned radiograph, and 1600px of it answers the same
#: question at a tenth the size.
MAX_IMAGE_BYTES = 1_500_000
DOWNSCALE_WIDTH = 1600

#: What the tool model is asked. American, like the rest of the clinical copy.
DESCRIBE_PROMPT = (
    "Describe this image for a colleague who cannot see it and must reason from "
    "your words alone.\n"
    "- Report what is visible: modality, body part, structures, colors, "
    "measurements, axes, arrows, and any text printed on the image, quoted exactly.\n"
    "- Describe abnormal findings in clinical terms. Do not name a diagnosis "
    "unless the image itself is labeled with one.\n"
    "- Say plainly if the image is a logo, page header, blank, or otherwise "
    "carries no clinical content.\n"
    "- Description only: no preamble, no answer to any question."
)

_catalogue: dict[str, bool | None] = {}
_catalogue_at: float = 0.0
_probes: dict[str, bool] = {}


class VisionUnavailable(RuntimeError):
    """No model available to look at an image the request depends on.

    Raised rather than returned so that no caller can carry on without noticing
    that the picture went unread. The message names the setting to change.
    """


@dataclass(frozen=True)
class Image:
    data: bytes
    media_type: str = "image/jpeg"
    caption: str = ""


@dataclass
class Handoff:
    """What actually ran, so nobody has to infer a second model call from a
    latency graph."""

    primary_model: str
    images: int = 0
    delegated: bool = False
    tool_model: str | None = None
    reason: str = ""
    cached: bool = False
    elapsed_ms: int = 0
    descriptions: list[str] = field(default_factory=list)

    def as_dict(self) -> dict:
        return {
            "delegated": self.delegated,
            "primary_model": self.primary_model,
            "tool_model": self.tool_model,
            "images": self.images,
            "reason": self.reason,
            "cached": self.cached,
            "elapsed_ms": self.elapsed_ms,
        }


def _redis():
    try:
        import redis as redis_lib

        return redis_lib.from_url(settings.REDIS_URL, decode_responses=True,
                                  socket_connect_timeout=1)
    except Exception:
        return None


def _cached(key: str) -> str | None:
    client = _redis()
    if client is None:
        return None
    try:
        return client.get(key)
    except Exception:
        return None


def _remember(key: str, value: str, ttl: int) -> None:
    client = _redis()
    if client is None:
        return
    try:
        client.set(key, value, ex=ttl)
    except Exception:
        # A cache that is down is a cost, never a failure.
        logger.debug("Could not cache %s", key, exc_info=True)


def catalogue() -> dict[str, bool | None]:
    """model name -> what the proxy says about vision, refreshed on a timer.

    A name can appear several times, once per deployment behind it: `best-chat`
    fans out to six. A single `false` among them decides the alias, because a
    request can land on any of them and an alias that fails one time in six is
    worse than one that always takes the slower path.
    """
    global _catalogue, _catalogue_at
    if _catalogue and time.monotonic() - _catalogue_at < CATALOGUE_TTL:
        return _catalogue

    import httpx

    base = (settings.LITELLM_API_BASE or "").rstrip("/").removesuffix("/v1")
    if not base:
        return {}
    headers = {"Authorization": f"Bearer {settings.LITELLM_API_KEY}"} if settings.LITELLM_API_KEY else {}
    try:
        response = httpx.get(f"{base}/model/info", headers=headers, timeout=10)
        response.raise_for_status()
        rows = response.json().get("data", [])
    except Exception:
        # Leave whatever is already known in place: an unreachable catalogue is
        # a reason to fall back to probing, not to forget what it said before.
        logger.warning("Could not read the model catalogue; vision capability "
                       "falls back to probing", exc_info=True)
        return _catalogue

    verdicts: dict[str, bool | None] = {}
    for row in rows:
        name = row.get("model_name")
        if not name:
            continue
        says = (row.get("model_info") or {}).get("supports_vision")
        known = verdicts.get(name, "absent")
        if says is False or known is False:
            verdicts[name] = False
        elif says is True:
            verdicts[name] = True
        elif known == "absent":
            verdicts[name] = None
    _catalogue, _catalogue_at = verdicts, time.monotonic()
    return _catalogue


def _probe_image() -> str:
    """An eight-pixel PNG as a data URL. Deliberately not one pixel: some
    providers reject an image below a minimum dimension, and a refusal of the
    probe would read as a refusal of images."""
    from PIL import Image as PILImage

    buffer = io.BytesIO()
    PILImage.new("RGB", (8, 8), (255, 255, 255)).save(buffer, format="PNG")
    return data_url(buffer.getvalue(), "image/png")


def _probe(model_id: str, api_key: str | None) -> bool:
    """Ask the model itself, for the models the catalogue has no opinion on.

    A 4xx is the proxy or the provider refusing the shape of the request, which
    for a one-token call carrying nothing but a white square means it will not
    take images; that verdict is worth keeping. A timeout or a 5xx says nothing
    about the model, so it is not cached and the caller takes the safe path
    this once.
    """
    cache_key = f"vision:probe:{model_id}"
    if model_id in _probes:
        return _probes[model_id]
    remembered = _cached(cache_key)
    if remembered is not None:
        _probes[model_id] = remembered == "1"
        return _probes[model_id]

    try:
        chat(
            model=model_id,
            max_tokens=1,
            messages=[{"role": "user", "content": [
                {"type": "image_url", "image_url": {"url": _probe_image()}},
                {"type": "text", "text": "Reply with: ok"},
            ]}],
            api_key=api_key,
            timeout=30,
        )
        verdict = True
    except Exception as error:
        status = getattr(error, "status_code", None)
        if not (isinstance(status, int) and 400 <= status < 500):
            logger.info("Vision probe of %s was inconclusive: %s", model_id, error)
            return False
        logger.info("Vision probe: %s refused an image (%s)", model_id, error)
        verdict = False

    _probes[model_id] = verdict
    _remember(cache_key, "1" if verdict else "0", PROBE_TTL)
    return verdict


def can_see(model_id: str, api_key: str | None = None) -> bool:
    """Whether this model may be handed an image. Cached both ways."""
    if not model_id:
        return False
    verdict = catalogue().get(model_id, None)
    if verdict is not None:
        return verdict
    return _probe(model_id, api_key)


def data_url(data: bytes, media_type: str) -> str:
    return f"data:{media_type};base64,{base64.b64encode(data).decode()}"


#: How a JPEG 2000 file starts — the JP2 container, then a bare codestream.
#: Sniffed rather than taken from the name, because the name is no guide: the
#: stem images extracted from these PDFs are `.jpx`, the slim base image ships
#: no MIME table so `mimetypes` reports nothing at all for them, and the
#: default that fills the gap is `image/jpeg`. The provider then answers "the
#: image data you provided does not represent a valid image" some seconds and
#: two fallback hops later.
JPEG_2000_SIGNATURES = (b"\x00\x00\x00\x0cjP  ", b"\xff\x4f\xff\x51")


def prepare(data: bytes, media_type: str | None = None, caption: str = "") -> Image | None:
    """Bytes as a vision API will take them, or None if they cannot be.

    Two conversions, both learned the hard way: JPEG 2000 comes out of the PDF
    pipeline and no vision API accepts it, and a full-resolution scan spends
    more on base64 than the extra pixels are worth.
    """
    if not data:
        return None
    media_type = media_type or "image/jpeg"
    if data.startswith(JPEG_2000_SIGNATURES):
        media_type = "image/jp2"
    if media_type in ("image/jp2", "image/jpx"):
        try:
            import fitz  # already a dependency of the PDF pipeline

            data, media_type = fitz.Pixmap(data).tobytes("png"), "image/png"
        except Exception:
            logger.info("Could not re-encode a JPEG 2000 image for a vision model",
                        exc_info=True)
            return None
    if len(data) > MAX_IMAGE_BYTES:
        from app.services import thumbnails

        smaller = thumbnails.render(data, DOWNSCALE_WIDTH)
        if smaller:
            data, media_type = smaller, "image/webp"
    return Image(data=data, media_type=media_type, caption=caption)


def load_image(key: str, caption: str = "") -> Image | None:
    """A stored upload, ready for a model. Never raises."""
    if not key:
        return None
    try:
        if re.match(r"^https?://", key, re.I):
            import httpx

            response = httpx.get(key, timeout=10, follow_redirects=True)
            response.raise_for_status()
            return prepare(response.content,
                           response.headers.get("content-type", "").split(";")[0] or None,
                           caption)
        from app.services import storage_service

        # Some rows hold the URL the browser uses rather than the storage key;
        # `/uploads/` is that route's prefix and nothing in the bucket has it.
        key = key.removeprefix("/uploads/")
        data = storage_service.s3_object(key) or storage_service.load(key)
        return prepare(data, mimetypes.guess_type(key)[0], caption) if data else None
    except Exception:
        logger.info("Could not load %s for a vision model", key, exc_info=True)
        return None


def image_part(image: Image) -> dict:
    return {"type": "image_url", "image_url": {"url": data_url(image.data, image.media_type)}}


def word_image(word: str) -> Image:
    """A picture of a word, for testing that a model can read a picture.

    Presence in the catalogue is not proof, the same way a transcription model
    listed by the proxy is not proof that anything can be transcribed. The test
    makes the tool model read something only an eye could have read.
    """
    from PIL import Image as PILImage, ImageDraw, ImageFont

    canvas = PILImage.new("RGB", (480, 160), (255, 255, 255))
    draw = ImageDraw.Draw(canvas)
    try:
        font = ImageFont.load_default(size=64)
    except TypeError:  # Pillow before 9.2 cannot size the default font
        font = ImageFont.load_default()
    draw.text((40, 45), word, fill=(0, 0, 0), font=font)
    buffer = io.BytesIO()
    canvas.save(buffer, format="PNG")
    return Image(data=buffer.getvalue(), media_type="image/png", caption="test image")


def describe(images: list[Image], model_id: str, api_key: str | None = None,
             context: str = "", use_cache: bool = True) -> list[str]:
    """What the tool model sees, one description per image.

    Cached on the bytes, the surrounding context and the model, because the
    same figure comes back on every turn of a tutor conversation and none of
    those turns changes what the picture shows.
    """
    return _describe(images, model_id, api_key, context, use_cache)[0]


def _describe(images: list[Image], model_id: str, api_key: str | None,
              context: str, use_cache: bool) -> tuple[list[str], int]:
    """The descriptions, and how many of them came from the cache."""
    out: list[str] = []
    hits = 0
    for image in images:
        digest = hashlib.sha256(
            image.data + context.encode() + model_id.encode()).hexdigest()[:24]
        cache_key = f"vision:description:{digest}"
        if use_cache:
            remembered = _cached(cache_key)
            if remembered:
                out.append(remembered)
                hits += 1
                continue

        prompt = DESCRIBE_PROMPT
        if image.caption:
            prompt += f"\n\nWhat this image is: {image.caption}"
        if context:
            prompt += f"\n\nWhere it appears:\n{context[:2000]}"

        text = (chat(
            model=model_id,
            temperature=0,
            max_tokens=700,
            messages=[{"role": "user", "content": [
                image_part(image), {"type": "text", "text": prompt},
            ]}],
            api_key=api_key,
            timeout=120,
        ) or "").strip()
        if not text:
            raise VisionUnavailable(
                f"{model_id} is configured as the tool model but returned no "
                "description of the image.")
        out.append(text)
        if use_cache:
            _remember(cache_key, text, DESCRIPTION_TTL)
    return out, hits


def image_context(db, images: list[Image], *, model_id: str, api_key: str | None = None,
                  context: str = "", tool: tuple[str, str | None] | None = None
                  ) -> tuple[list[dict], Handoff]:
    """Message content parts carrying these images, whatever the model can do.

    The caller splices the parts into a user message and does not need to know
    which of the two things happened: either the images themselves, or a
    description of each written by the tool model. The Handoff says which, for
    the log line and for the response.

    `tool` names the tool model outright, for callers running in a thread pool:
    a SQLAlchemy session belongs to one thread, so those resolve it once on the
    way in and pass it down rather than handing `db` to every worker.
    """
    handoff = Handoff(primary_model=model_id, images=len(images))
    if not images:
        return [], handoff

    started = time.monotonic()
    if can_see(model_id, api_key):
        handoff.reason = "the model reads images itself"
        handoff.elapsed_ms = int((time.monotonic() - started) * 1000)
        return [image_part(image) for image in images], handoff

    tool = tool or get_configured_model(db, "tool")
    if not tool:
        raise VisionUnavailable(
            f"{model_id} cannot read images and no tool model is configured. "
            "In Settings → AI models, allow a model that can see for the Tool "
            "job and select it.")
    tool_model, tool_key = tool
    if catalogue().get(tool_model) is False:
        raise VisionUnavailable(
            f"{model_id} cannot read images, and the tool model set to cover it "
            f"({tool_model}) cannot either. Choose a model that can see for the "
            "Tool job in Settings → AI models.")

    handoff.delegated = True
    handoff.tool_model = tool_model
    handoff.reason = f"{model_id} cannot read images"
    try:
        handoff.descriptions, hits = _describe(images, tool_model, tool_key, context, True)
    except VisionUnavailable:
        raise
    except Exception as error:
        raise VisionUnavailable(
            f"The tool model {tool_model} could not read the image: {error}") from error
    handoff.elapsed_ms = int((time.monotonic() - started) * 1000)
    handoff.cached = hits == len(images)

    parts = []
    for image, text in zip(images, handoff.descriptions):
        what = image.caption or "an image accompanying this request"
        parts.append({"type": "text", "text": (
            f"[Description of {what}. You cannot see images, so {tool_model} "
            f"looked at it and wrote this. Treat it as what the image shows.]\n{text}")})

    logger.info("Vision handoff: %s cannot see, %s described %d image(s) in %d ms",
                model_id, tool_model, len(images), handoff.elapsed_ms)
    return parts, handoff