from dataclasses import dataclass
from datetime import datetime
from json import loads
from unittest import TestCase

from avalan.entities import (
    ReasoningToken,
    Token,
    TokenDetail,
    ToolCall,
    ToolCallDiagnostic,
    ToolCallDiagnosticCode,
    ToolCallDiagnosticStage,
    ToolCallError,
    ToolCallResult,
    ToolCallToken,
)
from avalan.event import Event, EventType
from avalan.server.routers.responses import (
    ResponseState,
    _new_state,
    _switch_state,
    _token_to_sse,
    _tool_call_event_item,
)
from avalan.server.sse import sse_message
from avalan.utils import to_json


class ResponsesUtilsTestCase(TestCase):
    def test_token_to_sse_formats_tokens(self) -> None:
        rt = ReasoningToken(token="r")
        tc = ToolCallToken(token="t")
        tok = Token(token="a")
        detail = TokenDetail(token="b")

        rt_event = _token_to_sse(rt, 0)[0]
        self.assertIn("response.reasoning_text.delta", rt_event)
        self.assertEqual(loads(rt_event.split("data: ")[1])["delta"], "r")

        tc_event = _token_to_sse(tc, 1)[0]
        self.assertIn("response.custom_tool_call_input.delta", tc_event)
        self.assertEqual(loads(tc_event.split("data: ")[1])["delta"], "t")

        tok_event = _token_to_sse(tok, 2)[0]
        self.assertIn("response.output_text.delta", tok_event)
        self.assertEqual(loads(tok_event.split("data: ")[1])["delta"], "a")

        detail_event = _token_to_sse(detail, 3)[0]
        self.assertEqual(loads(detail_event.split("data: ")[1])["delta"], "b")

    def test_token_to_sse_handles_tool_result(self) -> None:
        call = ToolCall(id="c1", name="t", arguments={"p": 1})
        result = ToolCallResult(
            id="c1", call=call, name="t", arguments={"p": 1}, result={"v": 2}
        )
        event = Event(type=EventType.TOOL_RESULT, payload={"result": result})

        events = _token_to_sse(event, 0)
        self.assertEqual(len(events), 1)
        self.assertIn("response.function_call_arguments.delta", events[0])
        data = loads(events[0].split("data: ")[1])
        self.assertEqual(data["id"], "c1")
        self.assertEqual(data["result"], '{"v": 2}')

    def test_token_to_sse_handles_tool_result_error(self) -> None:
        call = ToolCall(id="c2", name="t", arguments={})
        error = ToolCallError(
            id="c2",
            name="t",
            arguments={},
            call=call,
            error=RuntimeError("boom"),
            message="boom",
        )
        event = Event(type=EventType.TOOL_RESULT, payload={"result": error})

        events = _token_to_sse(event, 1)
        self.assertEqual(len(events), 1)
        data = loads(events[0].split("data: ")[1])
        self.assertEqual(data["id"], "c2")
        self.assertEqual(
            data["error"],
            {"type": "RuntimeError", "message": "Tool call failed."},
        )
        delta = loads(data["delta"])
        self.assertEqual(
            delta["error"],
            {"type": "RuntimeError", "message": "Tool call failed."},
        )
        self.assertNotIn("boom", events[0])

    def test_token_to_sse_handles_tool_diagnostic(self) -> None:
        call = ToolCall(id="c-diagnostic", name="missing", arguments={})
        diagnostic = ToolCallDiagnostic(
            id="diag-1",
            call_id=call.id,
            requested_name="missing",
            code=ToolCallDiagnosticCode.UNKNOWN_TOOL,
            stage=ToolCallDiagnosticStage.RESOLVE,
            message="Unknown tool.",
        )
        event = Event(
            type=EventType.TOOL_DIAGNOSTIC,
            payload={"call": call, "diagnostic": diagnostic},
        )

        events = _token_to_sse(event, 3)

        self.assertEqual(len(events), 1)
        self.assertIn("response.tool_call_diagnostic.delta", events[0])
        data = loads(events[0].split("data: ")[1])
        self.assertEqual(data["id"], "c-diagnostic")
        self.assertEqual(data["diagnostic"]["code"], "tool.unknown")
        delta = loads(data["delta"])
        self.assertEqual(delta["diagnostic"]["call_id"], "c-diagnostic")

    def test_tool_call_event_item_handles_tool_result_diagnostic(
        self,
    ) -> None:
        call = ToolCall(id="c-result-diagnostic", name="missing", arguments={})
        diagnostic = ToolCallDiagnostic(
            id="diag-result",
            call_id=call.id,
            requested_name="missing",
            code=ToolCallDiagnosticCode.UNKNOWN_TOOL,
            stage=ToolCallDiagnosticStage.RESOLVE,
            message="Unknown tool.",
        )
        event = Event(
            type=EventType.TOOL_RESULT,
            payload={"call": call, "result": diagnostic},
        )

        item = _tool_call_event_item(event)

        self.assertEqual(item["id"], "c-result-diagnostic")
        self.assertEqual(item["diagnostic"]["code"], "tool.unknown")

    def test_tool_call_event_item_handles_diagnostic_timing_fields(
        self,
    ) -> None:
        diagnostic = ToolCallDiagnostic(
            id="diag-timing",
            requested_name="missing",
            code=ToolCallDiagnosticCode.UNKNOWN_TOOL,
            stage=ToolCallDiagnosticStage.RESOLVE,
            message="Unknown tool.",
            started_at=datetime(2026, 1, 1, 12, 0, 0),
            finished_at=datetime(2026, 1, 1, 12, 0, 1),
            duration_ms=1000.0,
        )
        event = Event(
            type=EventType.TOOL_DIAGNOSTIC,
            payload={"diagnostics": [diagnostic]},
        )

        item = _tool_call_event_item(event)

        self.assertEqual(item["id"], "diag-timing")
        self.assertEqual(item["name"], "missing")
        self.assertEqual(
            item["diagnostic"]["started_at"], "2026-01-01T12:00:00"
        )
        self.assertEqual(
            item["diagnostic"]["finished_at"], "2026-01-01T12:00:01"
        )
        self.assertEqual(item["diagnostic"]["duration_ms"], 1000.0)

    def test_token_to_sse_ignores_malformed_tool_diagnostic(self) -> None:
        event = Event(
            type=EventType.TOOL_DIAGNOSTIC,
            payload={"diagnostics": ["bad"]},
        )

        self.assertEqual(_token_to_sse(event, 4), [])

    def test_token_to_sse_handles_tool_result_without_payload(self) -> None:
        call = ToolCall(id="c3", name="tool", arguments={})
        result = ToolCallResult(
            id="c3",
            name="tool",
            arguments={},
            call=call,
            result=None,
        )
        event = Event(type=EventType.TOOL_RESULT, payload={"result": result})

        events = _token_to_sse(event, 2)
        self.assertEqual(len(events), 1)
        data = loads(events[0].split("data: ")[1])
        self.assertIsNone(data["result"])

    def test_token_to_sse_handles_current_none_tool_result(self) -> None:
        call = ToolCall(id="c-none", name="tool", arguments={"x": 1})
        event = Event(
            type=EventType.TOOL_RESULT,
            payload={0: call, "result": None},
        )

        events = _token_to_sse(event, 5)

        self.assertEqual(len(events), 1)
        data = loads(events[0].split("data: ")[1])
        self.assertEqual(data["id"], "c-none")
        self.assertEqual(data["name"], "tool")
        self.assertEqual(data["arguments"], {"x": 1})
        self.assertIsNone(data["result"])

    def test_tool_call_event_item_ignores_missing_tool_call(self) -> None:
        event = Event(
            type=EventType.TOOL_PROCESS,
            payload={"call": None},
        )

        self.assertIsNone(_tool_call_event_item(event))

    def test_tool_call_event_item_ignores_unexpected_payload(
        self,
    ) -> None:
        event = Event(type=EventType.TOOL_PROCESS, payload="bad")

        self.assertIsNone(_tool_call_event_item(event))

    def test_token_to_sse_handles_tool_call_token_with_call(self) -> None:
        call = ToolCall(id="c4", name="adder", arguments={"x": 1})
        token = ToolCallToken(token="ignored", call=call)

        events = _token_to_sse(token, 3)
        self.assertEqual(len(events), 1)
        data = loads(events[0].split("data: ")[1])
        self.assertEqual(data["id"], "c4")
        delta = loads(data["delta"])
        self.assertEqual(delta["arguments"], {"x": 1})

    def test_switch_state_generates_events(self) -> None:
        state = _new_state(ReasoningToken(token="r"))
        events = _switch_state(None, state, None, None)
        self.assertEqual(state, ResponseState.REASONING)
        names = [e.split("\n")[0].split(": ")[1] for e in events]
        self.assertEqual(
            names,
            ["response.output_item.added", "response.content_part.added"],
        )

        new_state = _new_state(ToolCallToken(token="t"))
        events = _switch_state(state, new_state, None, None)
        self.assertEqual(new_state, ResponseState.TOOL_CALLING)
        names = [e.split("\n")[0].split(": ")[1] for e in events]
        self.assertEqual(
            names,
            [
                "response.reasoning_text.done",
                "response.content_part.done",
                "response.output_item.done",
                "response.output_item.added",
                "response.content_part.added",
            ],
        )

        state = new_state
        new_state = _new_state("answer")
        events = _switch_state(state, new_state, None, None)
        self.assertEqual(new_state, ResponseState.ANSWERING)
        names = [e.split("\n")[0].split(": ")[1] for e in events]
        self.assertEqual(
            names,
            [
                "response.custom_tool_call_input.done",
                "response.content_part.done",
                "response.output_item.done",
                "response.output_item.added",
                "response.content_part.added",
            ],
        )

        state = new_state
        new_state = _new_state(None)
        events = _switch_state(state, new_state, None, None)
        self.assertIsNone(new_state)
        names = [e.split("\n")[0].split(": ")[1] for e in events]
        self.assertEqual(
            names,
            [
                "response.output_text.done",
                "response.content_part.done",
                "response.output_item.done",
            ],
        )

    def test_switch_state_handles_new_tool_call_id(self) -> None:
        events = _switch_state(
            ResponseState.TOOL_CALLING,
            ResponseState.TOOL_CALLING,
            "old",
            "new",
        )

        names = [e.split("\n")[0].split(": ")[1] for e in events]
        self.assertEqual(
            names,
            [
                "response.custom_tool_call_input.done",
                "response.content_part.done",
                "response.output_item.done",
                "response.output_item.added",
                "response.content_part.added",
            ],
        )

        data_items = [loads(e.split("data: ")[1]) for e in events]
        self.assertEqual(data_items[0]["id"], "old")
        self.assertEqual(data_items[1]["part"]["id"], "old")
        self.assertEqual(data_items[2]["item"]["id"], "old")
        self.assertEqual(data_items[3]["item"]["id"], "new")
        self.assertEqual(data_items[4]["part"]["id"], "new")

    def test_tool_call_event_item_handles_custom_result(self) -> None:
        call = ToolCall(id="c5", name="calc", arguments={"a": 2})

        @dataclass(frozen=True, slots=True)
        class DummyResult:
            call: ToolCall
            payload: dict[str, str]

        result = DummyResult(call=call, payload={"status": "ok"})
        event = Event(type=EventType.TOOL_RESULT, payload={"result": result})

        item = _tool_call_event_item(event)
        self.assertEqual(item["id"], "c5")
        self.assertIs(item["result"], result)

        events = _token_to_sse(event, 4)
        data = loads(events[0].split("data: ")[1])
        delta = loads(data["delta"])
        result_payload = loads(delta["result"])
        self.assertEqual(result_payload["payload"], {"status": "ok"})

    def test_tool_call_event_item_supports_indexed_mapping_payload(
        self,
    ) -> None:
        call = ToolCall(id="c6", name="calc", arguments={"n": 1})

        class IndexedPayload(dict[int, ToolCall]):
            pass

        event = Event(
            type=EventType.TOOL_PROCESS,
            payload=IndexedPayload({0: call}),
        )

        item = _tool_call_event_item(event)
        self.assertEqual(item["id"], "c6")
        self.assertEqual(item["name"], "calc")

    def test_sse_formats_event_and_data(self) -> None:
        result = sse_message(to_json({"a": 1}), event="test.event")
        self.assertEqual(result, 'event: test.event\ndata: {"a": 1}\n\n')

    def test_sse_message_handles_empty_payload(self) -> None:
        result = sse_message("", event="empty")
        self.assertEqual(result, "event: empty\ndata: \n\n")
