from __future__ import annotations

from collections.abc import AsyncIterable, Callable, Sequence
from dataclasses import dataclass
from typing import TYPE_CHECKING, Any

from pydantic import ValidationError

from pydantic_ai._instructions import AgentInstructions, SourcedInstruction, normalize_instructions
from pydantic_ai._utils import aclose_all, replace_no_init
from pydantic_ai.exceptions import ModelRetry
from pydantic_ai.messages import AgentStreamEvent, ModelResponse, ToolCallPart
from pydantic_ai.tools import (
    AgentDepsT,
    AgentNativeTool,
    DeferredToolRequests,
    DeferredToolResults,
    RunContext,
    ToolDefinition,
)
from pydantic_ai.toolsets import AbstractToolset, AgentToolset

from ._on_event import collect_on_event_methods, marked_listens_to
from .abstract import (
    AbstractCapability,
    AgentModel,
    AgentNode,
    CapabilityDescription,
    NodeResult,
    RawOutput,
    RawToolArgs,
    ValidatedToolArgs,
    WrapModelRequestHandler,
    WrapNodeRunHandler,
    WrapOutputProcessHandler,
    WrapOutputValidateHandler,
    WrapRunHandler,
    WrapToolExecuteHandler,
    WrapToolValidateHandler,
)

if TYPE_CHECKING:
    from pydantic_ai.agent.abstract import AbstractAgent, AgentModelSettings
    from pydantic_ai.models import KnownModelName, Model, ModelRequestContext, ModelResolutionContext
    from pydantic_ai.output import OutputContext
    from pydantic_ai.run import AgentRunResult


def _registers_children(wrapped: AbstractCapability[Any], collected: Sequence[AbstractCapability[Any]]) -> bool:
    """Whether `collected` -- what `wrapped.apply` yielded -- is more than `wrapped` itself."""
    return len(collected) != 1 or collected[0] is not wrapped


