"""Metrics collection and export module.

This module provides metrics collection infrastructure with support for
counters, gauges, and histograms. Metrics can be exported in Prometheus format.
"""

from __future__ import annotations

import threading
import time
from collections.abc import Callable
from dataclasses import dataclass, field
from typing import Any


@dataclass(frozen=True)
class MetricLabel:
    """A label key-value pair for metrics.

    Attributes:
        key: Label name
        value: Label value
    """

    key: str
    value: str

    def __str__(self) -> str:
        """Return Prometheus-format label string."""
        return f'{self.key}="{self.value}"'


@dataclass
class MetricValue:
    """A metric value with optional labels.

    Attributes:
        value: The numeric value
        labels: Tuple of label key-value pairs
        timestamp: Optional timestamp for the value
    """

    value: float
    labels: tuple[MetricLabel, ...] = field(default_factory=tuple)
    timestamp: float = field(default_factory=time.time)

    def get_label_string(self) -> str:
        """Get Prometheus-format label string."""
        if not self.labels:
            return ""
        return "{" + ",".join(str(label) for label in self.labels) + "}"


class Counter:
    """A counter metric that only increases.

    Counters are typically used to count events like requests served,
    tasks completed, or errors occurred.

    Example:
        counter = Counter("requests_total", "Total HTTP requests")
        counter.inc()  # Increment by 1
        counter.inc(5)  # Increment by 5
    """

    def __init__(self, name: str, description: str = "") -> None:
        """Initialize a counter metric.

        Args:
            name: Metric name (should follow Prometheus naming conventions)
            description: Human-readable description
        """
        self._name = name
        self._description = description
        self._values: dict[tuple[MetricLabel, ...], float] = {}
        self._lock = threading.Lock()

    def inc(self, amount: float = 1, labels: dict[str, str] | None = None) -> None:
        """Increment the counter.

        Args:
            amount: Amount to increment by (must be non-negative)
            labels: Optional label dictionary

        Raises:
            ValueError: If amount is negative
        """
        if amount < 0:
            raise ValueError("Counter cannot be decremented")

        label_tuple = self._labels_to_tuple(labels)
        with self._lock:
            self._values[label_tuple] = self._values.get(label_tuple, 0) + amount

    def get(self, labels: dict[str, str] | None = None) -> float:
        """Get current counter value.

        Args:
            labels: Optional label dictionary to filter by

        Returns:
            Current counter value
        """
        label_tuple = self._labels_to_tuple(labels)
        with self._lock:
            return self._values.get(label_tuple, 0)

    def _labels_to_tuple(
        self, labels: dict[str, str] | None
    ) -> tuple[MetricLabel, ...]:
        """Convert label dict to sorted tuple for consistent hashing."""
        if not labels:
            return ()
        return tuple(
            MetricLabel(k, v) for k, v in sorted(labels.items())
        )

    def collect(self) -> list[MetricValue]:
        """Collect all metric values.

        Returns:
            List of MetricValue objects with their labels
        """
        with self._lock:
            return [
                MetricValue(value, labels)
                for labels, value in self._values.items()
            ]

    def to_prometheus(self) -> str:
        """Export metric in Prometheus text format.

        Returns:
            Prometheus-formatted metric string
        """
        lines = []
        if self._description:
            lines.append(f"# HELP {self._name} {self._description}")
        lines.append(f"# TYPE {self._name} counter")

        for value in self.collect():
            label_str = value.get_label_string()
            lines.append(f"{self._name}{label_str} {value.value}")

        return "\n".join(lines)

    @property
    def name(self) -> str:
        """Return metric name."""
        return self._name


