"""Port wrappers that provide observability and control.

These wrappers intercept port method calls to:
1. **Observability** - Emit events (LLMPromptSent, LLMResponseReceived, etc.)
2. **Control** - Enable policy-based control (rate limiting, caching, retry, fallback)
3. **Metrics** - Track duration and performance
4. **Extensibility** - Provide hooks for custom behavior

This allows centralized control over infrastructure interactions without
polluting business logic in nodes and macros.
"""

from typing import Any

from hexdag.kernel.context import get_current_node_name, get_observer_manager
from hexdag.kernel.logging import get_logger
from hexdag.kernel.orchestration.events import (
    LLMPromptSent,
    LLMResponseReceived,
    ToolCalled,
    ToolCompleted,
)
from hexdag.kernel.ports.llm import (
    LLM,
    LLMResponse,
    MessageList,
    SupportsFunctionCalling,
    SupportsGeneration,
    SupportsUsageTracking,
)
from hexdag.kernel.ports.tool_router import ToolRouter
from hexdag.kernel.utils.node_timer import Timer

logger = get_logger(__name__)


class ObservableLLMWrapper:
    """Wraps an LLM port for observability and policy-based control.

    This wrapper provides:
    - Automatic event emission for all LLM calls
    - Policy evaluation before/after calls (future: rate limiting, caching)
    - Duration tracking and performance metrics
    - Transparent forwarding to underlying LLM
    """

    def __init__(self, llm: LLM):
        """Initialize wrapper with underlying LLM port.

        Parameters
        ----------
        llm : LLM
            The underlying LLM port to wrap
        """
        self._llm = llm

    async def aresponse(self, messages: MessageList, **kwargs: Any) -> str | None:
        """Call LLM with observability and policy control.

        This method:
        1. Emits LLMPromptSent event (for observability)
        2. Calls underlying LLM
        3. Emits LLMResponseReceived event with duration
        4. (Future: Policy evaluation for rate limiting, caching, retry)

        Parameters
        ----------
        messages : MessageList
            Messages to send to LLM
        **kwargs : Any
            Additional LLM parameters

        Returns
        -------
        str | None
            LLM response text
        """
        node_name = get_current_node_name() or "unknown"

        # Emit prompt sent event (OBSERVABILITY)
        if observer_mgr := get_observer_manager():
            await observer_mgr.notify(
                LLMPromptSent(
                    node_name=node_name,
                    messages=[{"role": m.role, "content": m.content} for m in messages],
                )
            )

        # TODO(hexdag-team): Add policy evaluation here (CONTROL) #noqa: TD003
        # if policy_mgr := get_policy_manager():
        #     policy_response = await policy_mgr.evaluate(LLMCallPolicy(...))
        #     if policy_response.signal == PolicySignal.SKIP:
        #         return policy_response.data  # Cached response
        #     elif policy_response.signal == PolicySignal.FAIL:
        #         raise RateLimitError("LLM rate limit exceeded")

        # Call underlying LLM
        llm_timer = Timer()

        # Check if the adapter supports generation
        if isinstance(self._llm, SupportsGeneration):
            response = await self._llm.aresponse(messages, **kwargs)
        else:
            raise NotImplementedError(
                f"LLM adapter {type(self._llm).__name__} does not support text generation. "
                "It must implement SupportsGeneration protocol."
            )

        duration_ms = llm_timer.duration_ms

        # Extract usage via SupportsUsageTracking protocol
        usage_dict = None
        if isinstance(self._llm, SupportsUsageTracking) and (
            last_usage := self._llm.get_last_usage()
        ):
            usage_dict = {
                "input_tokens": last_usage.input_tokens,
                "output_tokens": last_usage.output_tokens,
                "total_tokens": last_usage.total_tokens,
            }

        # Emit response received event (OBSERVABILITY)
        if observer_mgr := get_observer_manager():
            await observer_mgr.notify(
                LLMResponseReceived(
                    node_name=node_name,
                    response=response or "",
                    duration_ms=duration_ms,
                    usage=usage_dict,
                )
            )

        return response

    async def aresponse_with_tools(
        self,
        messages: MessageList,
        tools: list[dict[str, Any]],
        tool_choice: str | dict[str, Any] = "auto",
        **kwargs: Any,
    ) -> LLMResponse:
        """Call LLM with tools and emit events.

        Parameters
        ----------
        messages : MessageList
            Messages to send to LLM
        tools : list[dict[str, Any]]
            Tool definitions
        tool_choice : str | dict[str, Any]
            Tool choice strategy
        **kwargs : Any
            Additional LLM parameters

        Returns
        -------
        LLMResponse
            Structured response with tool calls
        """
        node_name = get_current_node_name() or "unknown"

        # Emit prompt sent event
        if observer_mgr := get_observer_manager():
            await observer_mgr.notify(
                LLMPromptSent(
                    node_name=node_name,
                    messages=[{"role": m.role, "content": m.content} for m in messages],
                )
            )

        # Call underlying LLM (must support function calling)
        if not isinstance(self._llm, SupportsFunctionCalling):
            raise NotImplementedError(
                f"LLM adapter {type(self._llm).__name__} does not support function calling. "
                "It must implement SupportsFunctionCalling protocol."
            )

        llm_timer = Timer()
        response = await self._llm.aresponse_with_tools(messages, tools, tool_choice, **kwargs)
        duration_ms = llm_timer.duration_ms

        # Extract usage from LLMResponse or via SupportsUsageTracking protocol
        usage_dict = None
        if response.usage:
            usage_dict = {
                "input_tokens": response.usage.input_tokens,
                "output_tokens": response.usage.output_tokens,
                "total_tokens": response.usage.total_tokens,
            }
        elif isinstance(self._llm, SupportsUsageTracking) and (
            last_usage := self._llm.get_last_usage()
        ):
            usage_dict = {
                "input_tokens": last_usage.input_tokens,
                "output_tokens": last_usage.output_tokens,
                "total_tokens": last_usage.total_tokens,
            }

        # Emit response received event
        if observer_mgr := get_observer_manager():
            await observer_mgr.notify(
                LLMResponseReceived(
                    node_name=node_name,
                    response=response.content or "",
                    duration_ms=duration_ms,
                    usage=usage_dict,
                )
            )

        return response

    def __getattr__(self, name: str) -> Any:
        """Forward all other attribute access to underlying LLM."""
        return getattr(self._llm, name)


