from __future__ import annotations as _annotations

from collections.abc import AsyncIterable, Awaitable, Callable, Iterator
from contextlib import contextmanager
from contextvars import ContextVar
from dataclasses import dataclass, replace
from functools import wraps
from typing import Any
from uuid import uuid4

import anyio
from pydantic_ai import Agent as PydanticAgent
from pydantic_ai.capabilities import (
    AbstractCapability,
    CombinedCapability,
    Hooks,
    HookTimeoutError,
)
from pydantic_ai.messages import ModelResponse, ToolCallPart
from pydantic_ai.models import ModelRequestContext

from ..awaitables import resolve_value
from ..hook_projection import HookEvent, HookProjectionMap
from ..session.state import SessionTranscriptUpdate
from ._compat import (
    build_capability_override_value,
    entry_func,
    entry_timeout,
    entry_tool_filters,
    hook_registry,
    replace_hook_entry,
    root_capability,
    root_capability_override,
)

__all__ = ("RegisteredHookInfo", "list_agent_hooks", "observe_agent_hooks")

_INTERNAL_EVENT_NAMES = {
    "_on_event": "event",
    "handle_deferred_tool_calls": "deferred_tool_calls",
    "on_model_request_error": "model_request_error",
    "on_node_run_error": "node_run_error",
    "on_output_process_error": "output_process_error",
    "on_output_validate_error": "output_validate_error",
    "on_run_error": "run_error",
    "on_tool_execute_error": "tool_execute_error",
    "on_tool_validate_error": "tool_validate_error",
    "wrap_model_request": "model_request",
    "wrap_node_run": "node_run",
    "wrap_output_process": "output_process",
    "wrap_output_validate": "output_validate",
    "wrap_run": "run",
    "wrap_run_event_stream": "run_event_stream",
    "wrap_tool_execute": "tool_execute",
    "wrap_tool_validate": "tool_validate",
}
_SKIPPED_HOOK_MODULE = "pydantic_acp.bridges.hooks"


@dataclass(slots=True)
class _HookUpdateEmitter:
    write_update: Callable[[SessionTranscriptUpdate], Awaitable[None]]
    projection_map: HookProjectionMap
    run_id: str
    next_sequence_id: int = 1

    async def emit_start(
        self,
        *,
        event: HookEvent,
    ) -> str | None:
        tool_call_id = f"hook:{self.run_id}:{self.next_sequence_id}"
        self.next_sequence_id += 1
        start_update = self.projection_map.build_start_update(
            tool_call_id=tool_call_id,
            event=event,
        )
        if start_update is None:
            return None
        await self.write_update(start_update)
        return tool_call_id

    async def emit_progress(
        self,
        *,
        tool_call_id: str | None,
        event: HookEvent,
    ) -> None:
        if tool_call_id is None:
            return
        progress_update = self.projection_map.build_progress_update(
            tool_call_id=tool_call_id,
            event=event,
        )
        if progress_update is None:
            return
        await self.write_update(progress_update)

    async def emit(
        self,
        *,
        event: HookEvent,
    ) -> None:
        tool_call_id = await self.emit_start(event=event)
        await self.emit_progress(
            tool_call_id=tool_call_id,
            event=event,
        )


@dataclass(slots=True, frozen=True, kw_only=True)
class RegisteredHookInfo:
    event_id: str
    hook_name: str
    tool_filters: tuple[str, ...]


@contextmanager
def observe_agent_hooks(
    agent: PydanticAgent[Any, Any],
    *,
    write_update: Callable[[SessionTranscriptUpdate], Awaitable[None]],
    projection_map: HookProjectionMap | None = None,
) -> Iterator[None]:
    resolved_root_capability = _root_capability(agent)
    override_root = _override_root_capability(agent)
    if resolved_root_capability is None or override_root is None:
        yield
        return
    emitter = _HookUpdateEmitter(
        write_update=write_update,
        projection_map=projection_map or HookProjectionMap(),
        run_id=uuid4().hex,
    )
    wrapped_capability, changed = _wrap_combined_capability(
        resolved_root_capability,
        emitter=emitter,
    )
    if not changed:
        yield
        return
    token = override_root.set(build_capability_override_value(wrapped_capability))
    try:
        yield
    finally:
        override_root.reset(token)


def _root_capability(agent: PydanticAgent[Any, Any]) -> CombinedCapability[Any] | None:
    return root_capability(agent)


