"""
Metrics and observability hooks for Proxilion.

Provides hooks for metrics and observability without requiring external
dependencies. Allows integration with Prometheus, OpenTelemetry, StatsD, etc.

Quick Start:
    >>> from proxilion.observability import (
    ...     ObservabilityHooks,
    ...     InMemoryMetricHook,
    ...     emit_counter,
    ...     emit_timing,
    ... )
    >>>
    >>> # Set up metrics hook
    >>> hooks = ObservabilityHooks.get_instance()
    >>> memory_hook = InMemoryMetricHook()
    >>> hooks.add_metric_hook(memory_hook)
    >>>
    >>> # Emit metrics
    >>> emit_counter("proxilion.auth.requests", tags={"user": "alice"})
    >>> emit_timing("proxilion.auth.latency_ms", 45.2, tags={"tool": "search"})
    >>>
    >>> # Check recorded metrics
    >>> memory_hook.get_counter("proxilion.auth.requests", {"user": "alice"})
    1.0

Integration with Prometheus:
    >>> from prometheus_client import Counter, Histogram
    >>>
    >>> class PrometheusMetricHook:
    ...     def __init__(self):
    ...         self.auth_requests = Counter(
    ...             "proxilion_auth_requests_total",
    ...             "Total authorization requests",
    ...             ["user", "tool"]
    ...         )
    ...         self.auth_latency = Histogram(
    ...             "proxilion_auth_latency_seconds",
    ...             "Authorization latency in seconds"
    ...         )
    ...
    ...     def increment(self, name, value=1.0, tags=None):
    ...         if name == "proxilion.auth.requests":
    ...             self.auth_requests.labels(**(tags or {})).inc(value)
    ...
    ...     def histogram(self, name, value, tags=None):
    ...         if name == "proxilion.auth.latency_ms":
    ...             self.auth_latency.observe(value / 1000)  # Convert to seconds
    ...
    ...     def gauge(self, name, value, tags=None): pass
    ...     def timing(self, name, duration_ms, tags=None):
    ...         self.histogram(name, duration_ms, tags)
    >>>
    >>> hooks = ObservabilityHooks.get_instance()
    >>> hooks.add_metric_hook(PrometheusMetricHook())

Integration with OpenTelemetry:
    >>> from opentelemetry import metrics
    >>>
    >>> class OpenTelemetryMetricHook:
    ...     def __init__(self):
    ...         meter = metrics.get_meter("proxilion")
    ...         self._counters = {}
    ...         self._histograms = {}
    ...         self._meter = meter
    ...
    ...     def increment(self, name, value=1.0, tags=None):
    ...         if name not in self._counters:
    ...             self._counters[name] = self._meter.create_counter(name)
    ...         self._counters[name].add(value, tags or {})
    ...
    ...     def histogram(self, name, value, tags=None):
    ...         if name not in self._histograms:
    ...             self._histograms[name] = self._meter.create_histogram(name)
    ...         self._histograms[name].record(value, tags or {})
    ...
    ...     def gauge(self, name, value, tags=None): pass
    ...     def timing(self, name, duration_ms, tags=None):
    ...         self.histogram(name, duration_ms, tags)
"""

from __future__ import annotations

import logging
from collections import defaultdict
from dataclasses import dataclass
from enum import Enum
from typing import Any, Protocol, runtime_checkable


class MetricType(Enum):
    """Types of metrics that can be emitted."""

    COUNTER = "counter"
    """A monotonically increasing counter."""

    GAUGE = "gauge"
    """A value that can go up or down."""

    HISTOGRAM = "histogram"
    """Distribution of values."""

    SUMMARY = "summary"
    """Similar to histogram but calculates quantiles."""

    TIMING = "timing"
    """Duration measurement (usually milliseconds)."""


