import hashlib
import json

import pytest
from langchain_core.messages import AIMessageChunk

from artemis.llm.argument_fingerprint import ArgumentFingerprint, MAX_PREFIXES, PREFIX_BYTES
from artemis.llm.diagnostics import RequestAttempt


@pytest.mark.parametrize("fragment_size", [1, 37, 4096, 100000])
def test_fingerprints_match_at_byte_offsets_despite_unicode_and_different_fragmentation(
    fragment_size,
):
    arguments = '{"target":"' + "私人内容🙂" * 8000 + '"}'
    expected = arguments.encode("utf-8")
    fingerprint = ArgumentFingerprint()
    for offset in range(0, len(arguments), fragment_size):
        fingerprint.observe(arguments[offset : offset + fragment_size])
    result = fingerprint.snapshot()
    assert result["argument_chars"] == len(arguments)
    assert result["argument_bytes"] == len(expected)
    assert result["argument_sha256"] == hashlib.sha256(expected).hexdigest()
    assert result["argument_encoding_valid"] is True
    assert result["argument_prefixes"] == [
        {"bytes": size, "sha256": hashlib.sha256(expected[:size]).hexdigest()}
        for size in range(PREFIX_BYTES, PREFIX_BYTES * MAX_PREFIXES + 1, PREFIX_BYTES)
    ]
    assert "私人" not in json.dumps(result, ensure_ascii=False)


def test_invalid_utf8_is_marked_without_failing_or_forging_a_digest():
    fingerprint = ArgumentFingerprint()
    for value in ["safe", "\ud800", "later"]:
        fingerprint.observe(value)
    expected = {
        "argument_chars": 10,
        "argument_bytes": None,
        "argument_sha256": None,
        "argument_deltas": 3,
        "argument_prefixes": [],
        "argument_encoding_valid": False,
    }
    assert {key: fingerprint.snapshot()[key] for key in expected} == expected


def test_shape_counts_cross_delta_whitespace_and_repetition_without_retaining_text():
    fingerprint = ArgumentFingerprint()
    for fragment in ['{"text":"private', " ", " \t", "", " \t", " \t", '\u00a0"}']:
        fingerprint.observe(fragment)
    stats = fingerprint.snapshot()
    assert stats["argument_json_whitespace_chars"] == 7
    assert stats["argument_json_whitespace_run_max"] == 7
    assert stats["argument_json_whitespace_run"] == 0
    assert stats["argument_delta_chars_min"] == 1
    assert stats["argument_delta_chars_max"] == len('{"text":"private')
    assert stats["argument_same_delta_run_max"] == 3
    assert stats["argument_same_delta_run"] == 1
    assert stats["argument_deltas"] == 6
    assert "private" not in repr(vars(fingerprint))
    assert "private" not in json.dumps(stats)


def test_tool_tracking_is_bounded_but_total_count_keeps_all_received_arguments():
    attempt = RequestAttempt(
        None, request_id="bounded", attempt=1, timeout=300, max_retries=0, metadata={}
    )
    for index in range(65):
        attempt.chunk(
            AIMessageChunk(
                content=[],
                tool_call_chunks=[{"name": "click", "args": "{}", "index": index}],
            )
        )
    assert len(attempt.data["tools"]) == len(attempt.tool_fingerprints) == 64
    assert attempt.data["tool_stats_truncated"] is True
    assert attempt.data["tool_argument_chars"] == 130