src/monitoring/metrics.py

"""Monitoring and metrics collection for Grant Analyst API.

Tracks:
- P50/P95 latency per endpoint
- Token usage per request
- Cache hit rates
- Model distribution (nano/mini/main)
- Exports to Prometheus format or JSON logs
"""

import time
import logging
from typing import Dict, List, Optional, Any
from dataclasses import dataclass, field
from collections import defaultdict
from datetime import datetime
import threading
import json

logger = logging.getLogger(__name__)


@dataclass
class LatencyMetrics:
    """Latency metrics for an endpoint."""
    samples: List[float] = field(default_factory=list)
    total_requests: int = 0

    def add_sample(self, latency_ms: float):
        """Add a latency sample."""
        self.samples.append(latency_ms)
        self.total_requests += 1

        # Keep only last 1000 samples to prevent memory growth
        if len(self.samples) > 1000:
            self.samples = self.samples[-1000:]

    def get_percentile(self, percentile: int) -> float:
        """Get percentile latency (50, 95, 99, etc.)."""
        if not self.samples:
            return 0.0

        sorted_samples = sorted(self.samples)
        idx = int(len(sorted_samples) * (percentile / 100.0))
        return sorted_samples[min(idx, len(sorted_samples) - 1)]

    def get_avg(self) -> float:
        """Get average latency."""
        if not self.samples:
            return 0.0
        return sum(self.samples) / len(self.samples)


@dataclass
class TokenMetrics:
    """Token usage metrics."""
    total_prompt_tokens: int = 0
    total_completion_tokens: int = 0
    total_requests: int = 0

    def add_usage(self, prompt_tokens: int, completion_tokens: int):
        """Record token usage for a request."""
        self.total_prompt_tokens += prompt_tokens
        self.total_completion_tokens += completion_tokens
        self.total_requests += 1

    @property
    def total_tokens(self) -> int:
        """Total tokens used."""
        return self.total_prompt_tokens + self.total_completion_tokens

    @property
    def avg_tokens_per_request(self) -> float:
        """Average tokens per request."""
        if self.total_requests == 0:
            return 0.0
        return self.total_tokens / self.total_requests


@dataclass
class CacheMetrics:
    """Cache hit/miss metrics."""
    hits: int = 0
    misses: int = 0

    def record_hit(self):
        """Record a cache hit."""
        self.hits += 1

    def record_miss(self):
        """Record a cache miss."""
        self.misses += 1

    @property
    def total(self) -> int:
        """Total cache accesses."""
        return self.hits + self.misses

    @property
    def hit_rate(self) -> float:
        """Cache hit rate as percentage."""
        if self.total == 0:
            return 0.0
        return (self.hits / self.total) * 100


@dataclass
class ModelMetrics:
    """Model usage distribution."""
    model_counts: Dict[str, int] = field(default_factory=lambda: defaultdict(int))

    def record_model_use(self, model_type: str):
        """Record a model usage."""
        self.model_counts[model_type] += 1

    @property
    def total_requests(self) -> int:
        """Total model requests."""
        return sum(self.model_counts.values())

    def get_distribution(self) -> Dict[str, float]:
        """Get model distribution as percentages."""
        if self.total_requests == 0:
            return {}

        return {
            model: (count / self.total_requests) * 100
            for model, count in self.model_counts.items()
        }