class Gauge:
    """A gauge metric that can go up and down.

    Gauges are typically used for measured values like temperatures,
    current memory usage, or the number of items in a queue.

    Example:
        gauge = Gauge("memory_usage_bytes", "Current memory usage")
        gauge.set(1024 * 1024 * 100)  # Set to 100MB
        gauge.inc(1024)  # Increment by 1KB
        gauge.dec(512)   # Decrement by 512 bytes
    """

    def __init__(self, name: str, description: str = "") -> None:
        """Initialize a gauge metric.

        Args:
            name: Metric name (should follow Prometheus naming conventions)
            description: Human-readable description
        """
        self._name = name
        self._description = description
        self._values: dict[tuple[MetricLabel, ...], float] = {}
        self._lock = threading.Lock()

    def set(self, value: float, labels: dict[str, str] | None = None) -> None:
        """Set the gauge to a specific value.

        Args:
            value: Value to set
            labels: Optional label dictionary
        """
        label_tuple = self._labels_to_tuple(labels)
        with self._lock:
            self._values[label_tuple] = value

    def inc(self, amount: float = 1, labels: dict[str, str] | None = None) -> None:
        """Increment the gauge.

        Args:
            amount: Amount to increment by
            labels: Optional label dictionary
        """
        label_tuple = self._labels_to_tuple(labels)
        with self._lock:
            self._values[label_tuple] = self._values.get(label_tuple, 0) + amount

    def dec(self, amount: float = 1, labels: dict[str, str] | None = None) -> None:
        """Decrement the gauge.

        Args:
            amount: Amount to decrement by
            labels: Optional label dictionary
        """
        self.inc(-amount, labels)

    def get(self, labels: dict[str, str] | None = None) -> float:
        """Get current gauge value.

        Args:
            labels: Optional label dictionary to filter by

        Returns:
            Current gauge value
        """
        label_tuple = self._labels_to_tuple(labels)
        with self._lock:
            return self._values.get(label_tuple, 0)

    def _labels_to_tuple(
        self, labels: dict[str, str] | None
    ) -> tuple[MetricLabel, ...]:
        """Convert label dict to sorted tuple for consistent hashing."""
        if not labels:
            return ()
        return tuple(
            MetricLabel(k, v) for k, v in sorted(labels.items())
        )

    def collect(self) -> list[MetricValue]:
        """Collect all metric values.

        Returns:
            List of MetricValue objects with their labels
        """
        with self._lock:
            return [
                MetricValue(value, labels)
                for labels, value in self._values.items()
            ]

    def to_prometheus(self) -> str:
        """Export metric in Prometheus text format.

        Returns:
            Prometheus-formatted metric string
        """
        lines = []
        if self._description:
            lines.append(f"# HELP {self._name} {self._description}")
        lines.append(f"# TYPE {self._name} gauge")

        for value in self.collect():
            label_str = value.get_label_string()
            lines.append(f"{self._name}{label_str} {value.value}")

        return "\n".join(lines)

    @property
    def name(self) -> str:
        """Return metric name."""
        return self._name


class Histogram:
    """A histogram metric for sampling observations.

    Histograms track the size and number of events in buckets.
    They are used for tracking request latencies, response sizes, etc.

    Example:
        histogram = Histogram(
            "request_duration_seconds",
            "HTTP request duration",
            buckets=[0.1, 0.5, 1.0, 2.0, 5.0]
        )
        histogram.observe(0.3)  # Observe a 300ms request
    """

    DEFAULT_BUCKETS = [0.005, 0.01, 0.025, 0.05, 0.1, 0.25, 0.5, 1.0, 2.5, 5.0, 10.0]

    def __init__(
        self,
        name: str,
        description: str = "",
        buckets: list[float] | None = None,
    ) -> None:
        """Initialize a histogram metric.

        Args:
            name: Metric name (should follow Prometheus naming conventions)
            description: Human-readable description
            buckets: Optional custom bucket boundaries (default: DEFAULT_BUCKETS)
        """
        self._name = name
        self._description = description
        self._buckets = sorted(buckets if buckets is not None else self.DEFAULT_BUCKETS)
        self._counts: dict[tuple[MetricLabel, ...], list[int]] = {}
        self._sums: dict[tuple[MetricLabel, ...], float] = {}
        self._lock = threading.Lock()

    def observe(self, value: float, labels: dict[str, str] | None = None) -> None:
        """Observe a value.

        Args:
            value: Value to observe
            labels: Optional label dictionary
        """
        label_tuple = self._labels_to_tuple(labels)

        with self._lock:
            if label_tuple not in self._counts:
                self._counts[label_tuple] = [0] * len(self._buckets)
                self._sums[label_tuple] = 0.0

            # Increment appropriate buckets
            for i, bucket in enumerate(self._buckets):
                if value <= bucket:
                    self._counts[label_tuple][i] += 1

            self._sums[label_tuple] += value

    def get_bucket_counts(
        self, labels: dict[str, str] | None = None
    ) -> list[tuple[float, int]]:
        """Get bucket counts.

        Args:
            labels: Optional label dictionary

        Returns:
            List of (bucket_boundary, count) tuples
        """
        label_tuple = self._labels_to_tuple(labels)
        with self._lock:
            counts = self._counts.get(label_tuple, [0] * len(self._buckets))
            return list(zip(self._buckets, counts))

    def get_sum(self, labels: dict[str, str] | None = None) -> float:
        """Get sum of all observed values.

        Args:
            labels: Optional label dictionary

        Returns:
            Sum of observed values
        """
        label_tuple = self._labels_to_tuple(labels)
        with self._lock:
            return self._sums.get(label_tuple, 0.0)

    def get_count(self, labels: dict[str, str] | None = None) -> int:
        """Get total count of observations.

        Args:
            labels: Optional label dictionary

        Returns:
            Total observation count
        """
        buckets = self.get_bucket_counts(labels)
        return buckets[-1][1] if buckets else 0

    def _labels_to_tuple(
        self, labels: dict[str, str] | None
    ) -> tuple[MetricLabel, ...]:
        """Convert label dict to sorted tuple for consistent hashing."""
        if not labels:
            return ()
        return tuple(
            MetricLabel(k, v) for k, v in sorted(labels.items())
        )

    def to_prometheus(self) -> str:
        """Export metric in Prometheus text format.

        Returns:
            Prometheus-formatted metric string
        """
        lines = []
        if self._description:
            lines.append(f"# HELP {self._name} {self._description}")
        lines.append(f"# TYPE {self._name} histogram")

        with self._lock:
            for labels, counts in self._counts.items():
                label_str = ""
                if labels:
                    label_str = "{" + ",".join(str(label) for label in labels) + ",}"

                cumulative = 0
                for bucket, count in zip(self._buckets, counts):
                    cumulative += count
                    bucket_labels = label_str.rstrip(",}") if label_str else ""
                    if bucket_labels:
                        bucket_labels += f',le="{bucket}"' + "}"
                    else:
                        bucket_labels = f'{{le="{bucket}"}}'
                    lines.append(f"{self._name}_bucket{bucket_labels} {cumulative}")

                # Add +Inf bucket
                inf_labels = label_str.rstrip(",}") if label_str else ""
                if inf_labels:
                    inf_labels += ',le="+Inf"' + "}"
                else:
                    inf_labels = '{le="+Inf"}'
                lines.append(f"{self._name}_bucket{inf_labels} {cumulative}")

                # Sum and count
                sum_label_str = ""
                if labels:
                    sum_label_str = "{" + ",".join(str(label) for label in labels) + "}"
                lines.append(f"{self._name}_sum{sum_label_str} {self._sums[labels]}")
                lines.append(f"{self._name}_count{sum_label_str} {cumulative}")

        return "\n".join(lines)

    @property
    def name(self) -> str:
        """Return metric name."""
        return self._name


