"""What an inline image costs the model, measured instead of guessed.

Every size decision in the compressor prices images at a flat 500 tokens.
On the local vision model that constant is off by 2x for an ordinary
screenshot and by 8x for a large photo, and the error is invisible: image
bytes are not in ``context_chars``, so nothing else in the estimate sees
them. A session that carries a handful of images therefore sits below the
90% trigger while the prompt it sends is already past the window — and once
Ollama truncates that prompt, the reported ``prompt_eval_count`` is what gets
cached as the baseline, so the count never recovers.

The measurement (Ollama 0.40.15, ``qwen3.6:35b``, ``prompt_eval_count`` of the
same prompt with and without an image, two calls each)::

    256x256   ->   66 tokens       1024x768  ->  770 tokens
    512x512   ->  258 tokens       1536x1536 -> 2306 tokens
    1024x1024 -> 1026 tokens       2048x2048 -> 4098 tokens
    800x600   ->  477 tokens       3072x3072 -> 4098 tokens (downscaled)

That is one token per 32x32 pixel block plus the two image markers, and a
downscale to at most 2048x2048 pixels before tokenizing. The estimate below
follows the same shape with *ceil* rounding, so it errs high rather than low:
over-counting costs an early compression, under-counting costs the session.

The per-32px tile size is a property of the model's vision encoder, not a
universal law (llava-style encoders charge a flat ~256 for any image, and a
cloud model may charge differently). A model that prices images unlike
qwen3.6 will therefore be estimated a little high; the real-prompt baseline
corrects the total after the first successful call, and this estimate only
supplies the delta — so the residual is small and always in the safe
direction.
"""

from __future__ import annotations

import base64
import math
import struct
from functools import lru_cache

from navi.config import settings

# One vision token per this many pixels per side (measured on qwen3.6).
_IMAGE_PATCH_PX = 32
# Ollama downscales anything larger to this budget before tokenizing.
_MAX_IMAGE_PIXELS = 2048 * 2048
# A floor keeps a tiny icon from pricing at ~0; a ceiling keeps a corrupt
# header from pricing an image at a whole window.
_MIN_IMAGE_TOKENS = 256
_MAX_IMAGE_TOKENS = 4096

# Enough of the base64 payload to hold a PNG IHDR, a WebP chunk header or a
# JPEG SOF marker for the images this server produces (canvas, PIL — neither
# keeps EXIF). Only this prefix is kept as the cache key, never the payload.
_HEAD_CHARS = 4096

_NO_DIMS: tuple[int, int] | None = None


def context_image_tokens(messages) -> int:
    """Total image cost of a message list, in real-prompt token units."""
    return sum(
        estimate_image_tokens(img)
        for m in messages
        for img in (m.images or [])
    )


def estimate_image_tokens(data: str) -> int:
    """Token cost of one base64 image (data-URI prefix accepted)."""
    if not data:
        return 0
    return _cost(len(data), data[:_HEAD_CHARS], int(settings.context_image_tokens))


@lru_cache(maxsize=256)
def _cost(encoded_len: int, head: str, fallback: int) -> int:
    """Cached by (length, header prefix) so one estimate does not re-parse.

    Two different images of the same encoded length and the same first 4 KB
    are, for every format this server writes, the same image — and even a
    collision only moves the estimate by a few hundred tokens.
    """
    dims = _dimensions(head)
    if dims is None:
        return max(1, min(int(fallback), _MAX_IMAGE_TOKENS))
    return tokens_for_size(*dims)


def tokens_for_size(width: int, height: int) -> int:
    """Estimated tokens for an image of ``width`` x ``height`` pixels."""
    if width <= 0 or height <= 0:
        return _MIN_IMAGE_TOKENS
    pixels = width * height
    if pixels > _MAX_IMAGE_PIXELS:
        scale = math.sqrt(_MAX_IMAGE_PIXELS / pixels)
        width = max(1, int(width * scale))
        height = max(1, int(height * scale))
    tiles = _ceil_div(width, _IMAGE_PATCH_PX) * _ceil_div(height, _IMAGE_PATCH_PX)
    return max(_MIN_IMAGE_TOKENS, min(_MAX_IMAGE_TOKENS, tiles + 2))