class ObservableToolRouterWrapper:
    """Wraps a ToolRouter port for observability and policy-based control.

    This wrapper provides:
    - Automatic event emission for all tool calls (ToolCalled, ToolCompleted)
    - Policy evaluation before/after calls (future: auth, rate limiting)
    - Duration tracking and performance metrics
    - Error handling and logging
    """

    def __init__(self, tool_router: Any):
        """Initialize wrapper with underlying tool router.

        Parameters
        ----------
        tool_router : Any
            The underlying tool router to wrap
        """
        self._tool_router = tool_router

    async def acall_tool(self, tool_name: str, params: dict[str, Any]) -> Any:
        """Call tool with observability and policy control.

        This method:
        1. Emits ToolCalled event (for observability)
        2. Calls underlying tool
        3. Emits ToolCompleted event with duration and result
        4. (Future: Policy evaluation for auth, rate limiting)

        Parameters
        ----------
        tool_name : str
            Name of the tool to call
        params : dict[str, Any]
            Tool parameters

        Returns
        -------
        Any
            Tool execution result
        """
        node_name = get_current_node_name() or "unknown"

        # Emit tool called event (OBSERVABILITY)
        if observer_mgr := get_observer_manager():
            await observer_mgr.notify(
                ToolCalled(node_name=node_name, tool_name=tool_name, params=params)
            )

        # TODO(hexdag-team): Add policy evaluation here (CONTROL) #noqa: TD003
        # if policy_mgr := get_policy_manager():
        #     policy_response = await policy_mgr.evaluate(ToolCallPolicy(...))
        #     if policy_response.signal == PolicySignal.FAIL:
        #         raise AuthorizationError(f"Tool '{tool_name}' not authorized")

        # Call underlying tool
        tool_timer = Timer()
        try:
            result = await self._tool_router.acall_tool(tool_name, params)

            # Emit tool completed event (OBSERVABILITY)
            if observer_mgr := get_observer_manager():
                await observer_mgr.notify(
                    ToolCompleted(
                        node_name=node_name,
                        tool_name=tool_name,
                        result=result,
                        duration_ms=tool_timer.duration_ms,
                    )
                )

            return result

        except Exception as e:
            logger.error(f"Tool '{tool_name}' failed in {tool_timer.duration_str}ms: {e}")

            # Still emit completed event with error info
            if observer_mgr := get_observer_manager():
                await observer_mgr.notify(
                    ToolCompleted(
                        node_name=node_name,
                        tool_name=tool_name,
                        result={"error": str(e)},
                        duration_ms=tool_timer.duration_ms,
                    )
                )

            raise

    def __getattr__(self, name: str) -> Any:
        """Forward all other attribute access to underlying tool router."""
        return getattr(self._tool_router, name)


def wrap_llm_port(llm: Any) -> Any:
    """Wrap LLM port with event emission if it implements SupportsGeneration.

    Parameters
    ----------
    llm : Any
        The LLM port to potentially wrap

    Returns
    -------
    Any
        Wrapped LLM if it supports generation, otherwise original object
    """
    if isinstance(llm, SupportsGeneration):
        return ObservableLLMWrapper(llm)
    return llm


def wrap_tool_router_port(tool_router: Any) -> Any:
    """Wrap ToolRouter port with event emission.

    Parameters
    ----------
    tool_router : Any
        The tool router to potentially wrap

    Returns
    -------
    Any
        Wrapped tool router if it is a ToolRouter, otherwise original object
    """
    if isinstance(tool_router, ToolRouter):
        return ObservableToolRouterWrapper(tool_router)
    return tool_router


def wrap_ports_with_observability(ports: dict[str, Any]) -> dict[str, Any]:
    """Wrap all ports with event-emitting wrappers.

    Uses ``isinstance`` protocol checks to detect wrappable ports,
    regardless of their dictionary key.

    Currently wraps:
    - LLM ports (``SupportsGeneration``): Emit LLMPromptSent/LLMResponseReceived
    - ToolRouter ports: Emit ToolCalled/ToolCompleted events

    Parameters
    ----------
    ports : dict[str, Any]
        Dictionary of port name to port adapter

    Returns
    -------
    dict[str, Any]
        Dictionary with wrapped ports
    """
    wrapped = {}
    for name, port in ports.items():
        if isinstance(port, SupportsGeneration):
            wrapped[name] = wrap_llm_port(port)
        elif isinstance(port, ToolRouter):
            wrapped[name] = wrap_tool_router_port(port)
        else:
            wrapped[name] = port
    return wrapped