@dataclass
class WrapperCapability(AbstractCapability[AgentDepsT]):
    """A capability that wraps another capability and delegates all methods.

    Analogous to [`WrapperToolset`][pydantic_ai.toolsets.WrapperToolset] for toolsets.
    Subclass and override specific methods to modify behavior while delegating the rest.

    When the wrapped capability returns a fresh instance from
    [`for_agent`][pydantic_ai.capabilities.AbstractCapability.for_agent] or
    [`for_run`][pydantic_ai.capabilities.AbstractCapability.for_run], the wrapper is rebound
    as a shallow copy holding the new `wrapped`: subclass state is carried over verbatim and
    `__init__`/`__post_init__` are not re-run. Compute values derived from `wrapped` on
    access (e.g. via a property) rather than caching them at construction, so they can't go
    stale across a rebind.
    """

    wrapped: AbstractCapability[AgentDepsT]

    def __post_init__(self) -> None:
        self.__adopt_wrapped_identity()

    # Name-mangled deliberately: this upholds a base-class invariant on rebinds, so a
    # subclass attribute of the same name must not be able to override it.
    def __adopt_wrapped_identity(self) -> None:
        # A wrapper is transparent by default: with no explicit `id` of its own, it adopts
        # the wrapped capability's `id` and `defer_loading`. This is what lets a wrapper sit
        # over a deferred capability without losing its deferral or its place in the load
        # catalog. `for_agent`/`for_run` re-run this on the rebound copy, so it re-resolves
        # against the new wrapped instance — e.g. one a `DynamicCapability` produced at run
        # time, whose `id` only becomes known once the factory has run.
        if self.id is None:
            self.id = self.wrapped.id
            self.defer_loading = self.wrapped.defer_loading

    def apply(self, visitor: Callable[[AbstractCapability[AgentDepsT]], None]) -> None:
        visitor(self)
        # Collected once and replayed rather than walking the subtree twice: two walks per level
        # turns a chain of `n` wrappers into `2**n` traversals, so a stack of `prefix_tools()`
        # calls stops resolving in any reasonable time. One walk per level keeps the cost of the
        # chain linear in its depth.
        wrapped_capabilities: list[AbstractCapability[AgentDepsT]] = []
        self.wrapped.apply(wrapped_capabilities.append)
        if _registers_children(self.wrapped, wrapped_capabilities):
            for capability in wrapped_capabilities:
                visitor(capability)

    def visit_and_replace(
        self, visitor: Callable[[AbstractCapability[AgentDepsT]], AbstractCapability[AgentDepsT] | None]
    ) -> AbstractCapability[AgentDepsT] | None:
        """Visit the wrapper first; a replaced or removed wrapper takes its subtree with it.

        When the wrapper survives, the visit descends into `wrapped` and this wrapper is rebuilt
        around whatever remains; see
        [`AbstractCapability.visit_and_replace`][pydantic_ai.capabilities.AbstractCapability.visit_and_replace]
        for the tree-walking contract.
        """
        replacement = visitor(self)
        if replacement is not self:
            # The wrapper is what's registered for the subtree, so replacing or removing it takes
            # the subtree with it — visiting children the caller just discarded would be pointless.
            return replacement
        # A wrapper over a leaf capability is the registered proxy for that leaf: if the wrapped
        # subtree registers nothing of its own, there is nothing beneath this wrapper to visit.
        wrapped_capabilities: list[AbstractCapability[AgentDepsT]] = []
        self.wrapped.apply(wrapped_capabilities.append)
        if not _registers_children(self.wrapped, wrapped_capabilities):
            return self
        new_wrapped = self.wrapped.visit_and_replace(visitor)
        if new_wrapped is None:
            # `wrapped` is required, and a wrapper whose subtree is gone has nothing left to modify.
            return None
        if new_wrapped is self.wrapped:
            return self
        new_self = replace_no_init(self, wrapped=new_wrapped)
        new_self.__adopt_wrapped_identity()
        return new_self

    @classmethod
    def get_serialization_name(cls) -> str | None:
        return None

    def get_description(self) -> CapabilityDescription[AgentDepsT] | None:
        return self.description if self.description is not None else self.wrapped.get_description()

    @property
    def _has_wrap_node_run(self) -> bool:
        return type(self).wrap_node_run is not WrapperCapability.wrap_node_run or self.wrapped._has_wrap_node_run

    @property
    def _has_on_node_run_error(self) -> bool:
        return (
            type(self).on_node_run_error is not WrapperCapability.on_node_run_error
            or self.wrapped._has_on_node_run_error
        )

    @property
    def _has_wrap_model_request(self) -> bool:
        return (
            type(self).wrap_model_request is not WrapperCapability.wrap_model_request
            or self.wrapped._has_wrap_model_request
        )

    @property
    def _has_on_model_request_error(self) -> bool:
        return (
            type(self).on_model_request_error is not WrapperCapability.on_model_request_error
            or self.wrapped._has_on_model_request_error
        )

    @property
    def has_wrap_run_event_stream(self) -> bool:
        return (
            type(self).wrap_run_event_stream is not WrapperCapability.wrap_run_event_stream
            or self.wrapped.has_wrap_run_event_stream
        )

    @property
    def has_on_event(self) -> bool:
        return (
            type(self).on_event is not WrapperCapability.on_event
            or bool(collect_on_event_methods(type(self)))
            or self.wrapped.has_on_event
        )

    def listens_to(self, event: AgentStreamEvent) -> bool:
        return (
            type(self).on_event is not WrapperCapability.on_event
            or marked_listens_to(type(self), event)
            or self.wrapped.listens_to(event)
        )

    @property
    def _emits_app_events(self) -> bool:
        # The `RunContext.emit` gate must see through wrappers: wrapping an app-facing
        # `Hooks`/`ProcessEventStream` must not revoke its user callbacks' permission to emit
        # `CustomEvent`s.
        return self.wrapped._emits_app_events

    def for_agent(self, agent: AbstractAgent[AgentDepsT, Any]) -> AbstractCapability[AgentDepsT]:
        new_wrapped = self.wrapped.for_agent(agent)
        if new_wrapped is self.wrapped:
            return self
        new_self = replace_no_init(self, wrapped=new_wrapped)
        new_self.__adopt_wrapped_identity()
        return new_self

    async def for_run(self, ctx: RunContext[AgentDepsT]) -> AbstractCapability[AgentDepsT]:
        new_wrapped = await self.wrapped.for_run(ctx)
        if new_wrapped is self.wrapped:
            return self
        new_self = replace_no_init(self, wrapped=new_wrapped)
        new_self.__adopt_wrapped_identity()
        return new_self

    def _prepare_run_context(self, ctx: RunContext[AgentDepsT]) -> None:
        self.wrapped._prepare_run_context(ctx)

    def _validate_runtime_capabilities(
        self, ctx: RunContext[AgentDepsT], capabilities: Sequence[AbstractCapability[AgentDepsT]]
    ) -> None:
        self.wrapped._validate_runtime_capabilities(ctx, capabilities)

    # --- Get methods ---

    def get_instructions(self) -> AgentInstructions[AgentDepsT] | None:
        return self.wrapped.get_instructions()

    def _collect_instructions(self) -> list[SourcedInstruction[AgentDepsT]]:
        if type(self).get_instructions is not WrapperCapability.get_instructions:
            relayed = self.wrapped._collect_instructions()
            return self._attribute_container_instructions(normalize_instructions(self.get_instructions()), relayed)
        # Pass through the wrapped capability's own attribution: a wrapper adopts the id of the
        # capability it wraps, but a wrapper over a container has none to adopt and would
        # otherwise flatten every leaf's contribution into one unaddressable part.
        return self.wrapped._collect_instructions()

    def get_model_settings(self) -> AgentModelSettings[AgentDepsT] | None:
        return self.wrapped.get_model_settings()

    def get_model(self) -> AgentModel[AgentDepsT] | None:
        return self.wrapped.get_model()

    @property
    def has_resolve_model_id(self) -> bool:
        return (
            type(self).resolve_model_id is not WrapperCapability.resolve_model_id or self.wrapped.has_resolve_model_id
        )

    async def resolve_model_id(
        self,
        ctx: ModelResolutionContext[AgentDepsT],
        *,
        model_id: KnownModelName | str,
    ) -> Model | None:
        return await self.wrapped.resolve_model_id(ctx, model_id=model_id)

    def get_toolset(self) -> AgentToolset[AgentDepsT] | None:
        return self.wrapped.get_toolset()

    def get_native_tools(self) -> Sequence[AgentNativeTool[AgentDepsT]]:
        return self.wrapped.get_native_tools()

    def get_wrapper_toolset(self, toolset: AbstractToolset[AgentDepsT]) -> AbstractToolset[AgentDepsT] | None:
        return self.wrapped.get_wrapper_toolset(toolset)

    async def prepare_tools(
        self,
        ctx: RunContext[AgentDepsT],
        tool_defs: list[ToolDefinition],
    ) -> list[ToolDefinition]:
        return await self.wrapped.prepare_tools(ctx, tool_defs)

    async def prepare_output_tools(
        self,
        ctx: RunContext[AgentDepsT],
        tool_defs: list[ToolDefinition],
    ) -> list[ToolDefinition]:
        return await self.wrapped.prepare_output_tools(ctx, tool_defs)

    # --- Run lifecycle hooks ---

    async def before_run(self, ctx: RunContext[AgentDepsT]) -> None:
        await self.wrapped.before_run(ctx)

    async def after_run(
        self,
        ctx: RunContext[AgentDepsT],
        *,
        result: AgentRunResult[Any],
    ) -> AgentRunResult[Any]:
        return await self.wrapped.after_run(ctx, result=result)

    async def wrap_run(
        self,
        ctx: RunContext[AgentDepsT],
        *,
        handler: WrapRunHandler,
    ) -> AgentRunResult[Any]:
        return await self.wrapped.wrap_run(ctx, handler=handler)

    async def on_run_error(
        self,
        ctx: RunContext[AgentDepsT],
        *,
        error: BaseException,
    ) -> AgentRunResult[Any]:
        return await self.wrapped.on_run_error(ctx, error=error)

    # --- Node run lifecycle hooks ---

    async def before_node_run(
        self,
        ctx: RunContext[AgentDepsT],
        *,
        node: AgentNode[AgentDepsT],
    ) -> AgentNode[AgentDepsT]:
        return await self.wrapped.before_node_run(ctx, node=node)

    async def after_node_run(
        self,
        ctx: RunContext[AgentDepsT],
        *,
        node: AgentNode[AgentDepsT],
        result: NodeResult[AgentDepsT],
    ) -> NodeResult[AgentDepsT]:
        return await self.wrapped.after_node_run(ctx, node=node, result=result)

    async def wrap_node_run(
        self,
        ctx: RunContext[AgentDepsT],
        *,
        node: AgentNode[AgentDepsT],
        handler: WrapNodeRunHandler[AgentDepsT],
    ) -> NodeResult[AgentDepsT]:
        return await self.wrapped.wrap_node_run(ctx, node=node, handler=handler)

    async def on_node_run_error(
        self,
        ctx: RunContext[AgentDepsT],
        *,
        node: AgentNode[AgentDepsT],
        error: Exception,
    ) -> NodeResult[AgentDepsT]:
        return await self.wrapped.on_node_run_error(ctx, node=node, error=error)

    # --- Event hooks ---

    async def on_event(self, ctx: RunContext[AgentDepsT], *, event: AgentStreamEvent) -> None:
        await super().on_event(ctx, event=event)
        if self.wrapped.listens_to(event):
            await self.wrapped.on_event(ctx, event=event)

    async def wrap_run_event_stream(
        self,
        ctx: RunContext[AgentDepsT],
        *,
        stream: AsyncIterable[AgentStreamEvent],
    ) -> AsyncIterable[AgentStreamEvent]:
        wrapped_stream = self.wrapped.wrap_run_event_stream(ctx, stream=stream)
        try:
            async for event in wrapped_stream:
                yield event
        finally:
            await aclose_all((wrapped_stream, stream))

    # --- Model request lifecycle hooks ---

    async def before_model_request(
        self,
        ctx: RunContext[AgentDepsT],
        request_context: ModelRequestContext,
    ) -> ModelRequestContext:
        return await self.wrapped.before_model_request(ctx, request_context)

    async def after_model_request(
        self,
        ctx: RunContext[AgentDepsT],
        *,
        request_context: ModelRequestContext,
        response: ModelResponse,
    ) -> ModelResponse:
        return await self.wrapped.after_model_request(ctx, request_context=request_context, response=response)

    async def wrap_model_request(
        self,
        ctx: RunContext[AgentDepsT],
        *,
        request_context: ModelRequestContext,
        handler: WrapModelRequestHandler,
    ) -> ModelResponse:
        return await self.wrapped.wrap_model_request(ctx, request_context=request_context, handler=handler)

    async def on_model_request_error(
        self,
        ctx: RunContext[AgentDepsT],
        *,
        request_context: ModelRequestContext,
        error: Exception,
    ) -> ModelResponse:
        return await self.wrapped.on_model_request_error(ctx, request_context=request_context, error=error)

    # --- Tool validate lifecycle hooks ---

    async def before_tool_validate(
        self,
        ctx: RunContext[AgentDepsT],
        *,
        call: ToolCallPart,
        tool_def: ToolDefinition,
        args: RawToolArgs,
    ) -> RawToolArgs:
        return await self.wrapped.before_tool_validate(ctx, call=call, tool_def=tool_def, args=args)

    async def after_tool_validate(
        self,
        ctx: RunContext[AgentDepsT],
        *,
        call: ToolCallPart,
        tool_def: ToolDefinition,
        args: ValidatedToolArgs,
    ) -> ValidatedToolArgs:
        return await self.wrapped.after_tool_validate(ctx, call=call, tool_def=tool_def, args=args)

    async def wrap_tool_validate(
        self,
        ctx: RunContext[AgentDepsT],
        *,
        call: ToolCallPart,
        tool_def: ToolDefinition,
        args: RawToolArgs,
        handler: WrapToolValidateHandler,
    ) -> ValidatedToolArgs:
        return await self.wrapped.wrap_tool_validate(ctx, call=call, tool_def=tool_def, args=args, handler=handler)

    async def on_tool_validate_error(
        self,
        ctx: RunContext[AgentDepsT],
        *,
        call: ToolCallPart,
        tool_def: ToolDefinition,
        args: RawToolArgs,
        error: ValidationError | ModelRetry,
    ) -> ValidatedToolArgs:
        return await self.wrapped.on_tool_validate_error(ctx, call=call, tool_def=tool_def, args=args, error=error)

    # --- Tool execute lifecycle hooks ---

    async def before_tool_execute(
        self,
        ctx: RunContext[AgentDepsT],
        *,
        call: ToolCallPart,
        tool_def: ToolDefinition,
        args: ValidatedToolArgs,
    ) -> ValidatedToolArgs:
        return await self.wrapped.before_tool_execute(ctx, call=call, tool_def=tool_def, args=args)

    async def after_tool_execute(
        self,
        ctx: RunContext[AgentDepsT],
        *,
        call: ToolCallPart,
        tool_def: ToolDefinition,
        args: ValidatedToolArgs,
        result: Any,
    ) -> Any:
        return await self.wrapped.after_tool_execute(ctx, call=call, tool_def=tool_def, args=args, result=result)

    async def wrap_tool_execute(
        self,
        ctx: RunContext[AgentDepsT],
        *,
        call: ToolCallPart,
        tool_def: ToolDefinition,
        args: ValidatedToolArgs,
        handler: WrapToolExecuteHandler,
    ) -> Any:
        return await self.wrapped.wrap_tool_execute(ctx, call=call, tool_def=tool_def, args=args, handler=handler)

    async def on_tool_execute_error(
        self,
        ctx: RunContext[AgentDepsT],
        *,
        call: ToolCallPart,
        tool_def: ToolDefinition,
        args: ValidatedToolArgs,
        error: Exception,
    ) -> Any:
        return await self.wrapped.on_tool_execute_error(ctx, call=call, tool_def=tool_def, args=args, error=error)

    # --- Output validate lifecycle hooks ---

    async def before_output_validate(
        self,
        ctx: RunContext[AgentDepsT],
        *,
        output_context: OutputContext,
        output: RawOutput,
    ) -> RawOutput:
        return await self.wrapped.before_output_validate(ctx, output_context=output_context, output=output)

    async def after_output_validate(
        self,
        ctx: RunContext[AgentDepsT],
        *,
        output_context: OutputContext,
        output: Any,
    ) -> Any:
        return await self.wrapped.after_output_validate(ctx, output_context=output_context, output=output)

    async def wrap_output_validate(
        self,
        ctx: RunContext[AgentDepsT],
        *,
        output_context: OutputContext,
        output: RawOutput,
        handler: WrapOutputValidateHandler,
    ) -> Any:
        return await self.wrapped.wrap_output_validate(
            ctx, output_context=output_context, output=output, handler=handler
        )

    async def on_output_validate_error(
        self,
        ctx: RunContext[AgentDepsT],
        *,
        output_context: OutputContext,
        output: RawOutput,
        error: ValidationError | ModelRetry,
    ) -> Any:
        return await self.wrapped.on_output_validate_error(
            ctx, output_context=output_context, output=output, error=error
        )

    # --- Output process lifecycle hooks ---

    async def before_output_process(
        self,
        ctx: RunContext[AgentDepsT],
        *,
        output_context: OutputContext,
        output: Any,
    ) -> Any:
        return await self.wrapped.before_output_process(ctx, output_context=output_context, output=output)

    async def after_output_process(
        self,
        ctx: RunContext[AgentDepsT],
        *,
        output_context: OutputContext,
        output: Any,
    ) -> Any:
        return await self.wrapped.after_output_process(ctx, output_context=output_context, output=output)

    async def wrap_output_process(
        self,
        ctx: RunContext[AgentDepsT],
        *,
        output_context: OutputContext,
        output: Any,
        handler: WrapOutputProcessHandler,
    ) -> Any:
        return await self.wrapped.wrap_output_process(
            ctx, output_context=output_context, output=output, handler=handler
        )

    async def on_output_process_error(
        self,
        ctx: RunContext[AgentDepsT],
        *,
        output_context: OutputContext,
        output: Any,
        error: Exception,
    ) -> Any:
        return await self.wrapped.on_output_process_error(
            ctx, output_context=output_context, output=output, error=error
        )

    async def handle_deferred_tool_calls(
        self,
        ctx: RunContext[AgentDepsT],
        *,
        requests: DeferredToolRequests,
    ) -> DeferredToolResults | None:
        return await self.wrapped.handle_deferred_tool_calls(ctx, requests=requests)