"""Metrics middleware for Prometheus integration."""

from __future__ import annotations

import time
from typing import Any, Dict

try:
    from prometheus_client import Counter, Histogram, Gauge
    HAS_PROMETHEUS = True
except ImportError:
    HAS_PROMETHEUS = False

from ..logging import get_logger
from .base import Middleware

logger = get_logger(__name__)


class MetricsMiddleware(Middleware):
    """Middleware for collecting Prometheus metrics.
    
    Tracks request counts, latencies, error rates, and active requests.
    """
    
    def __init__(self):
        """Initialize metrics middleware.
        
        Raises:
            ImportError: If prometheus_client is not installed
        """
        super().__init__("MetricsMiddleware")
        
        if not HAS_PROMETHEUS:
            logger.warning("prometheus_client_not_installed", 
                         message="Metrics middleware disabled")
            self.disable()
            return
        
        # Define metrics
        self.requests_total = Counter(
            "mcp_requests_total",
            "Total number of MCP requests",
            ["tool_name", "status"]
        )
        
        self.request_duration = Histogram(
            "mcp_request_duration_seconds",
            "MCP request duration in seconds",
            ["tool_name"],
            buckets=[0.01, 0.05, 0.1, 0.5, 1.0, 2.5, 5.0, 10.0]
        )
        
        self.requests_in_progress = Gauge(
            "mcp_requests_in_progress",
            "Number of MCP requests currently being processed",
            ["tool_name"]
        )
        
        self.errors_total = Counter(
            "mcp_errors_total",
            "Total number of MCP errors",
            ["tool_name", "error_type"]
        )
    
    async def before_request(self, request: Dict[str, Any]) -> Dict[str, Any]:
        """Record request start metrics.
        
        Args:
            request: The request data
            
        Returns:
            Request with timing data
        """
        if not self.is_enabled():
            return request
        
        tool_name = request.get("tool_name", "unknown")
        
        # Increment in-progress gauge
        self.requests_in_progress.labels(tool_name=tool_name).inc()
        
        # Store start time for duration calculation
        request["_metrics_start_time"] = time.time()
        
        return request
    
    async def after_response(self, response: Any, request: Dict[str, Any]) -> Any:
        """Record response metrics.
        
        Args:
            response: The response data
            request: The original request
            
        Returns:
            Unmodified response
        """
        if not self.is_enabled():
            return response
        
        tool_name = request.get("tool_name", "unknown")
        
        # Decrement in-progress gauge
        self.requests_in_progress.labels(tool_name=tool_name).dec()
        
        # Calculate duration
        start_time = request.get("_metrics_start_time")
        if start_time:
            duration = time.time() - start_time
            self.request_duration.labels(tool_name=tool_name).observe(duration)
        
        # Determine status
        is_success = True
        error_type = None
        
        if hasattr(response, "ok"):
            is_success = response.ok
            if not is_success and hasattr(response, "error"):
                error_info = response.error
                if isinstance(error_info, dict):
                    error_type = error_info.get("code", "unknown_error")
        elif isinstance(response, dict):
            is_success = response.get("ok", True)
            if not is_success:
                error_info = response.get("error", {})
                if isinstance(error_info, dict):
                    error_type = error_info.get("code", "unknown_error")
        
        # Record total requests
        status = "success" if is_success else "error"
        self.requests_total.labels(tool_name=tool_name, status=status).inc()
        
        # Record errors
        if not is_success and error_type:
            self.errors_total.labels(
                tool_name=tool_name,
                error_type=error_type
            ).inc()
        
        return response