@runtime_checkable
class MetricHook(Protocol):
    """
    Protocol for metric backends (Prometheus, StatsD, OpenTelemetry, etc.).

    Implement this protocol to integrate with your metrics infrastructure.
    The hooks will be called by Proxilion whenever metrics are emitted.

    Example:
        >>> class MyMetricHook:
        ...     def increment(self, name, value=1.0, tags=None):
        ...         # Send to your metrics backend
        ...         pass
        ...
        ...     def gauge(self, name, value, tags=None):
        ...         pass
        ...
        ...     def histogram(self, name, value, tags=None):
        ...         pass
        ...
        ...     def timing(self, name, duration_ms, tags=None):
        ...         pass
        >>>
        >>> from proxilion.observability import ObservabilityHooks
        >>> hooks = ObservabilityHooks.get_instance()
        >>> hooks.add_metric_hook(MyMetricHook())
    """

    def increment(self, name: str, value: float = 1.0, tags: dict[str, Any] | None = None) -> None:
        """
        Increment a counter metric.

        Args:
            name: The metric name (e.g., "proxilion.auth.requests").
            value: The amount to increment by (default 1.0).
            tags: Optional tags/labels for the metric.
        """
        ...

    def gauge(self, name: str, value: float, tags: dict[str, Any] | None = None) -> None:
        """
        Set a gauge metric.

        Args:
            name: The metric name.
            value: The gauge value.
            tags: Optional tags/labels for the metric.
        """
        ...

    def histogram(self, name: str, value: float, tags: dict[str, Any] | None = None) -> None:
        """
        Record a value in a histogram.

        Args:
            name: The metric name.
            value: The value to record.
            tags: Optional tags/labels for the metric.
        """
        ...

    def timing(self, name: str, duration_ms: float, tags: dict[str, Any] | None = None) -> None:
        """
        Record a timing/duration measurement.

        Args:
            name: The metric name.
            duration_ms: The duration in milliseconds.
            tags: Optional tags/labels for the metric.
        """
        ...


class LoggingMetricHook:
    """
    Simple hook that logs metrics (for development and debugging).

    Example:
        >>> import logging
        >>> logging.basicConfig(level=logging.DEBUG)
        >>> hook = LoggingMetricHook()
        >>> hook.increment("requests", 1.0, {"service": "api"})
        DEBUG:proxilion.metrics:COUNTER requests=1.0 tags={'service': 'api'}
    """

    def __init__(self, logger: logging.Logger | None = None, level: int = logging.DEBUG):
        """
        Initialize the logging hook.

        Args:
            logger: Logger instance to use. Defaults to "proxilion.metrics".
            level: Logging level for metric messages.
        """
        self.logger = logger or logging.getLogger("proxilion.metrics")
        self.level = level

    def increment(self, name: str, value: float = 1.0, tags: dict[str, Any] | None = None) -> None:
        """Log a counter increment."""
        self.logger.log(self.level, f"COUNTER {name}={value} tags={tags}")

    def gauge(self, name: str, value: float, tags: dict[str, Any] | None = None) -> None:
        """Log a gauge value."""
        self.logger.log(self.level, f"GAUGE {name}={value} tags={tags}")

    def histogram(self, name: str, value: float, tags: dict[str, Any] | None = None) -> None:
        """Log a histogram value."""
        self.logger.log(self.level, f"HISTOGRAM {name}={value} tags={tags}")

    def timing(self, name: str, duration_ms: float, tags: dict[str, Any] | None = None) -> None:
        """Log a timing value."""
        self.logger.log(self.level, f"TIMING {name}={duration_ms}ms tags={tags}")


@dataclass
class HistogramStats:
    """Statistics for a histogram metric."""

    count: int
    """Number of recorded values."""

    total: float
    """Sum of all values."""

    min: float
    """Minimum value."""

    max: float
    """Maximum value."""

    avg: float
    """Average value."""

    def to_dict(self) -> dict[str, Any]:
        """Convert to dictionary."""
        return {
            "count": self.count,
            "total": self.total,
            "min": self.min,
            "max": self.max,
            "avg": self.avg,
        }