class MetricsCollector:
    """Central metrics collector for the API."""

    def __init__(self):
        """Initialize metrics collector."""
        self._lock = threading.Lock()

        # Latency metrics per endpoint
        self.latency_metrics: Dict[str, LatencyMetrics] = defaultdict(LatencyMetrics)

        # Token usage metrics
        self.token_metrics = TokenMetrics()

        # Cache metrics
        self.cache_metrics = CacheMetrics()

        # Model distribution
        self.model_metrics = ModelMetrics()

        # Start time for uptime
        self.start_time = datetime.utcnow()

        logger.info("Metrics collector initialized")

    def record_request_latency(self, endpoint: str, latency_ms: float):
        """
        Record request latency for an endpoint.

        Args:
            endpoint: Endpoint path (e.g., "/qa", "/translate")
            latency_ms: Request latency in milliseconds
        """
        with self._lock:
            self.latency_metrics[endpoint].add_sample(latency_ms)

    def record_token_usage(self, prompt_tokens: int, completion_tokens: int):
        """
        Record token usage for a request.

        Args:
            prompt_tokens: Number of prompt tokens
            completion_tokens: Number of completion tokens
        """
        with self._lock:
            self.token_metrics.add_usage(prompt_tokens, completion_tokens)

    def record_cache_hit(self):
        """Record a cache hit."""
        with self._lock:
            self.cache_metrics.record_hit()

    def record_cache_miss(self):
        """Record a cache miss."""
        with self._lock:
            self.cache_metrics.record_miss()

    def record_model_use(self, model_type: str):
        """
        Record model usage.

        Args:
            model_type: Model identifier (e.g., "gpt-5-nano", "gpt-5-mini", "gpt-5")
        """
        with self._lock:
            self.model_metrics.record_model_use(model_type)

    def get_uptime_seconds(self) -> float:
        """Get uptime in seconds."""
        return (datetime.utcnow() - self.start_time).total_seconds()

    def get_summary(self) -> Dict[str, Any]:
        """
        Get metrics summary as JSON-serializable dict.

        Returns:
            Dict with all metrics
        """
        with self._lock:
            # Latency metrics per endpoint
            latency_summary = {}
            for endpoint, metrics in self.latency_metrics.items():
                latency_summary[endpoint] = {
                    "total_requests": metrics.total_requests,
                    "p50_ms": round(metrics.get_percentile(50), 2),
                    "p95_ms": round(metrics.get_percentile(95), 2),
                    "p99_ms": round(metrics.get_percentile(99), 2),
                    "avg_ms": round(metrics.get_avg(), 2),
                }

            # Token metrics
            token_summary = {
                "total_tokens": self.token_metrics.total_tokens,
                "total_prompt_tokens": self.token_metrics.total_prompt_tokens,
                "total_completion_tokens": self.token_metrics.total_completion_tokens,
                "total_requests": self.token_metrics.total_requests,
                "avg_tokens_per_request": round(self.token_metrics.avg_tokens_per_request, 2),
            }

            # Cache metrics
            cache_summary = {
                "hits": self.cache_metrics.hits,
                "misses": self.cache_metrics.misses,
                "total": self.cache_metrics.total,
                "hit_rate_percent": round(self.cache_metrics.hit_rate, 2),
            }

            # Model distribution
            model_summary = {
                "total_requests": self.model_metrics.total_requests,
                "distribution_percent": {
                    model: round(pct, 2)
                    for model, pct in self.model_metrics.get_distribution().items()
                },
                "counts": dict(self.model_metrics.model_counts),
            }

            return {
                "uptime_seconds": round(self.get_uptime_seconds(), 2),
                "start_time": self.start_time.isoformat(),
                "latency": latency_summary,
                "tokens": token_summary,
                "cache": cache_summary,
                "models": model_summary,
            }

    def export_prometheus(self) -> str:
        """
        Export metrics in Prometheus format.

        Returns:
            Prometheus-formatted metrics string
        """
        with self._lock:
            lines = []

            # Uptime
            lines.append("# HELP grant_analyst_uptime_seconds Time since server started")
            lines.append("# TYPE grant_analyst_uptime_seconds gauge")
            lines.append(f"grant_analyst_uptime_seconds {self.get_uptime_seconds()}")
            lines.append("")

            # Latency metrics
            lines.append("# HELP grant_analyst_request_latency_ms Request latency in milliseconds")
            lines.append("# TYPE grant_analyst_request_latency_ms summary")
            for endpoint, metrics in self.latency_metrics.items():
                safe_endpoint = endpoint.replace("/", "_").strip("_")
                lines.append(f'grant_analyst_request_latency_ms{{endpoint="{endpoint}",quantile="0.5"}} {metrics.get_percentile(50)}')
                lines.append(f'grant_analyst_request_latency_ms{{endpoint="{endpoint}",quantile="0.95"}} {metrics.get_percentile(95)}')
                lines.append(f'grant_analyst_request_latency_ms{{endpoint="{endpoint}",quantile="0.99"}} {metrics.get_percentile(99)}')
                lines.append(f'grant_analyst_request_latency_ms_count{{endpoint="{endpoint}"}} {metrics.total_requests}')
            lines.append("")

            # Token metrics
            lines.append("# HELP grant_analyst_tokens_total Total tokens used")
            lines.append("# TYPE grant_analyst_tokens_total counter")
            lines.append(f"grant_analyst_tokens_total {self.token_metrics.total_tokens}")
            lines.append("")

            lines.append("# HELP grant_analyst_prompt_tokens_total Total prompt tokens")
            lines.append("# TYPE grant_analyst_prompt_tokens_total counter")
            lines.append(f"grant_analyst_prompt_tokens_total {self.token_metrics.total_prompt_tokens}")
            lines.append("")

            lines.append("# HELP grant_analyst_completion_tokens_total Total completion tokens")
            lines.append("# TYPE grant_analyst_completion_tokens_total counter")
            lines.append(f"grant_analyst_completion_tokens_total {self.token_metrics.total_completion_tokens}")
            lines.append("")

            # Cache metrics
            lines.append("# HELP grant_analyst_cache_hits_total Total cache hits")
            lines.append("# TYPE grant_analyst_cache_hits_total counter")
            lines.append(f"grant_analyst_cache_hits_total {self.cache_metrics.hits}")
            lines.append("")

            lines.append("# HELP grant_analyst_cache_misses_total Total cache misses")
            lines.append("# TYPE grant_analyst_cache_misses_total counter")
            lines.append(f"grant_analyst_cache_misses_total {self.cache_metrics.misses}")
            lines.append("")

            lines.append("# HELP grant_analyst_cache_hit_rate Cache hit rate (0-100)")
            lines.append("# TYPE grant_analyst_cache_hit_rate gauge")
            lines.append(f"grant_analyst_cache_hit_rate {self.cache_metrics.hit_rate}")
            lines.append("")

            # Model distribution
            lines.append("# HELP grant_analyst_model_requests_total Total requests per model")
            lines.append("# TYPE grant_analyst_model_requests_total counter")
            for model, count in self.model_metrics.model_counts.items():
                lines.append(f'grant_analyst_model_requests_total{{model="{model}"}} {count}')
            lines.append("")

            return "\n".join(lines)

    def log_metrics(self):
        """Log metrics summary as JSON."""
        summary = self.get_summary()
        logger.info(f"📊 Metrics Summary: {json.dumps(summary, indent=2)}")