def list_agent_hooks(agent: PydanticAgent[Any, Any]) -> list[RegisteredHookInfo]:
    resolved_root_capability = _root_capability(agent)
    if resolved_root_capability is None:
        return []
    hook_infos: list[RegisteredHookInfo] = []
    for hooks in _iter_hooks(resolved_root_capability):
        registry = hook_registry(hooks)
        if registry is None:
            continue
        for registry_key, entries in registry.items():
            event_id = _INTERNAL_EVENT_NAMES.get(registry_key, registry_key)
            for entry in entries:
                func = entry_func(entry)
                if not callable(func):
                    continue
                if getattr(func, "__module__", "") == _SKIPPED_HOOK_MODULE:
                    continue
                hook_infos.append(
                    RegisteredHookInfo(
                        event_id=event_id,
                        hook_name=getattr(func, "__name__", "") or event_id,
                        tool_filters=_tool_filters(entry),
                    ),
                )
    return sorted(
        hook_infos,
        key=lambda hook_info: (
            hook_info.event_id,
            hook_info.hook_name,
            hook_info.tool_filters,
        ),
    )


def _override_root_capability(
    agent: PydanticAgent[Any, Any],
) -> ContextVar[Any] | None:
    return root_capability_override(agent)


def _wrap_capability(
    capability: AbstractCapability[Any],
    *,
    emitter: _HookUpdateEmitter,
) -> tuple[AbstractCapability[Any], bool]:
    if isinstance(capability, Hooks):
        return _wrap_hooks(capability, emitter=emitter)
    if isinstance(capability, CombinedCapability):
        changed = False
        wrapped_capabilities: list[AbstractCapability[Any]] = []
        for nested_capability in capability.capabilities:
            wrapped_capability, nested_changed = _wrap_capability(
                nested_capability,
                emitter=emitter,
            )
            wrapped_capabilities.append(wrapped_capability)
            changed = changed or nested_changed
        if not changed:
            return capability, False
        return replace(capability, capabilities=wrapped_capabilities), True
    return capability, False


def _iter_hooks(capability: AbstractCapability[Any]) -> Iterator[Hooks[Any]]:
    if isinstance(capability, Hooks):
        yield capability
        return
    if isinstance(capability, CombinedCapability):
        for nested_capability in capability.capabilities:
            yield from _iter_hooks(nested_capability)


def _wrap_combined_capability(
    capability: CombinedCapability[Any],
    *,
    emitter: _HookUpdateEmitter,
) -> tuple[CombinedCapability[Any], bool]:
    wrapped_capability, changed = _wrap_capability(capability, emitter=emitter)
    if not changed:
        return capability, False
    assert isinstance(wrapped_capability, CombinedCapability)
    return wrapped_capability, True


def _wrap_hooks(
    hooks: Hooks[Any],
    *,
    emitter: _HookUpdateEmitter,
) -> tuple[Hooks[Any], bool]:
    registry = hook_registry(hooks)
    if registry is None:
        return hooks, False
    wrapped_registry: dict[str, list[Any]] = {}
    changed = False
    for registry_key, entries in registry.items():
        wrapped_entries: list[Any] = []
        for entry in entries:
            wrapped_entry, entry_changed = _wrap_hook_entry(
                registry_key,
                entry,
                emitter=emitter,
            )
            wrapped_entries.append(wrapped_entry)
            changed = changed or entry_changed
        wrapped_registry[registry_key] = wrapped_entries
    if not changed:
        return hooks, False
    wrapped_hooks = Hooks[Any](ordering=hooks.get_ordering())
    wrapped_hooks._registry = wrapped_registry
    return wrapped_hooks, True


def _wrap_hook_entry(
    registry_key: str,
    entry: object,
    *,
    emitter: _HookUpdateEmitter,
) -> tuple[object, bool]:
    original_func = entry_func(entry)
    if not callable(original_func):
        return entry, False
    if getattr(original_func, "__module__", "") == _SKIPPED_HOOK_MODULE:
        return entry, False
    event_id = _INTERNAL_EVENT_NAMES.get(registry_key, registry_key)
    hook_name = getattr(original_func, "__name__", "") or event_id
    timeout = entry_timeout(entry)
    tool_filters = _tool_filters(entry)

    if registry_key == "wrap_run_event_stream":
        wrapped_func = _wrap_run_event_stream_entry(
            original_func,
            emitter=emitter,
            event_id=event_id,
            hook_name=hook_name,
            tool_filters=tool_filters,
        )
    else:
        wrapped_func = _wrap_standard_hook_entry(
            original_func,
            emitter=emitter,
            event_id=event_id,
            hook_name=hook_name,
            timeout=timeout,
            tool_filters=tool_filters,
        )

    wrapped_entry = replace_hook_entry(entry, func=wrapped_func, timeout=None)
    if wrapped_entry is None:
        return entry, False
    return wrapped_entry, True