class InMemoryMetricHook:
    """
    In-memory metrics for testing and simple use cases.

    Stores all metrics in memory and provides methods to query them.
    Useful for testing that metrics are being emitted correctly.

    Example:
        >>> hook = InMemoryMetricHook()
        >>> hook.increment("requests", tags={"user": "alice"})
        >>> hook.increment("requests", tags={"user": "alice"})
        >>> hook.increment("requests", tags={"user": "bob"})
        >>>
        >>> hook.get_counter("requests", {"user": "alice"})
        2.0
        >>> hook.get_counter("requests", {"user": "bob"})
        1.0
        >>>
        >>> hook.histogram("latency", 100.0)
        >>> hook.histogram("latency", 200.0)
        >>> stats = hook.get_histogram_stats("latency")
        >>> stats.avg
        150.0
    """

    def __init__(self) -> None:
        """Initialize the in-memory hook."""
        self.counters: dict[str, float] = defaultdict(float)
        self.gauges: dict[str, float] = {}
        self.histograms: dict[str, list[float]] = defaultdict(list)
        self.timings: dict[str, list[float]] = defaultdict(list)

    def _make_key(self, name: str, tags: dict[str, Any] | None) -> str:
        """Create a unique key from metric name and tags."""
        if not tags:
            return name
        sorted_tags = sorted(tags.items())
        tag_str = ",".join(f"{k}={v}" for k, v in sorted_tags)
        return f"{name}[{tag_str}]"

    def increment(self, name: str, value: float = 1.0, tags: dict[str, Any] | None = None) -> None:
        """Increment a counter."""
        key = self._make_key(name, tags)
        self.counters[key] += value

    def gauge(self, name: str, value: float, tags: dict[str, Any] | None = None) -> None:
        """Set a gauge value."""
        key = self._make_key(name, tags)
        self.gauges[key] = value

    def histogram(self, name: str, value: float, tags: dict[str, Any] | None = None) -> None:
        """Record a histogram value."""
        key = self._make_key(name, tags)
        self.histograms[key].append(value)

    def timing(self, name: str, duration_ms: float, tags: dict[str, Any] | None = None) -> None:
        """Record a timing value."""
        key = self._make_key(name, tags)
        self.timings[key].append(duration_ms)

    def get_counter(self, name: str, tags: dict[str, Any] | None = None) -> float:
        """
        Get the current value of a counter.

        Args:
            name: The metric name.
            tags: Optional tags/labels for the metric.

        Returns:
            The current counter value, or 0.0 if not found.
        """
        key = self._make_key(name, tags)
        return self.counters[key]

    def get_gauge(self, name: str, tags: dict[str, Any] | None = None) -> float | None:
        """
        Get the current value of a gauge.

        Args:
            name: The metric name.
            tags: Optional tags/labels for the metric.

        Returns:
            The current gauge value, or None if not set.
        """
        key = self._make_key(name, tags)
        return self.gauges.get(key)

    def get_histogram_stats(
        self, name: str, tags: dict[str, Any] | None = None
    ) -> HistogramStats | None:
        """
        Get statistics for a histogram.

        Args:
            name: The metric name.
            tags: Optional tags/labels for the metric.

        Returns:
            HistogramStats with count, sum, min, max, avg, or None if empty.
        """
        key = self._make_key(name, tags)
        values = self.histograms.get(key, [])
        if not values:
            return None
        return HistogramStats(
            count=len(values),
            total=sum(values),
            min=min(values),
            max=max(values),
            avg=sum(values) / len(values),
        )

    def get_timing_stats(
        self, name: str, tags: dict[str, Any] | None = None
    ) -> HistogramStats | None:
        """
        Get statistics for timing measurements.

        Args:
            name: The metric name.
            tags: Optional tags/labels for the metric.

        Returns:
            HistogramStats with count, sum, min, max, avg, or None if empty.
        """
        key = self._make_key(name, tags)
        values = self.timings.get(key, [])
        if not values:
            return None
        return HistogramStats(
            count=len(values),
            total=sum(values),
            min=min(values),
            max=max(values),
            avg=sum(values) / len(values),
        )

    def get_all_counters(self) -> dict[str, float]:
        """Get all counter values."""
        return dict(self.counters)

    def get_all_gauges(self) -> dict[str, float]:
        """Get all gauge values."""
        return dict(self.gauges)

    def reset(self) -> None:
        """Reset all metrics to initial state."""
        self.counters.clear()
        self.gauges.clear()
        self.histograms.clear()
        self.timings.clear()


