"""Vision / documents / citations smoke tests on Haiku 4.5: 1x1 PNG, invalid image bytes, text document with
citations (char_location), search_result citation, generated 1-page PDF (page_location), count_tokens for an image.
≈$0.004 per run.
"""
from __future__ import annotations

import base64
import struct
import zlib

import pytest


def tiny_png_b64() -> str:
    def chunk(t: bytes, d: bytes) -> bytes:
        return struct.pack(">I", len(d)) + t + d + struct.pack(">I", zlib.crc32(t + d) & 0xFFFFFFFF)
    png = (b"\x89PNG\r\n\x1a\n" + chunk(b"IHDR", struct.pack(">IIBBBBB", 1, 1, 8, 2, 0, 0, 0))
           + chunk(b"IDAT", zlib.compress(b"\x00\xff\x00\x00")) + chunk(b"IEND", b""))
    return base64.b64encode(png).decode()


def tiny_pdf_b64(text: str) -> str:
    stream = f"BT /F1 18 Tf 40 750 Td ({text}) Tj ET".encode()
    objs = [b"<< /Type /Catalog /Pages 2 0 R >>", b"<< /Type /Pages /Kids [3 0 R] /Count 1 >>",
            b"<< /Type /Page /Parent 2 0 R /MediaBox [0 0 612 792] /Contents 4 0 R /Resources << /Font << /F1 5 0 R >> >> >>",
            b"<< /Length %d >>\nstream\n" % len(stream) + stream + b"\nendstream", b"<< /Type /Font /Subtype /Type1 /BaseFont /Helvetica >>"]
    out, offs = b"%PDF-1.4\n", []
    for i, o in enumerate(objs, 1):
        offs.append(len(out)); out += f"{i} 0 obj\n".encode() + o + b"\nendobj\n"
    xref = len(out)
    out += f"xref\n0 {len(objs)+1}\n0000000000 65535 f \n".encode() + b"".join(f"{o:010d} 00000 n \n".encode() for o in offs)
    out += f"trailer\n<< /Size {len(objs)+1} /Root 1 0 R >>\nstartxref\n{xref}\n%%EOF\n".encode()
    return base64.b64encode(out).decode()


Q = {"type": "text", "text": "What color is grass? Answer in one sentence and cite."}


def _msg(model: str, content: list, max_tokens: int = 60) -> dict:
    return {"model": model, "max_tokens": max_tokens, "messages": [{"role": "user", "content": content}]}


def _first_citation(body: dict) -> dict:
    for b in body["content"]:
        if b["type"] == "text" and b.get("citations"):
            return b["citations"][0]
    raise AssertionError(f"no citation in {body['content']}")


def test_base64_image_and_token_count(anthropic, models):
    img = {"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": tiny_png_b64()}}
    st, body, _ = anthropic("POST", "/v1/messages", _msg(models["anthropic"], [img, {"type": "text", "text": "What color is this pixel? One word."}], 10),
                            est_cost_usd=0.00005, note="test_vision 1x1 png")
    assert st == 200 and "red" in body["content"][0]["text"].lower()
    st, cnt, _ = anthropic("POST", "/v1/messages/count_tokens", {"model": models["anthropic"], "messages": [{"role": "user", "content": [img, {"type": "text", "text": "Describe."}]}]},
                           note="test_vision count_tokens image")
    assert st == 200 and 5 < cnt["input_tokens"] < 100


def test_invalid_image_bytes_400(anthropic, models):
    bad = {"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": base64.b64encode(b"notapng").decode()}}
    st, body, _ = anthropic("POST", "/v1/messages", _msg(models["anthropic"], [bad, {"type": "text", "text": "Describe."}], 10), note="test_vision invalid image")
    assert st == 400 and "image" in body["error"]["message"].lower()


def test_text_document_char_location_citation(anthropic, models):
    doc = {"type": "document", "source": {"type": "text", "media_type": "text/plain", "data": "The sky is blue. Grass is green."}, "title": "Colors", "citations": {"enabled": True}}
    st, body, _ = anthropic("POST", "/v1/messages", _msg(models["anthropic"], [doc, Q]), est_cost_usd=0.0007, note="test_vision citations text")
    assert st == 200, body
    c = _first_citation(body)
    assert c["type"] == "char_location" and "green" in c["cited_text"].lower() and c["end_char_index"] > c["start_char_index"]


def test_search_result_citation(anthropic, models):
    sr = {"type": "search_result", "source": "kb://colors/1", "title": "Color facts",
          "content": [{"type": "text", "text": "The sky is blue."}, {"type": "text", "text": "Grass is green."}], "citations": {"enabled": True}}
    st, body, _ = anthropic("POST", "/v1/messages", _msg(models["anthropic"], [sr, Q]), est_cost_usd=0.0007, note="test_vision citations search_result")
    assert st == 200, body
    c = _first_citation(body)
    assert c["type"] == "search_result_location" and c["source"] == "kb://colors/1" and c["search_result_index"] == 0


def test_pdf_page_location_citation(anthropic, models):
    pdf = {"type": "document", "source": {"type": "base64", "media_type": "application/pdf", "data": tiny_pdf_b64("The sky is blue. Grass is green.")},
           "title": "Colors PDF", "citations": {"enabled": True}}
    st, body, _ = anthropic("POST", "/v1/messages", _msg(models["anthropic"], [pdf, Q]), est_cost_usd=0.0025, note="test_vision citations pdf")
    assert st == 200, body
    c = _first_citation(body)
    assert c["type"] == "page_location" and c["start_page_number"] == 1 and c["end_page_number"] == 2


def test_mixed_citation_settings_400(anthropic, models):
    on = {"type": "document", "source": {"type": "text", "media_type": "text/plain", "data": "A."}, "citations": {"enabled": True}}
    off = {"type": "document", "source": {"type": "text", "media_type": "text/plain", "data": "B."}, "citations": {"enabled": False}}
    st, body, _ = anthropic("POST", "/v1/messages", _msg(models["anthropic"], [on, off, Q], 10), note="test_vision mixed citations")
    assert st == 400 and "itations" in body["error"]["message"]