from __future__ import annotations

import asyncio

from loushang.ai.types import AssistantMessage, TextPart, Usage
from loushang.coding.session.agent_event_router import AgentEventRouter


def _usage() -> Usage:
    return Usage(
        input=0,
        output=0,
        cache_read=0,
        cache_write=0,
        total_tokens=0,
        cost={},
    )


def _assistant_message(*, stop_reason: str = "stop", error_message: str | None = None) -> AssistantMessage:
    return AssistantMessage(
        role="assistant",
        content=[TextPart(type="text", text="reply")],
        api="anthropic-messages",
        provider="faux",
        model="faux-model",
        response_id=None,
        usage=_usage(),
        stop_reason=stop_reason,
        error_message=error_message,
        timestamp=0.0,
    )


class _RetryRecorder:
    def __init__(self, order: list[str], *, retryable: bool = False, did_retry: bool = False) -> None:
        self.order = order
        self.retryable = retryable
        self.did_retry = did_retry

    async def finish_success_if_needed(self, assistant_message: AssistantMessage) -> None:
        del assistant_message
        self.order.append("retry_finish_success")

    def should_prepare_retry(self, assistant_message: AssistantMessage) -> bool:
        del assistant_message
        self.order.append("retry_should_prepare")
        return self.retryable

    def ensure_future(self) -> None:
        self.order.append("retry_ensure_future")

    def is_retryable_error(self, assistant_message: AssistantMessage) -> bool:
        del assistant_message
        self.order.append("retry_is_retryable")
        return self.retryable

    async def handle_retryable_error(self, assistant_message: AssistantMessage) -> bool:
        del assistant_message
        self.order.append("retry_handle")
        return self.did_retry


class _CompactionRecorder:
    def __init__(self, order: list[str]) -> None:
        self.order = order

    def clear_overflow_recovery_attempted(self) -> None:
        self.order.append("compaction_clear_overflow")


def test_agent_event_router_preserves_assistant_message_end_ordering() -> None:
    order: list[str] = []
    assistant = _assistant_message()

    async def scenario() -> None:
        router = AgentEventRouter(
            append_message=lambda message: order.append(f"append:{message.role}"),
            dispatch_event=lambda event: _record_async(order, f"dispatch:{event['type']}"),
            emit_extension_agent_event=lambda event: _record_async(order, f"extension:{event['type']}"),
            record_tool_execution_error=lambda event: order.append(f"tool_error:{event['type']}"),
            retry_controller=_RetryRecorder(order),
            compaction_controller=_CompactionRecorder(order),
            sync_extension_diagnostics=lambda **kwargs: order.append(f"sync:{kwargs['phase']}"),
            record_assistant_response_error=lambda message: order.append(f"assistant_error:{message.role}"),
            check_auto_compaction=lambda message: _record_async(order, f"auto_compact:{message.role}"),
        )

        await router.handle({"type": "message_end", "message": assistant}, object())

    asyncio.run(scenario())

    assert order == [
        "append:assistant",
        "dispatch:message_end",
        "extension:message_end",
        "retry_finish_success",
        "compaction_clear_overflow",
    ]


def test_agent_event_router_records_tool_errors_before_dispatch_and_extension_mirror() -> None:
    order: list[str] = []

    async def scenario() -> None:
        router = AgentEventRouter(
            append_message=lambda message: order.append(f"append:{message.role}"),
            dispatch_event=lambda event: _record_async(order, f"dispatch:{event['type']}"),
            emit_extension_agent_event=lambda event: _record_async(order, f"extension:{event['type']}"),
            record_tool_execution_error=lambda event: order.append(f"tool_error:{event['type']}"),
            retry_controller=_RetryRecorder(order),
            compaction_controller=_CompactionRecorder(order),
            sync_extension_diagnostics=lambda **kwargs: order.append(f"sync:{kwargs['phase']}"),
            record_assistant_response_error=lambda message: order.append(f"assistant_error:{message.role}"),
            check_auto_compaction=lambda message: _record_async(order, f"auto_compact:{message.role}"),
        )

        await router.handle({"type": "tool_execution_end", "is_error": True}, object())

    asyncio.run(scenario())

    assert order == [
        "tool_error:tool_execution_end",
        "dispatch:tool_execution_end",
        "extension:tool_execution_end",
    ]


def test_agent_event_router_retries_before_auto_compaction_and_short_circuits() -> None:
    order: list[str] = []
    assistant = _assistant_message(stop_reason="error", error_message="503 service unavailable")

    async def scenario() -> None:
        router = AgentEventRouter(
            append_message=lambda message: order.append(f"append:{message.role}"),
            dispatch_event=lambda event: _record_async(order, f"dispatch:{event['type']}"),
            emit_extension_agent_event=lambda event: _record_async(order, f"extension:{event['type']}"),
            record_tool_execution_error=lambda event: order.append(f"tool_error:{event['type']}"),
            retry_controller=_RetryRecorder(order, retryable=True, did_retry=True),
            compaction_controller=_CompactionRecorder(order),
            sync_extension_diagnostics=lambda **kwargs: order.append(f"sync:{kwargs['phase']}"),
            record_assistant_response_error=lambda message: order.append(f"assistant_error:{message.role}"),
            check_auto_compaction=lambda message: _record_async(order, f"auto_compact:{message.role}"),
        )

        await router.handle({"type": "agent_end", "messages": [assistant]}, object())

    asyncio.run(scenario())

    assert order == [
        "dispatch:agent_end",
        "extension:agent_end",
        "sync:runtime",
        "assistant_error:assistant",
        "retry_should_prepare",
        "retry_ensure_future",
        "retry_is_retryable",
        "retry_handle",
    ]


async def _record_async(order: list[str], value: str) -> None:
    order.append(value)