class ObservabilityHooks:
    """
    Central registry for observability hooks.

    This is a singleton that manages all metric hooks. Use `get_instance()`
    to get the global instance.

    Example:
        >>> from proxilion.observability import ObservabilityHooks, InMemoryMetricHook
        >>>
        >>> hooks = ObservabilityHooks.get_instance()
        >>> memory_hook = InMemoryMetricHook()
        >>> hooks.add_metric_hook(memory_hook)
        >>>
        >>> # Emit metrics (these will be sent to all registered hooks)
        >>> hooks.emit_counter("requests", tags={"service": "api"})
        >>> hooks.emit_timing("latency", 45.2)
        >>>
        >>> # Check the recorded metrics
        >>> memory_hook.get_counter("requests", {"service": "api"})
        1.0
    """

    _instance: ObservabilityHooks | None = None

    def __init__(self) -> None:
        """Initialize the observability hooks registry."""
        self._metric_hooks: list[MetricHook] = []

    @classmethod
    def get_instance(cls) -> ObservabilityHooks:
        """
        Get the global ObservabilityHooks instance.

        Returns:
            The singleton instance.
        """
        if cls._instance is None:
            cls._instance = cls()
        return cls._instance

    @classmethod
    def reset_instance(cls) -> None:
        """
        Reset the global instance.

        Useful for testing to ensure a clean state.
        """
        cls._instance = None

    def add_metric_hook(self, hook: MetricHook) -> None:
        """
        Register a metric hook.

        Args:
            hook: The metric hook to register.
        """
        self._metric_hooks.append(hook)

    def remove_metric_hook(self, hook: MetricHook) -> bool:
        """
        Remove a metric hook.

        Args:
            hook: The metric hook to remove.

        Returns:
            True if the hook was removed, False if not found.
        """
        try:
            self._metric_hooks.remove(hook)
            return True
        except ValueError:
            return False

    def clear_metric_hooks(self) -> None:
        """Remove all registered metric hooks."""
        self._metric_hooks.clear()

    @property
    def metric_hooks(self) -> list[MetricHook]:
        """Get a copy of the registered metric hooks."""
        return list(self._metric_hooks)

    def emit_metric(
        self,
        metric_type: MetricType,
        name: str,
        value: float,
        tags: dict[str, Any] | None = None,
    ) -> None:
        """
        Emit a metric to all registered hooks.

        Args:
            metric_type: The type of metric.
            name: The metric name.
            value: The metric value.
            tags: Optional tags/labels for the metric.
        """
        for hook in self._metric_hooks:
            if metric_type == MetricType.COUNTER:
                hook.increment(name, value, tags)
            elif metric_type == MetricType.GAUGE:
                hook.gauge(name, value, tags)
            elif metric_type == MetricType.HISTOGRAM:
                hook.histogram(name, value, tags)
            elif metric_type in (MetricType.TIMING, MetricType.SUMMARY):
                hook.timing(name, value, tags)

    def emit_counter(
        self, name: str, value: float = 1.0, tags: dict[str, Any] | None = None
    ) -> None:
        """
        Emit a counter metric.

        Args:
            name: The metric name.
            value: The amount to increment by (default 1.0).
            tags: Optional tags/labels for the metric.
        """
        self.emit_metric(MetricType.COUNTER, name, value, tags)

    def emit_gauge(self, name: str, value: float, tags: dict[str, Any] | None = None) -> None:
        """
        Emit a gauge metric.

        Args:
            name: The metric name.
            value: The gauge value.
            tags: Optional tags/labels for the metric.
        """
        self.emit_metric(MetricType.GAUGE, name, value, tags)

    def emit_histogram(self, name: str, value: float, tags: dict[str, Any] | None = None) -> None:
        """
        Emit a histogram metric.

        Args:
            name: The metric name.
            value: The value to record.
            tags: Optional tags/labels for the metric.
        """
        self.emit_metric(MetricType.HISTOGRAM, name, value, tags)

    def emit_timing(
        self, name: str, duration_ms: float, tags: dict[str, Any] | None = None
    ) -> None:
        """
        Emit a timing metric.

        Args:
            name: The metric name.
            duration_ms: The duration in milliseconds.
            tags: Optional tags/labels for the metric.
        """
        self.emit_metric(MetricType.TIMING, name, duration_ms, tags)


# Standard metric names emitted by Proxilion
# These are the metric names that Proxilion components emit

