"""Structured metrics for agent tool calls."""

import asyncio
import logging
import time
from typing import Any, Callable

from langchain.agents.middleware import AgentMiddleware
from langchain.tools.tool_node import ToolCallRequest

from agent.agent_config import AgentConfig
from utils.structured_log import emit_question_metric

logger = logging.getLogger("app")


class ToolMetricsMiddleware(AgentMiddleware):
    """Emit structured timing metrics for every tool call."""

    def __init__(self, config: AgentConfig):
        self.config = config

    def _emit_tool_metric(
        self,
        request: ToolCallRequest,
        *,
        started_at: float,
        status: str,
        error_type: str | None = None,
    ) -> None:
        tool_call = request.tool_call or {}
        tool_name = tool_call.get("name") or "unknown_tool"
        tool_call_id = tool_call.get("id")
        duration_ms = max(int((time.monotonic() - started_at) * 1000), 0)

        try:
            emit_question_metric(
                stage="catalog_agent.tool_call",
                status=status,
                duration_ms=duration_ms,
                trace_id=self.config.trace_id,
                ai_id=self.config.bot_id,
                session_id=self.config.session_id,
                robot_type="agent",
                model=self.config.model_name,
                stream=self.config.stream,
                error_type=error_type,
                extra={
                    "bot_id": self.config.bot_id,
                    "tool_name": tool_name,
                    "tool_call_id": tool_call_id,
                    "tool_response": self.config.tool_response,
                    "enable_thinking": self.config.enable_thinking,
                },
            )
        except Exception:
            logger.exception("Failed to emit tool metric for tool_name=%s", tool_name)

    def wrap_tool_call(
        self,
        request: ToolCallRequest,
        handler: Callable[[ToolCallRequest], Any],
    ) -> Any:
        started_at = time.monotonic()
        try:
            result = handler(request)
        except Exception as exc:
            self._emit_tool_metric(
                request,
                started_at=started_at,
                status="error",
                error_type=type(exc).__name__,
            )
            raise

        self._emit_tool_metric(request, started_at=started_at, status="success")
        return result

    async def awrap_tool_call(
        self,
        request: ToolCallRequest,
        handler: Callable[[ToolCallRequest], Any],
    ) -> Any:
        started_at = time.monotonic()
        try:
            result = await handler(request)
        except asyncio.CancelledError:
            self._emit_tool_metric(request, started_at=started_at, status="cancel")
            raise
        except Exception as exc:
            self._emit_tool_metric(
                request,
                started_at=started_at,
                status="error",
                error_type=type(exc).__name__,
            )
            raise

        self._emit_tool_metric(request, started_at=started_at, status="success")
        return result