"""Agent 扩展契约与注册表。

扩展只接触本次运行的快照与 ToolResult，不持有数据库会话、密钥或安全治理对象。
"""

from __future__ import annotations

import asyncio
from collections.abc import Awaitable, Callable
from dataclasses import dataclass, field, replace
from typing import Any

from app.ai_provider.types.chat import TokenUsage
from app.ai_provider.types.messages import BaseMessage
from app.modules.tool_runtime.definition import ToolResult


class CancellationToken:
    """协作式取消信号；阻塞 await 仍由活动 asyncio.Task.cancel() 负责中断。"""

    def __init__(self) -> None:
        self._cancelled = False

    @property
    def cancelled(self) -> bool:
        return self._cancelled

    def cancel(self) -> None:
        self._cancelled = True

    def raise_if_cancelled(self) -> None:
        if self._cancelled:
            raise asyncio.CancelledError


@dataclass(frozen=True)
class AgentRunContext:
    """扩展可见的最小运行上下文。"""

    run_id: str | None = None
    session_id: str | None = None
    character_id: str | None = None
    trace_id: str | None = None
    metadata: dict[str, Any] = field(default_factory=dict)
    cancellation: CancellationToken = field(default_factory=CancellationToken)


@dataclass(frozen=True)
class AgentContext:
    """每轮发送给模型前可转换的上下文快照。"""

    messages: tuple[BaseMessage, ...]
    tools: tuple[dict[str, Any], ...]
    model_hint: str | None = None


@dataclass(frozen=True)
class AgentEvent:
    """Agent Core 内部生命周期事件。"""

    type: str
    payload: dict[str, Any] = field(default_factory=dict)


@dataclass(frozen=True)
class ToolCallContext:
    """工具调用 Hook 的稳定输入。"""

    tool_call_id: str
    name: str
    input: dict[str, Any]
    messages: tuple[BaseMessage, ...]


@dataclass(frozen=True)
class BeforeToolCallDecision:
    """执行前决策；扩展只能收紧或改参，不能绕过 ToolRuntime。"""

    allowed: bool = True
    reason: str | None = None
    input: dict[str, Any] | None = None


@dataclass(frozen=True)
class AfterToolCallDecision:
    """执行后结果覆盖与提前终止提示。"""

    result: ToolResult | None = None
    terminate: bool = False


@dataclass(frozen=True)
class TurnSnapshot:
    """一轮模型调用及其工具执行结果摘要。"""

    iteration: int
    content: str
    reasoning: str
    tool_calls: tuple[dict[str, Any], ...]
    tool_results: tuple[dict[str, Any], ...]
    usage: TokenUsage = field(default_factory=TokenUsage)


ContextTransformer = Callable[[AgentContext, AgentRunContext], Awaitable[AgentContext]]
BeforeToolCallHook = Callable[
    [ToolCallContext, AgentRunContext], Awaitable[BeforeToolCallDecision | None]
]
AfterToolCallHook = Callable[
    [ToolCallContext, ToolResult, AgentRunContext], Awaitable[AfterToolCallDecision | None]
]
TurnPolicy = Callable[[TurnSnapshot, AgentRunContext], Awaitable[bool]]
AgentEventListener = Callable[[AgentEvent, AgentRunContext], Awaitable[None]]


class ExtensionRegistry:
    """受约束的 Agent 扩展注册表；按注册顺序执行 Hook。"""

    def __init__(self) -> None:
        self._context_transformers: list[ContextTransformer] = []
        self._before_tool_hooks: list[BeforeToolCallHook] = []
        self._after_tool_hooks: list[AfterToolCallHook] = []
        self._turn_policies: list[TurnPolicy] = []
        self._event_listeners: list[AgentEventListener] = []

    def add_context_transformer(self, transformer: ContextTransformer) -> None:
        self._context_transformers.append(transformer)

    def add_before_tool_hook(self, hook: BeforeToolCallHook) -> None:
        self._before_tool_hooks.append(hook)

    def add_after_tool_hook(self, hook: AfterToolCallHook) -> None:
        self._after_tool_hooks.append(hook)

    def add_turn_policy(self, policy: TurnPolicy) -> None:
        self._turn_policies.append(policy)

    def add_event_listener(self, listener: AgentEventListener) -> Callable[[], None]:
        self._event_listeners.append(listener)

        def unsubscribe() -> None:
            if listener in self._event_listeners:
                self._event_listeners.remove(listener)

        return unsubscribe

    def extend(self, registry: ExtensionRegistry) -> None:
        """按原注册顺序合并另一组扩展。"""
        self._context_transformers.extend(registry._context_transformers)
        self._before_tool_hooks.extend(registry._before_tool_hooks)
        self._after_tool_hooks.extend(registry._after_tool_hooks)
        self._turn_policies.extend(registry._turn_policies)
        self._event_listeners.extend(registry._event_listeners)

    async def transform_context(
        self, context: AgentContext, run: AgentRunContext
    ) -> AgentContext:
        for transformer in self._context_transformers:
            run.cancellation.raise_if_cancelled()
            context = await transformer(context, run)
        return context

    async def before_tool_call(
        self, call: ToolCallContext, run: AgentRunContext
    ) -> BeforeToolCallDecision:
        current = call
        for hook in self._before_tool_hooks:
            run.cancellation.raise_if_cancelled()
            decision = await hook(current, run)
            if decision is None:
                continue
            if not decision.allowed:
                return decision
            if decision.input is not None:
                current = replace(current, input=decision.input)
        return BeforeToolCallDecision(input=current.input)

    async def after_tool_call(
        self,
        call: ToolCallContext,
        result: ToolResult,
        run: AgentRunContext,
    ) -> AfterToolCallDecision:
        current = result
        terminate = False
        for hook in self._after_tool_hooks:
            run.cancellation.raise_if_cancelled()
            decision = await hook(call, current, run)
            if decision is None:
                continue
            if decision.result is not None:
                current = decision.result
            terminate = terminate or decision.terminate
        return AfterToolCallDecision(result=current, terminate=terminate)

    async def should_stop(self, turn: TurnSnapshot, run: AgentRunContext) -> bool:
        for policy in self._turn_policies:
            run.cancellation.raise_if_cancelled()
            if await policy(turn, run):
                return True
        return False

    async def publish(self, event: AgentEvent, run: AgentRunContext) -> None:
        for listener in tuple(self._event_listeners):
            await listener(event, run)

    @property
    def enabled(self) -> bool:
        return any(
            (
                self._context_transformers,
                self._before_tool_hooks,
                self._after_tool_hooks,
                self._turn_policies,
                self._event_listeners,
            )
        )


__all__ = [
    "AfterToolCallDecision",
    "AgentContext",
    "AgentEvent",
    "AgentRunContext",
    "BeforeToolCallDecision",
    "CancellationToken",
    "ExtensionRegistry",
    "ToolCallContext",
    "TurnSnapshot",
]