"""Pi agent-core 风格的有状态 Agent 门面。"""

from __future__ import annotations

import asyncio
from collections.abc import Awaitable, Callable
from contextlib import suppress
from dataclasses import dataclass, field, replace
from enum import StrEnum
from typing import Any

from app.ai_provider.types.chat import ChatChunk, TokenUsage
from app.ai_provider.types.messages import BaseMessage, HumanMessage
from app.application.agent.confirmation import ConfirmationRegistry
from app.application.agent.extensions import (
    AgentEvent,
    AgentRunContext,
    CancellationToken,
    ExtensionRegistry,
)
from app.application.agent.runtime import AgentTurnResult, run_agent_turn
from app.modules.tool_runtime.definition import ToolRuntimeLike


class AgentStatus(StrEnum):
    idle = "idle"
    running = "running"
    completed = "completed"
    failed = "failed"
    cancelled = "cancelled"


class QueueMode(StrEnum):
    all = "all"
    one_at_a_time = "one-at-a-time"


@dataclass
class AgentState:
    """有状态 Agent 的进程内快照。"""

    messages: list[BaseMessage] = field(default_factory=list)
    status: AgentStatus = AgentStatus.idle
    streaming_message: str = ""
    pending_tool_calls: set[str] = field(default_factory=set)
    error_message: str | None = None
    usage: TokenUsage = field(default_factory=TokenUsage)


class PendingMessageQueue:
    """steering/follow-up 共用的小型 FIFO。"""

    def __init__(self, mode: QueueMode = QueueMode.one_at_a_time) -> None:
        self.mode = mode
        self._messages: list[BaseMessage] = []

    def put(self, message: BaseMessage) -> None:
        self._messages.append(message)

    def drain(self) -> list[BaseMessage]:
        if not self._messages:
            return []
        if self.mode is QueueMode.all:
            messages, self._messages = self._messages, []
            return messages
        return [self._messages.pop(0)]

    def clear(self) -> None:
        self._messages.clear()

    def __bool__(self) -> bool:
        return bool(self._messages)


@dataclass(frozen=True)
class AgentConfig:
    """Stateful Agent 的稳定运行配置。"""

    tools: tuple[dict[str, Any], ...]
    tool_runtime: ToolRuntimeLike | None
    model_hint: str | None = None
    thinking_enabled: bool | None = None
    reasoning_effort: str | None = None
    max_iterations: int = 8
    confirmation_registry: ConfirmationRegistry | None = None
    confirmation_timeout: float = 300.0
    on_turn: Callable[[], Awaitable[None]] | None = None
    on_chunk: Callable[[ChatChunk], Awaitable[None]] | None = None
    on_tool_event: Callable[[str, dict[str, Any]], Awaitable[None]] | None = None
    on_confirm_timeout: Callable[[str], Awaitable[None]] | None = None


