"""A provider's declared image edge holds for every image it is sent.

A profile's `imageEdge` stands in for Claude's 2000px limit, which only the
real Anthropic upstream declares.
"""

import base64
import hashlib
import json
import re
import sqlite3
import struct
import unittest
import urllib.error
import zlib

from harness import (
    LOOPBACK_REMOTE,
    Albedo,
    Provider,
    exclusive,
    operation_id,
    python,
    text,
)


def png(width, height, metadata=b""):
    def chunk(kind, data):
        body = kind + data
        return struct.pack(">I", len(data)) + body + struct.pack(">I", zlib.crc32(body))

    rows = b"".join(b"\x00" + b"\x80\x40\x20" * width for _ in range(height))
    return (
        b"\x89PNG\r\n\x1a\n"
        + chunk(b"IHDR", struct.pack(">IIBBBBB", width, height, 8, 2, 0, 0, 0))
        + (chunk(b"npAD", metadata) if metadata else b"")
        + chunk(b"IDAT", zlib.compress(rows))
        + chunk(b"IEND", b"")
    )


WIDE = png(1200, 10)
WIDE_DATA = base64.b64encode(WIDE).decode()
NOTE = "An earlier image was scaled from 1200x10 to 1000x"


def widths(request):
    """The pixel width of every PNG the request carries."""
    found = re.findall(r"data:image/png;base64,([A-Za-z0-9+/=]+)", json.dumps(request))
    return [struct.unpack(">I", base64.b64decode(data)[16:20])[0] for data in found]


def cell_result(output):
    """The python tool's result: plain text, or the text part beside images."""
    body = output["output"]
    if isinstance(body, str):
        return json.loads(body)
    [text] = [part["text"] for part in body if part.get("type") == "input_text"]
    return json.loads(text)