# Authorization metrics
METRIC_AUTH_REQUESTS = "proxilion.auth.requests"
"""Counter: Total authorization requests."""

METRIC_AUTH_ALLOWED = "proxilion.auth.allowed"
"""Counter: Authorization requests that were allowed."""

METRIC_AUTH_DENIED = "proxilion.auth.denied"
"""Counter: Authorization requests that were denied."""

METRIC_AUTH_LATENCY = "proxilion.auth.latency_ms"
"""Histogram: Authorization check latency in milliseconds."""

# Rate limiting metrics
METRIC_RATE_LIMIT_REQUESTS = "proxilion.rate_limit.requests"
"""Counter: Total rate limit checks."""

METRIC_RATE_LIMIT_EXCEEDED = "proxilion.rate_limit.exceeded"
"""Counter: Rate limit exceeded events."""

# Tool execution metrics
METRIC_TOOL_CALLS = "proxilion.tool.calls"
"""Counter: Total tool call executions."""

METRIC_TOOL_LATENCY = "proxilion.tool.latency_ms"
"""Histogram: Tool execution latency in milliseconds."""

METRIC_TOOL_ERRORS = "proxilion.tool.errors"
"""Counter: Tool execution errors."""

# Cost metrics
METRIC_COST_USD = "proxilion.cost.usd"
"""Counter: Total cost in USD."""

METRIC_TOKENS_INPUT = "proxilion.tokens.input"
"""Counter: Total input tokens processed."""

METRIC_TOKENS_OUTPUT = "proxilion.tokens.output"
"""Counter: Total output tokens generated."""

# Circuit breaker metrics
METRIC_CIRCUIT_BREAKER_OPEN = "proxilion.circuit_breaker.open"
"""Counter: Circuit breaker open events."""

METRIC_CIRCUIT_BREAKER_HALF_OPEN = "proxilion.circuit_breaker.half_open"
"""Counter: Circuit breaker half-open events."""

METRIC_CIRCUIT_BREAKER_CLOSED = "proxilion.circuit_breaker.closed"
"""Counter: Circuit breaker closed events."""


# Convenience functions for emitting metrics
def emit_counter(name: str, value: float = 1.0, tags: dict[str, Any] | None = None) -> None:
    """
    Emit a counter metric to all registered hooks.

    Args:
        name: The metric name (e.g., "proxilion.auth.requests").
        value: The amount to increment by (default 1.0).
        tags: Optional tags/labels for the metric.

    Example:
        >>> from proxilion.observability import emit_counter
        >>> emit_counter("proxilion.auth.requests", tags={"user": "alice"})
    """
    ObservabilityHooks.get_instance().emit_counter(name, value, tags)


def emit_gauge(name: str, value: float, tags: dict[str, Any] | None = None) -> None:
    """
    Emit a gauge metric to all registered hooks.

    Args:
        name: The metric name.
        value: The gauge value.
        tags: Optional tags/labels for the metric.

    Example:
        >>> from proxilion.observability import emit_gauge
        >>> emit_gauge("proxilion.connections.active", 42)
    """
    ObservabilityHooks.get_instance().emit_gauge(name, value, tags)


def emit_histogram(name: str, value: float, tags: dict[str, Any] | None = None) -> None:
    """
    Emit a histogram metric to all registered hooks.

    Args:
        name: The metric name.
        value: The value to record.
        tags: Optional tags/labels for the metric.

    Example:
        >>> from proxilion.observability import emit_histogram
        >>> emit_histogram("proxilion.response.size_bytes", 1024)
    """
    ObservabilityHooks.get_instance().emit_histogram(name, value, tags)


def emit_timing(name: str, duration_ms: float, tags: dict[str, Any] | None = None) -> None:
    """
    Emit a timing metric to all registered hooks.

    Args:
        name: The metric name.
        duration_ms: The duration in milliseconds.
        tags: Optional tags/labels for the metric.

    Example:
        >>> from proxilion.observability import emit_timing
        >>> emit_timing("proxilion.auth.latency_ms", 45.2, tags={"tool": "search"})
    """
    ObservabilityHooks.get_instance().emit_timing(name, duration_ms, tags)