def _ceil_div(value: int, divisor: int) -> int:
    return -(-value // divisor)


def _dimensions(data: str) -> tuple[int, int] | None:
    raw = _prefix_bytes(data)
    if not raw:
        return _NO_DIMS
    for parser in (_png_size, _jpeg_size, _webp_size, _gif_size, _bmp_size):
        dims = parser(raw)
        if dims is not None:
            return dims
    return _NO_DIMS


def _prefix_bytes(data: str, max_bytes: int = 48 * 1024) -> bytes:
    """Decode the head of a base64 payload — enough for every header we read."""
    s = data.strip()
    if s.startswith("data:"):
        comma = s.find(",")
        if comma == -1:
            return b""
        s = s[comma + 1:]
    take = min(len(s), _HEAD_CHARS)
    chunk = "".join(s[:take].split())
    chunk = chunk[: len(chunk) - len(chunk) % 4]
    if not chunk:
        return b""
    try:
        return base64.b64decode(chunk)[:max_bytes]
    except ValueError:  # binascii.Error (bad padding / alphabet) is a ValueError
        return b""


def _png_size(b: bytes) -> tuple[int, int] | None:
    if len(b) < 24 or b[:8] != b"\x89PNG\r\n\x1a\n" or b[12:16] != b"IHDR":
        return None
    return struct.unpack(">II", b[16:24])


def _jpeg_size(b: bytes) -> tuple[int, int] | None:
    if len(b) < 4 or b[:2] != b"\xff\xd8":
        return None
    i, n = 2, len(b)
    while i + 4 <= n:
        if b[i] != 0xFF:
            i += 1
            continue
        marker = b[i + 1]
        if marker == 0xFF:  # fill byte before the real marker
            i += 1
            continue
        if marker in (0x01, 0xD8) or 0xD0 <= marker <= 0xD7:
            i += 2  # markers with no length field
            continue
        if marker in (0xD9, 0xDA):  # end of image / scan: no SOF before this
            return None
        seg_len = int.from_bytes(b[i + 2:i + 4], "big")
        if seg_len < 2:
            return None
        # SOF0..SOF15, minus DHT (C4), JPG (C8) and DAC (CC).
        if 0xC0 <= marker <= 0xCF and marker not in (0xC4, 0xC8, 0xCC):
            if i + 9 > n:
                return None
            height = int.from_bytes(b[i + 5:i + 7], "big")
            width = int.from_bytes(b[i + 7:i + 9], "big")
            return width, height
        i += 2 + seg_len
    return None


def _webp_size(b: bytes) -> tuple[int, int] | None:
    if len(b) < 30 or b[:4] != b"RIFF" or b[8:12] != b"WEBP":
        return None
    fmt = b[12:16]
    if fmt == b"VP8X":
        width = 1 + int.from_bytes(b[24:27], "little")
        height = 1 + int.from_bytes(b[27:30], "little")
        return width, height
    if fmt == b"VP8 ":
        if b[23:26] != b"\x9d\x01\x2a":
            return None
        width = int.from_bytes(b[26:28], "little") & 0x3FFF
        height = int.from_bytes(b[28:30], "little") & 0x3FFF
        return width, height
    if fmt == b"VP8L":
        if b[20] != 0x2F:
            return None
        bits = int.from_bytes(b[21:25], "little")
        return (bits & 0x3FFF) + 1, ((bits >> 14) & 0x3FFF) + 1
    return None


def _gif_size(b: bytes) -> tuple[int, int] | None:
    if len(b) < 10 or b[:6] not in (b"GIF87a", b"GIF89a"):
        return None
    return struct.unpack("<HH", b[6:10])


def _bmp_size(b: bytes) -> tuple[int, int] | None:
    if len(b) < 26 or b[:2] != b"BM":
        return None
    width, height = struct.unpack("<ii", b[18:26])
    return abs(width), abs(height)