import json
import gzip
from types import SimpleNamespace

import httpx2
import pytest

from artemis.llm.responses_diagnostics import MAX_FRAME_BYTES, ResponsesObservation
from artemis.llm.diagnostics import ACTIVE_ATTEMPT, observe_response_headers
from artemis.llm.tool_constraints import MAX_REQUEST_BYTES, request_constraints, tool_constraints


def frame(kind, sequence, **fields):
    return (
        "data: "
        + json.dumps({"type": kind, "sequence_number": sequence, **fields}, ensure_ascii=False)
        + "\r\n\r\n"
    ).encode()


@pytest.mark.parametrize("size", [1, 37, 100000])
def test_raw_frames_split_unicode_and_done_snapshots_are_not_double_counted(size):
    arguments = '{"text":"私人🙂"}'
    data = b"".join(
        [
            frame("response.created", 0, response={"id": "resp_one"}),
            frame(
                "response.output_item.added",
                1,
                output_index=2,
                item={
                    "type": "function_call",
                    "id": "fc_one",
                    "call_id": "call_one",
                    "name": "click",
                    "arguments": "",
                },
            ),
            frame(
                "response.function_call_arguments.delta",
                2,
                output_index=2,
                item_id="fc_one",
                delta=arguments,
            ),
            frame(
                "response.function_call_arguments.done",
                3,
                output_index=2,
                item_id="fc_one",
                arguments=arguments,
            ),
            frame(
                "response.output_item.done",
                4,
                output_index=2,
                item={"type": "function_call", "id": "fc_one", "arguments": arguments},
            ),
            frame(
                "response.completed",
                5,
                response={
                    "id": "resp_one",
                    "tools": [{"type": "function", "name": "click", "strict": True}],
                },
            ),
        ]
    )
    observer = ResponsesObservation()
    for offset in range(0, len(data), size):
        observer.feed(data[offset : offset + size])
    observer.end(eof=True)
    stats = observer.snapshot()
    assert stats["body_bytes"] == len(data)
    assert stats["parse_errors"] == 0
    assert stats["incomplete_sse_frame"] is False
    (tool,) = stats["tools"]
    assert tool["argument_chars"] == len(arguments)
    assert (
        tool["delta_events"]
        == tool["arguments_done_events"]
        == tool["output_item_done_events"]
        == 1
    )
    assert tool["call_id"] == "call_one"
    (response,) = stats["responses"]
    assert response["sequence_first"] == 0 and response["sequence_last"] == 5
    assert response["sequence_forward_gaps"] == response["sequence_non_increasing"] == 0
    assert response["returned_tool_constraints"]["tools"][0]["strict"] == "true"
    assert "私人" not in json.dumps(stats, ensure_ascii=False)
    assert not observer.line and not observer.lines


def test_sequence_statistics_are_per_response_and_describe_observed_jumps():
    observer = ResponsesObservation()
    observer.event({"type": "response.in_progress", "sequence_number": 1})
    observer.event({"type": "response.created", "sequence_number": 0, "response": {"id": "resp_a"}})
    for seq in [2, 2, 1, 3, None]:
        observer.event({"type": "response.in_progress", "sequence_number": seq})
    observer.event({"type": "response.created", "sequence_number": 0, "response": {"id": "resp_b"}})
    a, b = observer.snapshot()["responses"]
    assert a["sequence_forward_gaps"] == a["sequence_forward_missing_total"] == 2
    assert a["sequence_non_increasing"] == 2
    assert a["sequence_equal_previous"] == 1
    assert a["sequence_missing_events"] == 1
    assert b["sequence_non_increasing"] == b["sequence_forward_gaps"] == 0
    assert observer.data["unknown_response_events"] == 1


def test_oversized_frame_recovery_and_partial_eof_keep_memory_bounded():
    observer = ResponsesObservation()
    observer.feed(b"data: " + b"x" * MAX_FRAME_BYTES)
    observer.feed(b"\r\n\r\n")
    observer.feed(b"data: invalid\n\n")
    observer.feed(frame("response.created", 0, response={"id": "resp_valid"}))
    observer.feed(b'data: {"private":')
    observer.end(eof=True)
    stats = observer.snapshot()
    assert stats["oversized_frames"] == 1
    assert stats["parse_errors"] == 1
    assert len(stats["responses"]) == 1
    assert stats["incomplete_sse_frame"] is True
    assert "private" not in repr(vars(observer))
    assert not observer.line and not observer.lines
    observer.end()  # HTTP close can follow EOF; retain the original coverage gap.
    assert observer.data["incomplete_sse_frame"] is True


