from __future__ import annotations

import asyncio
from typing import Any
from unittest.mock import MagicMock

import pytest

from app.ai_provider.types.chat import TokenUsage
from app.ai_provider.types.messages import HumanMessage
from app.application.agent import (
    AfterToolCallDecision,
    Agent,
    AgentConfig,
    AgentContext,
    AgentEvent,
    AgentRunContext,
    AgentTurnResult,
    BeforeToolCallDecision,
    ExtensionRegistry,
    PendingMessageQueue,
    QueueMode,
    ToolCallContext,
    TurnSnapshot,
)
from app.modules.tool_runtime.definition import ToolResult


def test_pending_message_queue_supports_both_modes() -> None:
    queue = PendingMessageQueue(QueueMode.one_at_a_time)
    queue.put(HumanMessage(content="一"))
    queue.put(HumanMessage(content="二"))

    assert [message.content for message in queue.drain()] == ["一"]
    assert [message.content for message in queue.drain()] == ["二"]

    queue.mode = QueueMode.all
    queue.put(HumanMessage(content="三"))
    queue.put(HumanMessage(content="四"))
    assert [message.content for message in queue.drain()] == ["三", "四"]


async def test_extension_registry_orders_context_and_tool_hooks() -> None:
    registry = ExtensionRegistry()
    run = AgentRunContext(run_id="run_1")

    async def first(context: AgentContext, _run: AgentRunContext) -> AgentContext:
        return AgentContext(
            messages=(*context.messages, HumanMessage(content="first")),
            tools=context.tools,
            model_hint=context.model_hint,
        )

    async def second(context: AgentContext, _run: AgentRunContext) -> AgentContext:
        return AgentContext(
            messages=(*context.messages, HumanMessage(content="second")),
            tools=context.tools,
            model_hint=context.model_hint,
        )

    async def rewrite(
        _call: ToolCallContext, _run: AgentRunContext
    ) -> BeforeToolCallDecision:
        return BeforeToolCallDecision(input={"value": 2})

    async def terminate(
        _call: ToolCallContext,
        result: ToolResult,
        _run: AgentRunContext,
    ) -> AfterToolCallDecision:
        return AfterToolCallDecision(result=result, terminate=True)

    registry.add_context_transformer(first)
    registry.add_context_transformer(second)
    registry.add_before_tool_hook(rewrite)
    registry.add_after_tool_hook(terminate)

    transformed = await registry.transform_context(AgentContext((), ()), run)
    assert [message.content for message in transformed.messages] == ["first", "second"]

    call = ToolCallContext("call_1", "tool", {"value": 1}, ())
    before = await registry.before_tool_call(call, run)
    assert before.input == {"value": 2}

    result = ToolResult(True, 0, "ok", "", 1, False, False, "trace_1")
    after = await registry.after_tool_call(call, result, run)
    assert after.result is result
    assert after.terminate is True


async def test_extension_registry_extend_preserves_registration_order() -> None:
    registry = ExtensionRegistry()
    skill_registry = ExtensionRegistry()

    async def base(context: AgentContext, _run: AgentRunContext) -> AgentContext:
        return AgentContext((*context.messages, HumanMessage(content="base")), context.tools)

    async def skill(context: AgentContext, _run: AgentRunContext) -> AgentContext:
        return AgentContext((*context.messages, HumanMessage(content="skill")), context.tools)

    registry.add_context_transformer(base)
    skill_registry.add_context_transformer(skill)
    registry.extend(skill_registry)

    transformed = await registry.transform_context(AgentContext((), ()), AgentRunContext())

    assert [message.content for message in transformed.messages] == ["base", "skill"]


async def test_stateful_agent_consumes_follow_up_and_publishes_events() -> None:
    calls: list[list[str]] = []
    events: list[str] = []

    async def runner(*, messages: list[Any], **_kwargs: object) -> AgentTurnResult:
        calls.append([message.content for message in messages])
        messages.append(HumanMessage(content=f"assistant-{len(calls)}"))
        return AgentTurnResult(
            content=f"done-{len(calls)}",
            reasoning="",
            iterations=1,
            tool_history=[{"name": f"tool-{len(calls)}"}],
            usage=TokenUsage(total_tokens=10),
        )

    extensions = ExtensionRegistry()

    async def listen(event: AgentEvent, _run: AgentRunContext) -> None:
        events.append(event.type)

    extensions.add_event_listener(listen)
    agent = Agent(
        AgentConfig(tools=(), tool_runtime=MagicMock(), max_iterations=2),
        extensions=extensions,
        runner=runner,
    )

    async def enqueue_follow_up(event: AgentEvent, _run: AgentRunContext) -> None:
        if event.type == "agent.started":
            assert agent.follow_up("再做一步") is True

    extensions.add_event_listener(enqueue_follow_up)
    result = await agent.prompt("开始")

    assert result.content == "done-2"
    assert calls == [["开始"], ["开始", "assistant-1", "再做一步"]]
    assert result.iterations == 2
    assert result.usage.total_tokens == 20
    assert [item["name"] for item in result.tool_history] == ["tool-1", "tool-2"]
    assert events == ["agent.started", "agent.ended", "agent.settled"]
    assert agent.state.status == "completed"
    assert agent.state.streaming_message == "done-2"