class Agent:
    """管理状态、消息队列、取消与扩展 Hook；低层循环仍由 run_agent_turn 负责。"""

    def __init__(
        self,
        config: AgentConfig,
        *,
        messages: list[BaseMessage] | None = None,
        run_context: AgentRunContext | None = None,
        extensions: ExtensionRegistry | None = None,
        runner: Callable[..., Awaitable[AgentTurnResult]] = run_agent_turn,
    ) -> None:
        self.config = config
        self.state = AgentState(messages=list(messages or []))
        self.run_context = run_context or AgentRunContext()
        self.extensions = extensions or ExtensionRegistry()
        self._runner = runner
        self.steering = PendingMessageQueue()
        self.follow_ups = PendingMessageQueue()
        self._active_task: asyncio.Task[AgentTurnResult] | None = None
        self._accepting_messages = False
        self._idle = asyncio.Event()
        self._idle.set()

    def steer(self, message: str | BaseMessage) -> bool:
        if not self._accepting_messages:
            return False
        self.steering.put(self._as_message(message))
        return True

    def follow_up(self, message: str | BaseMessage) -> bool:
        if not self._accepting_messages:
            return False
        self.follow_ups.put(self._as_message(message))
        return True

    async def prompt(self, message: str | BaseMessage) -> AgentTurnResult:
        if self.state.status is AgentStatus.running:
            raise RuntimeError("Agent 正在运行，请使用 steer 或 follow_up")
        self.state.messages.append(self._as_message(message))
        return await self._run()

    async def continue_run(self) -> AgentTurnResult:
        if self.state.status is AgentStatus.running:
            raise RuntimeError("Agent 正在运行")
        return await self._run()

    async def abort(self) -> bool:
        task = self._active_task
        if task is None or task.done():
            return False
        self.run_context.cancellation.cancel()
        task.cancel()
        return True

    async def wait_for_idle(self) -> None:
        await self._idle.wait()

    async def reset(self) -> None:
        await self.abort()
        await self.wait_for_idle()
        self.state = AgentState()
        self.steering.clear()
        self.follow_ups.clear()

    async def _run(self) -> AgentTurnResult:
        if self.run_context.cancellation.cancelled:
            self.run_context = replace(self.run_context, cancellation=CancellationToken())
        self.state.status = AgentStatus.running
        self.state.streaming_message = ""
        self.state.pending_tool_calls.clear()
        self.state.error_message = None
        self._accepting_messages = True
        self._idle.clear()
        try:
            await self.extensions.publish(AgentEvent("agent.started"), self.run_context)
            task = asyncio.create_task(self._run_loop())
            self._active_task = task
            result = await task
            self._accepting_messages = False
            self.state.status = AgentStatus.completed
            self.state.streaming_message = result.content
            self.state.usage = result.usage
            await self.extensions.publish(
                AgentEvent("agent.ended", {"iterations": result.iterations}), self.run_context
            )
            return result
        except asyncio.CancelledError:
            self._accepting_messages = False
            self.state.status = AgentStatus.cancelled
            await self.extensions.publish(AgentEvent("agent.cancelled"), self.run_context)
            raise
        except Exception as exc:
            self._accepting_messages = False
            self.state.status = AgentStatus.failed
            self.state.error_message = str(exc)
            with suppress(Exception):
                await self.extensions.publish(
                    AgentEvent("agent.failed", {"error": str(exc)}), self.run_context
                )
            raise
        finally:
            self._accepting_messages = False
            with suppress(Exception):
                await self.extensions.publish(
                    AgentEvent("agent.settled", {"status": self.state.status.value}),
                    self.run_context,
                )
            self._active_task = None
            self._idle.set()

    async def _run_loop(self) -> AgentTurnResult:
        remaining = self.config.max_iterations
        total_iterations = 0
        total_usage = TokenUsage()
        total_tool_history: list[dict[str, Any]] = []
        final = AgentTurnResult(content="", reasoning="")
        while remaining > 0:
            working_messages = list(self.state.messages)
            final = await self._runner(
                messages=working_messages,
                tools=list(self.config.tools),
                tool_runtime=self.config.tool_runtime,
                model_hint=self.config.model_hint,
                thinking_enabled=self.config.thinking_enabled,
                reasoning_effort=self.config.reasoning_effort,
                max_iterations=remaining,
                run_id=self.run_context.run_id,
                character_id=self.run_context.character_id,
                trace_id=self.run_context.trace_id,
                on_turn=self.config.on_turn,
                on_chunk=self._on_chunk,
                on_tool_event=self._on_tool_event,
                confirmation_registry=self.config.confirmation_registry,
                confirmation_timeout=self.config.confirmation_timeout,
                on_confirm_timeout=self.config.on_confirm_timeout,
                extensions=self.extensions,
                run_context=self.run_context,
                get_steering_messages=self.steering.drain,
            )
            self.state.messages = working_messages
            total_iterations += final.iterations
            total_usage += final.usage
            total_tool_history.extend(final.tool_history)
            remaining -= final.iterations
            if final.truncated:
                return AgentTurnResult(
                    content=final.content,
                    reasoning=final.reasoning,
                    tool_history=total_tool_history,
                    iterations=total_iterations,
                    truncated=True,
                    usage=total_usage,
                )
            if remaining <= 0:
                return AgentTurnResult(
                    content=final.content,
                    reasoning=final.reasoning,
                    tool_history=total_tool_history,
                    iterations=total_iterations,
                    truncated=bool(self.follow_ups),
                    usage=total_usage,
                )
            follow_ups = self.follow_ups.drain()
            if not follow_ups:
                return AgentTurnResult(
                    content=final.content,
                    reasoning=final.reasoning,
                    tool_history=total_tool_history,
                    iterations=total_iterations,
                    truncated=final.truncated,
                    usage=total_usage,
                )
            self.state.messages.extend(follow_ups)
        return final

    async def _on_chunk(self, chunk: ChatChunk) -> None:
        if chunk.delta_content:
            self.state.streaming_message += chunk.delta_content
        if self.config.on_chunk is not None:
            await self.config.on_chunk(chunk)

    async def _on_tool_event(self, event_type: str, payload: dict[str, Any]) -> None:
        call_id = str(payload.get("tool_call_id", ""))
        if event_type == "tool.started" and call_id:
            self.state.pending_tool_calls.add(call_id)
        elif event_type in ("tool.completed", "tool.failed") and call_id:
            self.state.pending_tool_calls.discard(call_id)
        # 先转发 SSE（用户可见流优先），再把同一事件发布进扩展事件流，
        # 让 AgentEventListener 也能观测工具级生命周期。
        # 扩展 listener 异常被吞掉，避免观测副菜阻断工具事件主流程。
        if self.config.on_tool_event is not None:
            await self.config.on_tool_event(event_type, payload)
        with suppress(Exception):
            await self.extensions.publish(
                AgentEvent(event_type, payload), self.run_context
            )

    @staticmethod
    def _as_message(message: str | BaseMessage) -> BaseMessage:
        return HumanMessage(content=message) if isinstance(message, str) else message


__all__ = [
    "Agent",
    "AgentConfig",
    "AgentState",
    "AgentStatus",
    "PendingMessageQueue",
    "QueueMode",
]