class ImageLimitsTest(unittest.TestCase):
    def test_overflow_images_are_saved_and_can_be_shown_in_another_cell(self):
        small = base64.b64encode(png(2, 3)).decode()
        code = f"""
image = base64.b64decode({small!r})
notice = str(show_image(image, image, image, image, image, image))
print(notice)
paths = re.findall(r"saved to (.*?); show_image", notice)
assert len(paths) == 2, notice
assert all(Path(path).read_bytes() == image for path in paths)
print('six images preserved')
"""
        calls = 0

        def reply(request):
            nonlocal calls
            calls += 1
            if calls == 1:
                return python(code)
            if calls == 2:
                return python("show_image(*paths)")
            return text("done")

        provider = Provider(reply)
        self.addCleanup(provider.close)
        with Albedo(provider, protocol="responses") as app:
            session = app.session()
            app.prompt(session, "show six images").close()
            app.idle(session)
            outputs = [
                item
                for item in provider.requests[-1]["request"]["input"]
                if item.get("type") == "function_call_output"
            ]
            self.assertEqual(len(outputs), 2)
            results = [cell_result(item) for item in outputs]
            self.assertTrue(all(r["status"] == "ok" for r in results), results)
            self.assertIn("six images preserved", results[0]["output"])
            self.assertIn("in another cell", results[0]["output"])
            counts = [
                len(re.findall("data:image/png;base64,", json.dumps(item)))
                for item in outputs
            ]
            self.assertEqual(counts, [4, 2])

    def test_image_byte_budget_saves_overflow_without_failing_the_cell(self):
        import inspect

        code = (
            "import struct, zlib\n"
            + inspect.getsource(png)
            + """
image = png(2, 3, b'x' * (3 * 1024 * 1024))
notice = str(show_image(image, image))
print(notice)
path = re.search(r"saved to (.*?); show_image", notice)[1]
assert Path(path).read_bytes() == image
large = png(2, 3, b'x' * (6 * 1024 * 1024))
original = Path('large.png')
original.write_bytes(large)
for source in (original, large):
    try:
        show_image(source)
    except ValueError as error:
        message = str(error)
        assert str(len(large)) in message, message
        assert '5 MiB' in message, message
        assert 'downscale or crop' in message, message
    else:
        raise AssertionError('oversize image must not suggest showing it again')
assert original.read_bytes() == large
print('byte overflow preserved')
"""
        )

        def reply(request):
            if request["input"][-1].get("type") == "function_call_output":
                return text("done")
            return python(code)

        provider = Provider(reply)
        self.addCleanup(provider.close)
        with Albedo(provider, protocol="responses") as app:
            session = app.session()
            app.prompt(session, "show large images").close()
            app.idle(session)
            output = next(
                item
                for item in provider.requests[-1]["request"]["input"]
                if item.get("type") == "function_call_output"
            )
            result = cell_result(output)
            self.assertEqual(result["status"], "ok", result)
            self.assertIn("byte overflow preserved", result["output"])
            self.assertEqual(
                len(re.findall("data:image/png;base64,", json.dumps(output))), 1
            )

    @exclusive
    def test_large_valid_image_remains_available_after_durable_reload(self):
        # A valid ancillary chunk makes encoded length exceed 5 MiB while the
        # decoded PNG remains below the advertised file-size boundary.
        image = png(2, 3, b"x" * (4 * 1024 * 1024))
        data = base64.b64encode(image).decode()
        self.assertLess(len(image), 5 * 1024 * 1024)
        self.assertGreater(len(data), 5 * 1024 * 1024)
        provider = Provider(lambda _: text("image accepted"))
        self.addCleanup(provider.close)
        with Albedo(provider, protocol="responses") as app:
            session = app.session()
            app.api(
                f"/sessions/{session}/inputs/{operation_id()}",
                {
                    "kind": "message",
                    "text": "look",
                    "images": [{"mime_type": "image/png", "data": data}],
                },
                method="PUT",
            ).close()
            app.idle(session)
            app.restart()
            app.prompt(session, "look again").close()
            app.idle(session)
            sent = json.dumps(provider.requests[-1]["request"])
            encoded = re.findall(r"data:image/png;base64,([A-Za-z0-9+/=]+)", sent)
            self.assertEqual(len(encoded), 1)
            self.assertEqual(base64.b64decode(encoded[0]), image)

    def test_a_declared_edge_refuses_new_images_and_fits_history(self):
        def reply(request):
            inputs = request["input"]
            if inputs[-1].get("type") == "function_call_output":
                return text("done")
            prompts = [
                item.get("content")
                for item in inputs
                if item.get("role") == "user"
                and "<system-note" not in json.dumps(item.get("content"))
            ]
            if prompts[-1] == "show":
                return python(
                    "import base64\nshow_image(base64.b64decode('" + WIDE_DATA + "'))"
                )
            return text("ok")

        provider = Provider(reply)
        loose, strict = f"loose-{provider.route}", f"strict-{provider.route}"
        profile = {
            "baseUrl": provider.url,
            "apiKey": "fixture-key",
            "model": "fixture-model",
            "protocol": "responses",
        }
        try:
            with Albedo(
                provider,
                protocol="responses",
                providers={loose: profile, strict: dict(profile, imageEdge=1000)},
            ) as app:
                session = app.session()
                attach = {
                    "kind": "message",
                    "text": "look",
                    "images": [{"mime_type": "image/png", "data": WIDE_DATA}],
                }

                def switch(name):
                    with app.api(f"/sessions/{session}?view=configuration") as response:
                        json.load(response)
                        revision = response.headers["ETag"]
                    app.api(
                        f"/sessions/{session}?view=configuration",
                        {"provider_profile": name, "model": "fixture-model"},
                        method="PATCH",
                        headers={"If-Match": revision},
                    ).close()

                def attach_image():
                    return app.api(
                        f"/sessions/{session}/inputs/{operation_id()}",
                        attach,
                        method="PUT",
                    )

                def send(content):
                    start = len(provider.requests)
                    app.prompt(session, content).close()
                    app.idle(session)
                    return [r["request"] for r in provider.requests[start:]]

                attach_image().close()
                app.idle(session)
                self.assertEqual(widths(provider.requests[-1]["request"]), [1200])

                switch(strict)
                sent = len(provider.requests)
                with self.assertRaises(urllib.error.HTTPError) as refused:
                    attach_image().close()
                self.assertEqual(refused.exception.code, 409)
                self.assertIn(
                    "1000px edge limit", json.load(refused.exception)["detail"]
                )
                self.assertEqual(len(provider.requests), sent)

                strict_requests = send("show") + send("again")
                self.assertEqual(len(strict_requests), 3)
                for request in strict_requests:
                    self.assertEqual(widths(request), [1000])
                    self.assertEqual(json.dumps(request).count(NOTE), 1)
                output = next(
                    item
                    for item in strict_requests[1]["input"]
                    if item.get("type") == "function_call_output"
                )
                # show_image refuses it at the call, so the cell fails there.
                result = cell_result(output)
                self.assertEqual(result["status"], "error")
                self.assertIn(
                    "ValueError: 1200x10 image is over this model's 1000px edge limit",
                    result["output"],
                )
                self.assertNotIn("image_errors", result)
                entries = app.history(session)["items"]
                notes = [
                    part["text"]
                    for entry in entries
                    if entry["kind"] == "image_fit"
                    for part in entry["content"]
                    if part["kind"] == "text"
                ]
                self.assertEqual(len(notes), 1, notes)
                self.assertIn(NOTE, notes[0])
                original = [
                    entry
                    for entry in entries
                    if entry["kind"] == "user"
                    and any(
                        part["kind"] == "text" and part["text"] == "look"
                        for part in entry["content"]
                    )
                ]
                self.assertEqual(len(original), 1)
                [image] = [
                    part["image"]
                    for part in original[0]["content"]
                    if part["kind"] == "image"
                ]
                self.assertEqual(image["width"], 1200)
                self.assertEqual(image["height"], 10)
                self.assertEqual(image["original_bytes"], len(WIDE))

                source = hashlib.sha256(WIDE_DATA.encode()).hexdigest()
                database = f"file:{app.home / 'albedo.sqlite'}?mode=ro"

                def rows(session, needle):
                    with sqlite3.connect(database, uri=True) as db:
                        return [
                            seq
                            for (seq,) in db.execute(
                                "SELECT seq FROM transcript WHERE session=? "
                                "AND instr(payload, CAST(? AS BLOB))>0 ORDER BY seq",
                                (session, needle),
                            )
                        ]

                # The original row keeps its image; the fit is a row of its own.
                self.assertEqual(len(rows(session, source)), 2)
                self.assertEqual(len(rows(session, "image_fit")), 1)

                # A fit is permanent: a looser provider keeps the copy.
                switch(loose)
                [back] = send("back")
                self.assertEqual(widths(back), [1000])
                self.assertEqual(json.dumps(back).count(NOTE), 1)

                # A fork from before the fit starts from the original.
                checkpoint = next(
                    entry["checkpoint_id"]
                    for entry in entries
                    if entry["kind"] == "assistant"
                )
                branch = operation_id()
                app.api(
                    f"/sessions/{branch}",
                    {
                        "kind": "fork",
                        "source_session_id": session,
                        "checkpoint_id": checkpoint,
                    },
                    method="PUT",
                    headers={"If-None-Match": "*"},
                ).close()
                start = len(provider.requests)
                app.prompt(branch, "branch").close()
                app.idle(branch)
                [forked] = [r["request"] for r in provider.requests[start:]]
                self.assertEqual(widths(forked), [1200])
                self.assertNotIn(NOTE, json.dumps(forked))
        finally:
            provider.close()

    def test_an_image_cut_short_is_refused_where_it_is_attached(self):
        # A PNG cut off by a capped read kept a readable header, so it was
        # attached, and every later request carried an image the provider
        # could not decode: the session failed until it was forked.
        cut = png(2, 3, b"x" * 1000)[:500]
        cut_data = base64.b64encode(cut).decode()

        def reply(request):
            if request["input"][-1].get("type") == "function_call_output":
                return text("done")
            return python(
                "import base64\nshow_image(base64.b64decode('" + cut_data + "'))"
            )

        provider = Provider(reply)
        self.addCleanup(provider.close)
        with Albedo(provider, protocol="responses") as app:
            session = app.session()
            with self.assertRaises(urllib.error.HTTPError) as refused:
                app.api(
                    f"/sessions/{session}/inputs/{operation_id()}",
                    {
                        "kind": "message",
                        "text": "look",
                        "images": [{"mime_type": "image/png", "data": cut_data}],
                    },
                    method="PUT",
                ).close()
            self.assertEqual(refused.exception.code, 400)

            app.prompt(session, "show").close()
            app.idle(session)
            request = provider.requests[-1]["request"]
            output = next(
                item
                for item in request["input"]
                if item.get("type") == "function_call_output"
            )
            result = cell_result(output)
            self.assertEqual(result["status"], "error")
            self.assertIn(
                "ValueError: this 500-byte image/png stops before its IEND chunk",
                result["output"],
            )
            self.assertNotIn("data:image/png", json.dumps(request))

    def test_a_remote_image_is_shown_under_the_same_limits(self):
        cell = (
            LOOPBACK_REMOTE
            + """try:
    print(await rem.show_image("small.png"))
    await rem.show_image("wide.png")
finally:
    await rem.close()
"""
        )

        def reply(request):
            if request["input"][-1].get("type") == "function_call_output":
                return text("done")
            return python(cell)

        provider = Provider(reply)
        profile = {
            "baseUrl": provider.url,
            "apiKey": "fixture-key",
            "model": "fixture-model",
            "protocol": "responses",
            "imageEdge": 1000,
        }
        try:
            with Albedo(
                provider,
                protocol="responses",
                providers={f"strict-{provider.route}": profile},
            ) as app:
                (app.workspace / "small.png").write_bytes(png(10, 10))
                (app.workspace / "wide.png").write_bytes(WIDE)
                session = app.session()
                app.prompt(session, "remote").close()
                app.idle(session)
                [request] = [
                    r["request"]
                    for r in provider.requests
                    if r["request"]["input"][-1].get("type") == "function_call_output"
                ]
                output = next(
                    item
                    for item in request["input"]
                    if item.get("type") == "function_call_output"
                )
                result = cell_result(output)
                # The small image came over the connection and went on to the
                # model; the wide one failed the cell at its call.
                self.assertIn("attached image/png", result["output"])
                self.assertEqual(result["status"], "error")
                self.assertIn(
                    "ValueError: 1200x10 image is over this model's 1000px edge limit",
                    result["output"],
                )
                self.assertEqual(widths(request), [10])
        finally:
            provider.close()