async def test_stateful_agent_does_not_continue_after_runner_truncates() -> None:
    calls = 0

    async def runner(*, messages: list[Any], **_kwargs: object) -> AgentTurnResult:
        nonlocal calls
        calls += 1
        messages.append(HumanMessage(content="assistant"))
        return AgentTurnResult(content="partial", reasoning="", iterations=1, truncated=True)

    extensions = ExtensionRegistry()
    agent = Agent(
        AgentConfig(tools=(), tool_runtime=None, max_iterations=3),
        extensions=extensions,
        runner=runner,
    )

    async def enqueue(event: AgentEvent, _run: AgentRunContext) -> None:
        if event.type == "agent.started":
            assert agent.follow_up("稍后处理") is True

    extensions.add_event_listener(enqueue)
    result = await agent.prompt("开始")

    assert calls == 1
    assert result.truncated is True
    assert [message.content for message in agent.follow_ups.drain()] == ["稍后处理"]


async def test_follow_up_is_preserved_when_iteration_budget_is_exhausted() -> None:
    async def runner(*, messages: list[Any], **_kwargs: object) -> AgentTurnResult:
        messages.append(HumanMessage(content="assistant"))
        return AgentTurnResult(content="done", reasoning="", iterations=1)

    extensions = ExtensionRegistry()
    agent = Agent(
        AgentConfig(tools=(), tool_runtime=None, max_iterations=1),
        extensions=extensions,
        runner=runner,
    )

    async def enqueue(event: AgentEvent, _run: AgentRunContext) -> None:
        if event.type == "agent.started":
            assert agent.follow_up("稍后处理") is True

    extensions.add_event_listener(enqueue)
    result = await agent.prompt("开始")

    assert result.truncated is True
    assert bool(agent.follow_ups) is True
    assert [message.content for message in agent.follow_ups.drain()] == ["稍后处理"]


async def test_started_listener_failure_restores_idle_state() -> None:
    extensions = ExtensionRegistry()

    async def fail(_event: AgentEvent, _run: AgentRunContext) -> None:
        raise RuntimeError("event store unavailable")

    extensions.add_event_listener(fail)
    agent = Agent(
        AgentConfig(tools=(), tool_runtime=MagicMock()),
        extensions=extensions,
    )

    with pytest.raises(RuntimeError, match="event store unavailable"):
        await agent.prompt("开始")
    await agent.wait_for_idle()
    assert agent.state.status == "failed"


async def test_stateful_agent_abort_cancels_active_run() -> None:
    started = asyncio.Event()
    events: list[str] = []

    async def runner(**_kwargs: object) -> AgentTurnResult:
        started.set()
        await asyncio.Event().wait()
        raise AssertionError("unreachable")

    extensions = ExtensionRegistry()

    async def listen(event: AgentEvent, _run: AgentRunContext) -> None:
        events.append(event.type)

    extensions.add_event_listener(listen)
    agent = Agent(
        AgentConfig(tools=(), tool_runtime=MagicMock()),
        extensions=extensions,
        runner=runner,
    )
    task = asyncio.create_task(agent.prompt("开始"))
    await started.wait()

    assert await agent.abort() is True
    assert agent.run_context.cancellation.cancelled is True
    with pytest.raises(asyncio.CancelledError):
        await task
    await agent.wait_for_idle()
    assert agent.state.status == "cancelled"
    assert events == ["agent.started", "agent.cancelled", "agent.settled"]


async def test_agent_replaces_cancelled_token_on_next_run() -> None:
    async def runner(**_kwargs: object) -> AgentTurnResult:
        return AgentTurnResult(content="done", reasoning="", iterations=1)

    agent = Agent(AgentConfig(tools=(), tool_runtime=None), runner=runner)
    agent.run_context.cancellation.cancel()

    await agent.prompt("重新开始")

    assert agent.run_context.cancellation.cancelled is False