def test_raw_tool_and_response_statistics_have_fixed_caps():
    observer = ResponsesObservation()
    observer.event({"type": "response.created", "response": {"id": "resp_one"}})
    for index in range(100):
        observer.event(
            {
                "type": "response.function_call_arguments.delta",
                "item_id": f"fc_{index}",
                "output_index": index,
                "delta": " ",
            }
        )
    for index in range(20):
        observer.event({"type": "response.created", "response": {"id": f"resp_{index}"}})
    assert len(observer.tools) == len(observer.fingerprints) == len(observer.aliases) == 64
    assert len(observer.responses) == 16
    assert observer.data["statistics_truncated"] is True


@pytest.mark.parametrize("size", [1, 37, MAX_FRAME_BYTES * 2])
def test_oversized_crlf_frame_recovers_across_arbitrary_read_boundaries(size):
    observer = ResponsesObservation()
    data = b"data: " + b"x" * MAX_FRAME_BYTES + b"\r\n\r\n"
    data += frame("response.created", 0, response={"id": "resp_after_large"})
    for offset in range(0, len(data), size):
        observer.feed(data[offset : offset + size])
    assert observer.data["oversized_frames"] == 1
    assert observer.data["parse_errors"] == 0
    assert list(observer.responses) == ["resp_after_large"]
    assert not observer.line and not observer.lines and not observer.dropping


def test_alias_statistics_are_bounded_even_when_item_keeps_changing_index():
    observer = ResponsesObservation()
    observer.event({"response": {"id": "resp_one"}})
    for index in range(1000):
        observer.event(
            {
                "type": "response.function_call_arguments.delta",
                "item_id": "fc_same",
                "output_index": index,
                "delta": " ",
            }
        )
    assert len(observer.tools) == 1
    assert len(observer.aliases) == 128
    assert observer.data["statistics_truncated"] is True


def test_malformed_event_type_cannot_break_passive_observation():
    observer = ResponsesObservation()
    observer.feed(frame([], 0, response={"id": "resp_one"}))
    assert observer.snapshot()["responses"][0]["events"] == 1


@pytest.mark.asyncio
async def test_compressed_body_is_explicitly_skipped_and_sdk_decoding_is_unchanged():
    body = frame("response.created", 0, response={"id": "resp_one"})
    compressed = gzip.compress(body)

    class Stream(httpx2.AsyncByteStream):
        async def __aiter__(self):
            yield compressed

    attempt = SimpleNamespace(
        wire_streams=[],
        data={},
        headers=lambda _: None,
        stream_work=lambda *_: None,
        publish_if_due=lambda: None,
    )
    response = httpx2.Response(
        200,
        headers={"content-type": "text/event-stream", "content-encoding": "gzip"},
        stream=Stream(),
    )
    token = ACTIVE_ATTEMPT.set(attempt)
    try:
        await observe_response_headers(response)
        assert b"".join([chunk async for chunk in response.aiter_bytes()]) == body
        stats = attempt.wire_streams[0].observation.snapshot()
        assert stats["sse_observation"] == "skipped_content_encoding"
        assert stats["body_bytes"] == len(compressed)
        assert not stats["responses"] and not stats["parse_errors"]
    finally:
        await response.aclose()
        ACTIVE_ATTEMPT.reset(token)


@pytest.mark.parametrize("strict,expected", [(None, "omitted"), (False, "false"), (True, "true")])
def test_actual_strict_state_and_optional_schema_are_observed_without_mutation(strict, expected):
    tool = {
        "type": "function",
        "name": "click",
        "parameters": {
            "type": "object",
            "properties": {"target": {"anyOf": [{"type": "string"}, {"type": "null"}]}},
            "required": [],
        },
    }
    if strict is not None:
        tool["strict"] = strict
    body = json.dumps(
        {"tools": [tool], "input": "private prompt", "max_output_tokens": 2048}
    ).encode()
    result = request_constraints(body)
    assert result["max_output_tokens"] == 2048
    (observed,) = result["tools"]
    assert observed["strict"] == expected
    assert observed["schema"]["optional_properties"] == observed["schema"]["open_objects"] == 1
    assert "strict" not in tool if strict is None else tool["strict"] is strict
    assert "private" not in json.dumps(result)
    assert json.loads(body)["tools"][0] == tool


def test_constraints_bound_large_requests_namespaces_and_deep_schemas():
    assert request_constraints(b"x" * (MAX_REQUEST_BYTES + 1)) == {"observation": "body_too_large"}
    assert request_constraints(b"no JSON") == {"observation": "invalid_json"}
    tools = [
        {
            "type": "namespace",
            "name": "functions",
            "tools": [{"type": "function", "name": "click"}],
        },
        {"type": "function", "name": "drag", "strict": False},
    ]
    result = tool_constraints(tools)
    assert {t["name"]: t["strict"] for t in result["tools"]} == {
        "click": "omitted",
        "drag": "false",
    }
    schema = {"type": "object", "properties": {str(i): {"type": "string"} for i in range(2000)}}
    result = tool_constraints([{"type": "function", "name": "wide", "parameters": schema}])
    assert result["tools"][0]["schema"]["truncated"] is True
    assert result["tools"][0]["schema"]["nodes"] <= 1024