import itertools
import json

from langchain_core.messages import AIMessageChunk
import langchain_core.messages.ai as ai
import pytest

from artemis.llm.chunk_accumulator import ChunkAccumulator
from artemis.llm.response_guard import ResponseIntegrityError, ensure_response_complete


def test_batching_preserves_interleaved_tools_unicode_metadata_and_usage():
    arguments = [json.dumps({"target": "中文🙂" * 80}), json.dumps({"target": 3})]
    chunks = [AIMessageChunk(content="Live text", id="resp_one")]
    for index in range(2):
        chunks.append(
            AIMessageChunk(
                content=[],
                tool_call_chunks=[
                    {"name": "click", "args": "", "id": f"call_{index}", "index": index}
                ],
            )
        )
    for pair in itertools.zip_longest(*arguments):
        for index, value in enumerate(pair):
            if value is not None:
                chunks.append(
                    AIMessageChunk(content=[], tool_call_chunks=[{"args": value, "index": index}])
                )
    chunks.append(
        AIMessageChunk(
            content=[],
            response_metadata={"id": "resp_one", "status": "completed"},
            usage_metadata={"input_tokens": 10, "output_tokens": 5, "total_tokens": 15},
            chunk_position="last",
        )
    )
    expected = chunks[0]
    accumulator = ChunkAccumulator()
    max_pending = 0
    for index, chunk in enumerate(chunks):
        if index:
            expected = expected + chunk
        accumulator.append(chunk)
        if accumulator.ready:
            accumulator.flush()
        max_pending = max(max_pending, len(accumulator.pending))
    result = accumulator.flush()
    assert result.model_dump() == expected.model_dump()
    assert [t["args"] for t in result.tool_calls] == [json.loads(a) for a in arguments]
    assert result.usage_metadata["total_tokens"] == 15
    assert max_pending < 256
    assert not accumulator.pending


def test_partial_json_is_not_reparsed_for_every_single_character(monkeypatch):
    chunks = [
        AIMessageChunk(
            content=[], tool_call_chunks=[{"name": "click", "args": '{"target":x', "index": 0}]
        )
    ]
    chunks.extend(
        AIMessageChunk(content=[], tool_call_chunks=[{"args": "x", "index": 0}]) for _ in range(800)
    )
    original = ai.parse_partial_json
    parsed_lengths = []

    def counted(arguments, *args, **kwargs):
        parsed_lengths.append(len(arguments))
        return original(arguments, *args, **kwargs)

    monkeypatch.setattr(ai, "parse_partial_json", counted)
    accumulator = ChunkAccumulator()
    for chunk in chunks:
        accumulator.append(chunk)
        if accumulator.ready:
            accumulator.flush()
    result = accumulator.flush()
    assert result.tool_call_chunks[0]["args"] == '{"target":x' + "x" * 800
    assert len(parsed_lengths) < 20
    result.response_metadata["status"] = "completed"
    with pytest.raises(ResponseIntegrityError):
        ensure_response_complete(result)