TEXT_ONLY = {
    "fixture-cloud": {
        "id": "fixture-cloud",
        "name": "Fixture Cloud",
        "env": ["FIXTURE_API_KEY"],
        "api": "http://127.0.0.1/v1",
        "models": {
            "fixture": {
                "id": "fixture",
                "limit": {"context": 200000, "output": 8000},
                "modalities": {"input": ["text"]},
            }
        },
    }
}


@exclusive
class TextOnlyModelTest(unittest.TestCase):
    def test_a_text_only_model_is_told_and_sent_notes_not_images(self):
        def reply(request):
            if request["input"][-1].get("type") == "function_call_output":
                return text("done")
            if "show" in json.dumps(request["input"][-1]):
                return python(
                    "import base64\nshow_image(base64.b64decode('" + WIDE_DATA + "'))"
                )
            return text("ok")

        provider = Provider(reply, catalog=TEXT_ONLY)
        self.addCleanup(provider.close)

        def prepare(app):
            (app.home / "models.json").unlink(missing_ok=True)
            app.write_extensions(
                {"models": {"url": provider.url + "/models.json", "refreshHours": 24}}
            )

        with Albedo(
            provider,
            protocol="responses",
            prepare=prepare,
            providers={
                "fixture": {
                    "baseUrl": provider.url + "/v1",
                    "apiKey": "fixture-key",
                    "model": "fixture",
                    "protocol": "responses",
                }
            },
        ) as app:
            session = app.session()
            app.api(
                f"/sessions/{session}/inputs/{operation_id()}",
                {
                    "kind": "message",
                    "text": "look",
                    "images": [{"mime_type": "image/png", "data": WIDE_DATA}],
                },
                method="PUT",
            ).close()
            app.idle(session)
            app.prompt(session, "show").close()
            app.idle(session)
            for entry in provider.requests:
                sent = json.dumps(entry["request"])
                self.assertEqual(widths(entry["request"]), [], sent[:2000])
                self.assertIn(
                    "This model reads text only", entry["request"]["instructions"]
                )
            last = provider.requests[-1]["request"]
            # The user's image is a note; show_image refused at the call, so
            # the cell carries no image to omit.
            self.assertEqual(
                json.dumps(last).count("image(s) omitted: this model reads text only"),
                1,
            )
            output = next(
                item
                for item in last["input"]
                if item.get("type") == "function_call_output"
            )
            result = cell_result(output)
            self.assertEqual(result["status"], "error")
            self.assertIn("ValueError: this model reads text only", result["output"])