"""Pricing of inline images, measured against qwen3.6.

The numbers in the parametrized table are the ones measured on the local
vision model (``prompt_eval_count`` of the same prompt with and without one
image); the estimator must reproduce them, because every window decision —
the compression trigger, the per-turn image budget, ``check_context_size`` —
is taken in these units.
"""

import base64
import struct

import pytest

from navi.config import Settings
from navi.core import image_tokens as image_tokens_module
from navi.core.image_tokens import (
    context_image_tokens,
    estimate_image_tokens,
    tokens_for_size,
)
from navi.llm.base import Message


def _b64(raw: bytes) -> str:
    return base64.b64encode(raw).decode()


def _png_header(width: int, height: int) -> bytes:
    return (
        b"\x89PNG\r\n\x1a\n"
        + b"\x00\x00\x00\x0dIHDR"
        + struct.pack(">II", width, height)
        + b"\x08\x02\x00\x00\x00"
    )


def _jpeg_header(width: int, height: int) -> bytes:
    return (
        b"\xff\xd8"                       # SOI
        + b"\xff\xe0\x00\x04\x00\x00"    # APP0, 2 payload bytes
        + b"\xff\xc0\x00\x11\x08"        # SOF0, 17-byte segment, 8-bit
        + struct.pack(">HH", height, width)
        + b"\x03\x01\x22\x00\x02\x11\x01\x03\x11\x01"
    )


class TestTokensForSize:
    @pytest.mark.parametrize(
        ("width", "height", "expected"),
        [
            (512, 512, 258),
            (800, 600, 477),
            (1024, 768, 770),
            (1024, 1024, 1026),
            (1536, 1536, 2306),
            (3072, 1024, 3074),
            (4096, 512, 2050),
        ],
    )
    def test_matches_the_measured_costs(self, width, height, expected):
        assert tokens_for_size(width, height) == expected

    def test_a_downscaled_image_costs_the_same_as_a_full_size_one(self):
        # 3072x3072 is scaled to 2048x2048 before tokenizing (4098, capped).
        assert tokens_for_size(3072, 3072) == tokens_for_size(2048, 2048) == 4096
        assert tokens_for_size(4096, 1024) == 4096

    def test_a_tiny_icon_is_floored_not_priced_at_almost_nothing(self):
        # The formula alone says 66; the floor keeps small images visible in
        # the estimate instead of letting a wall of icons cost nothing.
        assert tokens_for_size(256, 256) == 256
        assert tokens_for_size(1, 1) == 256

    def test_a_degenerate_size_falls_back_to_the_floor(self):
        assert tokens_for_size(0, 100) == 256
        assert tokens_for_size(-5, -5) == 256


class TestHappensOnRealEncodedImages:
    @pytest.mark.parametrize(
        ("fmt", "size", "expected"),
        [
            ("PNG", (1024, 768), 770),
            ("JPEG", (512, 512), 258),
        ],
    )
    def test_a_genuinely_encoded_image_is_priced_by_its_header(self, fmt, size, expected):
        import io

        from PIL import Image

        buf = io.BytesIO()
        Image.new("RGB", size, (10, 20, 30)).save(buf, format=fmt)
        assert estimate_image_tokens(_b64(buf.getvalue())) == expected

    def test_a_data_uri_prefix_is_accepted(self):
        raw = _png_header(1024, 1024)
        assert estimate_image_tokens("data:image/png;base64," + _b64(raw)) == 1026


class TestHeaderParsers:
    def test_png(self):
        assert estimate_image_tokens(_b64(_png_header(800, 600))) == 477

    def test_jpeg(self):
        assert estimate_image_tokens(_b64(_jpeg_header(1024, 768))) == 770

    def test_gif(self):
        raw = b"GIF89a" + struct.pack("<HH", 512, 512) + b"\x00" * 8
        assert estimate_image_tokens(_b64(raw)) == 258

    def test_bmp(self):
        raw = b"BM" + b"\x00" * 16 + struct.pack("<ii", 1024, -768)
        assert estimate_image_tokens(_b64(raw)) == 770

    def test_webp_extended_header(self):
        raw = (
            b"RIFF" + b"\x00" * 4 + b"WEBP" + b"VP8X" + b"\x00" * 4
            + b"\x00" * 4
            + (512 - 1).to_bytes(3, "little")
            + (512 - 1).to_bytes(3, "little")
        )
        assert estimate_image_tokens(_b64(raw)) == 258


class TestFallbackWhenTheHeaderCannotBeRead:
    def test_an_unknown_format_is_charged_the_setting(self, monkeypatch):
        monkeypatch.setattr(
            image_tokens_module,
            "settings",
            Settings(_env_file=None, context_image_tokens=777),
        )
        assert estimate_image_tokens("this is not an encoded image") == 777

    def test_malformed_base64_does_not_raise(self, monkeypatch):
        monkeypatch.setattr(
            image_tokens_module,
            "settings",
            Settings(_env_file=None, context_image_tokens=640),
        )
        assert estimate_image_tokens("!!!not-base64!!!") == 640

    def test_the_fallback_is_capped_at_a_whole_window_scale(self, monkeypatch):
        monkeypatch.setattr(
            image_tokens_module,
            "settings",
            Settings(_env_file=None, context_image_tokens=10_000_000),
        )
        assert estimate_image_tokens("still not an image") == 4096

    def test_an_empty_payload_costs_nothing(self):
        assert estimate_image_tokens("") == 0


class TestContextImageTokens:
    def test_sums_every_image_of_every_message(self):
        png = _b64(_png_header(1024, 1024))  # 1026
        jpeg = _b64(_jpeg_header(512, 512))  # 258
        messages = [
            Message(role="user", content="a", images=[png]),
            Message(role="user", content="b", images=[png, jpeg]),
            Message(role="user", content="no images"),
        ]
        assert context_image_tokens(messages) == 1026 * 2 + 258

    def test_no_images_is_zero(self):
        assert context_image_tokens([Message(role="user", content="hi")]) == 0