# Global metrics collector instance
_metrics_collector: Optional[MetricsCollector] = None


def get_metrics_collector() -> MetricsCollector:
    """Get or create the global metrics collector."""
    global _metrics_collector
    if _metrics_collector is None:
        _metrics_collector = MetricsCollector()
    return _metrics_collector


class RequestTimer:
    """Context manager for timing requests."""

    def __init__(self, endpoint: str):
        """
        Initialize request timer.

        Args:
            endpoint: Endpoint path to track
        """
        self.endpoint = endpoint
        self.start_time = None
        self.collector = get_metrics_collector()

    def __enter__(self):
        """Start timing."""
        self.start_time = time.perf_counter()
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        """Stop timing and record latency."""
        if self.start_time is not None:
            latency_ms = (time.perf_counter() - self.start_time) * 1000
            self.collector.record_request_latency(self.endpoint, latency_ms)


# Convenience functions for recording metrics
def record_latency(endpoint: str, latency_ms: float):
    """Record request latency."""
    get_metrics_collector().record_request_latency(endpoint, latency_ms)


def record_tokens(prompt_tokens: int, completion_tokens: int):
    """Record token usage."""
    get_metrics_collector().record_token_usage(prompt_tokens, completion_tokens)


def record_cache_hit():
    """Record cache hit."""
    get_metrics_collector().record_cache_hit()


def record_cache_miss():
    """Record cache miss."""
    get_metrics_collector().record_cache_miss()


def record_model_use(model_type: str):
    """Record model usage."""
    get_metrics_collector().record_model_use(model_type)


def get_metrics_summary() -> Dict[str, Any]:
    """Get metrics summary."""
    return get_metrics_collector().get_summary()


def export_prometheus() -> str:
    """Export metrics in Prometheus format."""
    return get_metrics_collector().export_prometheus()