async def test_extension_registry_should_stop_short_circuits_on_first_policy() -> None:
    registry = ExtensionRegistry()
    run = AgentRunContext(run_id="run_1")
    turn = TurnSnapshot(
        iteration=1,
        content="done",
        reasoning="",
        tool_calls=(),
        tool_results=(),
    )
    calls: list[str] = []

    async def keep_going(snapshot: TurnSnapshot, _run: AgentRunContext) -> bool:
        calls.append("keep_going")
        assert snapshot is turn
        return False

    async def stop_now(_snapshot: TurnSnapshot, _run: AgentRunContext) -> bool:
        calls.append("stop_now")
        return True

    async def never_called(_snapshot: TurnSnapshot, _run: AgentRunContext) -> bool:
        calls.append("never_called")
        return False

    registry.add_turn_policy(keep_going)
    registry.add_turn_policy(stop_now)
    registry.add_turn_policy(never_called)

    assert await registry.should_stop(turn, run) is True
    assert calls == ["keep_going", "stop_now"]


async def test_extension_registry_should_stop_false_when_all_policies_decline() -> None:
    registry = ExtensionRegistry()
    turn = TurnSnapshot(
        iteration=2,
        content="",
        reasoning="",
        tool_calls=({"name": "tool", "args": "{}"},),
        tool_results=(),
    )

    async def decline(_snapshot: TurnSnapshot, _run: AgentRunContext) -> bool:
        return False

    registry.add_turn_policy(decline)
    registry.add_turn_policy(decline)

    assert await registry.should_stop(turn, AgentRunContext()) is False


async def test_extension_registry_should_stop_respects_cancellation() -> None:
    registry = ExtensionRegistry()
    run = AgentRunContext()
    run.cancellation.cancel()
    turn = TurnSnapshot(
        iteration=1,
        content="",
        reasoning="",
        tool_calls=(),
        tool_results=(),
    )

    async def policy(_snapshot: TurnSnapshot, _run: AgentRunContext) -> bool:
        return False

    # raise_if_cancelled 在 policy 遍历循环内，故至少注册一个 policy 才会触发；
    # 这与 transform_context / before_tool_call / after_tool_call 一致：空注册时
    # 方法无 await 点，不额外检查取消信号。
    registry.add_turn_policy(policy)
    with pytest.raises(asyncio.CancelledError):
        await registry.should_stop(turn, run)


async def test_on_tool_event_publishes_tool_lifecycle_to_extensions() -> None:
    """_on_tool_event 既转发 SSE，又把同一事件发布进扩展事件流。"""
    events: list[tuple[str, dict[str, Any]]] = []
    sse_events: list[tuple[str, dict[str, Any]]] = []

    async def on_tool_event(event_type: str, payload: dict[str, Any]) -> None:
        sse_events.append((event_type, payload))

    async def listen(event: AgentEvent, _run: AgentRunContext) -> None:
        events.append((event.type, dict(event.payload)))

    extensions = ExtensionRegistry()
    extensions.add_event_listener(listen)
    agent = Agent(
        AgentConfig(tools=(), tool_runtime=None, on_tool_event=on_tool_event),
        extensions=extensions,
    )

    await agent._on_tool_event(
        "tool.started", {"tool_call_id": "c1", "tool_name": "search"}
    )
    await agent._on_tool_event(
        "tool.completed", {"tool_call_id": "c1", "status": "success"}
    )
    await agent._on_tool_event(
        "tool.completed", {"tool_call_id": "c2", "status": "denied"}
    )

    assert [item[0] for item in sse_events] == [
        "tool.started",
        "tool.completed",
        "tool.completed",
    ]
    assert [item[0] for item in events] == [
        "tool.started",
        "tool.completed",
        "tool.completed",
    ]
    assert events[1][1]["status"] == "success"
    # pending 状态维护未被桥接改动破坏
    assert "c1" not in agent.state.pending_tool_calls


async def test_on_tool_event_swallows_listener_failure() -> None:
    """扩展 listener 抛异常不应阻断 SSE 转发或工具事件主流程。"""
    sse: list[str] = []

    async def on_tool_event(event_type: str, _payload: dict[str, Any]) -> None:
        sse.append(event_type)

    async def boom(_event: AgentEvent, _run: AgentRunContext) -> None:
        raise RuntimeError("listener boom")

    extensions = ExtensionRegistry()
    extensions.add_event_listener(boom)
    agent = Agent(
        AgentConfig(tools=(), tool_runtime=None, on_tool_event=on_tool_event),
        extensions=extensions,
    )

    # 不抛：异常被吞，SSE 仍照常转发
    await agent._on_tool_event("tool.started", {"tool_call_id": "c1"})
    assert sse == ["tool.started"]