def _wrap_run_event_stream_entry(
    original_func: Callable[..., Any],
    *,
    emitter: _HookUpdateEmitter,
    event_id: str,
    hook_name: str,
    tool_filters: tuple[str, ...],
) -> Callable[..., AsyncIterable[Any]]:
    @wraps(original_func)
    def wrapped_stream_hook(*args: Any, **kwargs: Any) -> AsyncIterable[Any]:
        tool_name = _tool_name(kwargs)
        start_event = HookEvent(
            event_id=event_id,
            hook_name=hook_name,
            tool_name=tool_name,
            tool_filters=tool_filters,
        )

        async def instrumented_stream() -> AsyncIterable[Any]:
            tool_call_id = await emitter.emit_start(event=start_event)
            try:
                result = original_func(*args, **kwargs)
                if not isinstance(result, AsyncIterable):
                    raise TypeError("run_event_stream hooks must return an async iterable")
                async for item in result:
                    yield item
            except BaseException as error:
                await emitter.emit_progress(
                    tool_call_id=tool_call_id,
                    event=HookEvent(
                        event_id=event_id,
                        hook_name=hook_name,
                        tool_name=tool_name,
                        tool_filters=tool_filters,
                        raw_output=_summarize_error(error),
                        status="failed",
                    ),
                )
                raise
            await emitter.emit_progress(
                tool_call_id=tool_call_id,
                event=HookEvent(
                    event_id=event_id,
                    hook_name=hook_name,
                    tool_name=tool_name,
                    tool_filters=tool_filters,
                    raw_output="stream wrapped",
                    status="completed",
                ),
            )

        return instrumented_stream()

    return wrapped_stream_hook


def _wrap_standard_hook_entry(
    original_func: Callable[..., Any],
    *,
    emitter: _HookUpdateEmitter,
    event_id: str,
    hook_name: str,
    timeout: float | None,
    tool_filters: tuple[str, ...],
) -> Callable[..., Awaitable[Any]]:
    @wraps(original_func)
    async def wrapped(*args: Any, **kwargs: Any) -> Any:
        tool_name = _tool_name(kwargs)
        start_event = HookEvent(
            event_id=event_id,
            hook_name=hook_name,
            tool_name=tool_name,
            tool_filters=tool_filters,
        )
        tool_call_id = await emitter.emit_start(event=start_event)
        try:
            result = await _call_hook_func(
                original_func,
                *args,
                timeout=timeout,
                hook_name=event_id,
                **kwargs,
            )
        except BaseException as error:
            await emitter.emit_progress(
                tool_call_id=tool_call_id,
                event=HookEvent(
                    event_id=event_id,
                    hook_name=hook_name,
                    tool_name=tool_name,
                    tool_filters=tool_filters,
                    raw_output=_summarize_error(error),
                    status="failed",
                ),
            )
            raise
        await emitter.emit_progress(
            tool_call_id=tool_call_id,
            event=HookEvent(
                event_id=event_id,
                hook_name=hook_name,
                tool_name=tool_name,
                tool_filters=tool_filters,
                raw_output=_summarize_result(result),
                status="completed",
            ),
        )
        return result

    return wrapped


async def _call_hook_func(
    func: Callable[..., Any],
    *args: Any,
    timeout: float | None,
    hook_name: str,
    **kwargs: Any,
) -> Any:
    if timeout is None:
        return await resolve_value(func(*args, **kwargs))
    try:
        with anyio.fail_after(timeout):
            return await resolve_value(func(*args, **kwargs))
    except TimeoutError:
        raise HookTimeoutError(
            hook_name=hook_name,
            func_name=getattr(func, "__name__", repr(func)),
            timeout=timeout,
        ) from None


def _tool_filters(entry: object) -> tuple[str, ...]:
    return entry_tool_filters(entry)


def _tool_name(kwargs: dict[str, Any]) -> str | None:
    call = kwargs.get("call")
    tool_name = getattr(call, "tool_name", None)
    if isinstance(tool_name, str):
        return tool_name
    return None


def _summarize_error(error: BaseException) -> str:
    return str(error) or type(error).__name__


def _summarize_result(result: Any) -> str:
    if result is None:
        return "completed"
    if isinstance(result, str):
        return result
    if isinstance(result, bool | int | float):
        return str(result)
    if isinstance(result, ModelRequestContext):
        return f"messages={len(result.messages)}"
    if isinstance(result, ModelResponse):
        return f"parts={len(result.parts)}"
    if isinstance(result, ToolCallPart):
        return result.tool_name
    return type(result).__name__