class MetricsCollector:
    """Collector for managing multiple metrics.

    The MetricsCollector provides a central registry for all metrics
    and can export them in Prometheus format.

    Example:
        collector = MetricsCollector()

        # Create metrics
        requests = collector.counter("requests_total", "Total requests")
        latency = collector.histogram("request_latency", "Request latency")

        # Export all metrics
        print(collector.to_prometheus())
    """

    def __init__(self) -> None:
        """Initialize the metrics collector."""
        self._metrics: dict[str, Counter | Gauge | Histogram] = {}
        self._lock = threading.Lock()

    def counter(self, name: str, description: str = "") -> Counter:
        """Create or get a counter metric.

        Args:
            name: Metric name
            description: Human-readable description

        Returns:
            Counter metric instance
        """
        with self._lock:
            if name not in self._metrics:
                self._metrics[name] = Counter(name, description)
            return self._metrics[name]  # type: ignore[return-value]

    def gauge(self, name: str, description: str = "") -> Gauge:
        """Create or get a gauge metric.

        Args:
            name: Metric name
            description: Human-readable description

        Returns:
            Gauge metric instance
        """
        with self._lock:
            if name not in self._metrics:
                self._metrics[name] = Gauge(name, description)
            return self._metrics[name]  # type: ignore[return-value]

    def histogram(
        self,
        name: str,
        description: str = "",
        buckets: list[float] | None = None,
    ) -> Histogram:
        """Create or get a histogram metric.

        Args:
            name: Metric name
            description: Human-readable description
            buckets: Optional custom bucket boundaries

        Returns:
            Histogram metric instance
        """
        with self._lock:
            if name not in self._metrics:
                self._metrics[name] = Histogram(name, description, buckets)
            return self._metrics[name]  # type: ignore[return-value]

    def get_metric(self, name: str) -> Counter | Gauge | Histogram | None:
        """Get a metric by name.

        Args:
            name: Metric name

        Returns:
            Metric instance or None if not found
        """
        with self._lock:
            return self._metrics.get(name)

    def remove_metric(self, name: str) -> bool:
        """Remove a metric by name.

        Args:
            name: Metric name

        Returns:
            True if metric was removed, False if not found
        """
        with self._lock:
            if name in self._metrics:
                del self._metrics[name]
                return True
            return False

    def clear(self) -> None:
        """Remove all metrics."""
        with self._lock:
            self._metrics.clear()

    def to_prometheus(self) -> str:
        """Export all metrics in Prometheus text format.

        Returns:
            Prometheus-formatted metrics string
        """
        with self._lock:
            if not self._metrics:
                return ""
            return "\n\n".join(
                metric.to_prometheus() for metric in self._metrics.values()
            )

    def get_all_names(self) -> list[str]:
        """Get all registered metric names.

        Returns:
            List of metric names
        """
        with self._lock:
            return list(self._metrics.keys())


def timing(metric: Histogram) -> Callable[..., Any]:
    """Decorator to time function execution.

    Args:
        metric: Histogram metric to record timing

    Returns:
        Decorator function

    Example:
        latency = collector.histogram("request_latency_seconds")

        @timing(latency)
        def handle_request():
            # ... handle request
            pass
    """
    def decorator(func: Callable[..., Any]) -> Callable[..., Any]:
        def wrapper(*args: Any, **kwargs: Any) -> Any:
            start = time.time()
            try:
                return func(*args, **kwargs)
            finally:
                metric.observe(time.time() - start)
        return wrapper
    return decorator
