"""Stage metrics, dashboards, and alert definitions for the v3 pipeline.

Tracks latency, errors, batch size, queue depth, route metrics, field
accuracy, evidence coverage, calibration, fast-path rate, adjudication
reasons, GPU memory, GPU utilization, and GPU-seconds per document.
"""

from __future__ import annotations

import enum
from dataclasses import dataclass, field
from datetime import datetime, timezone
from typing import Any


class AlertSeverity(str, enum.Enum):
    """Alert severity levels."""

    INFO = "info"
    WARNING = "warning"
    CRITICAL = "critical"


class MetricType(str, enum.Enum):
    """Types of collected metrics."""

    COUNTER = "counter"
    GAUGE = "gauge"
    HISTOGRAM = "histogram"
    SUMMARY = "summary"


@dataclass(frozen=True)
class MetricAlert:
    """Alert definition for pipeline metrics."""

    name: str
    metric_name: str
    condition: str  # e.g., "> 0.05", "< 0.6"
    severity: AlertSeverity
    description: str
    threshold: float
    window_seconds: int = 300

    def evaluate(self, current_value: float) -> bool:
        """Check if the alert condition is triggered.

        Returns True if the alert should fire.
        """
        if self.condition.startswith(">"):
            return current_value > self.threshold
        elif self.condition.startswith("<"):
            return current_value < self.threshold
        elif self.condition.startswith(">="):
            return current_value >= self.threshold
        elif self.condition.startswith("<="):
            return current_value <= self.threshold
        return False


@dataclass
class StageMetrics:
    """Metrics for a single pipeline stage."""

    stage_name: str
    total_invocations: int = 0
    total_errors: int = 0
    total_latency_ms: float = 0.0
    max_latency_ms: float = 0.0
    min_latency_ms: float = float("inf")
    total_tokens_in: int = 0
    total_tokens_out: int = 0
    total_batch_items: int = 0
    total_batches: int = 0
    gpu_seconds: float = 0.0
    gpu_memory_peak_mb: float = 0.0
    gpu_utilization_avg: float = 0.0

    def record_invocation(
        self,
        latency_ms: float,
        tokens_in: int = 0,
        tokens_out: int = 0,
        error: bool = False,
        batch_size: int = 1,
        gpu_seconds: float = 0.0,
        gpu_memory_mb: float = 0.0,
        gpu_utilization: float = 0.0,
    ) -> None:
        """Record a single stage invocation."""
        self.total_invocations += 1
        self.total_latency_ms += latency_ms
        self.max_latency_ms = max(self.max_latency_ms, latency_ms)
        self.min_latency_ms = min(self.min_latency_ms, latency_ms)
        self.total_tokens_in += tokens_in
        self.total_tokens_out += tokens_out
        self.total_batch_items += batch_size
        self.total_batches += 1
        self.gpu_seconds += gpu_seconds
        self.gpu_memory_peak_mb = max(self.gpu_memory_peak_mb, gpu_memory_mb)

        if error:
            self.total_errors += 1

        # Running average for GPU utilization
        if gpu_utilization > 0:
            n = self.total_invocations
            self.gpu_utilization_avg = (
                self.gpu_utilization_avg * (n - 1) + gpu_utilization
            ) / n

    @property
    def avg_latency_ms(self) -> float:
        if self.total_invocations == 0:
            return 0.0
        return self.total_latency_ms / self.total_invocations

    @property
    def error_rate(self) -> float:
        if self.total_invocations == 0:
            return 0.0
        return self.total_errors / self.total_invocations

    @property
    def avg_batch_size(self) -> float:
        if self.total_batches == 0:
            return 0.0
        return self.total_batch_items / self.total_batches

    @property
    def avg_tokens_per_doc(self) -> float:
        if self.total_invocations == 0:
            return 0.0
        return (self.total_tokens_in + self.total_tokens_out) / self.total_invocations

    @property
    def gpu_seconds_per_doc(self) -> float:
        if self.total_invocations == 0:
            return 0.0
        return self.gpu_seconds / self.total_invocations

    def to_dict(self) -> dict[str, Any]:
        return {
            "stage_name": self.stage_name,
            "total_invocations": self.total_invocations,
            "total_errors": self.total_errors,
            "error_rate": self.error_rate,
            "avg_latency_ms": self.avg_latency_ms,
            "max_latency_ms": self.max_latency_ms,
            "avg_batch_size": self.avg_batch_size,
            "gpu_seconds_per_doc": self.gpu_seconds_per_doc,
            "gpu_memory_peak_mb": self.gpu_memory_peak_mb,
        }


# Default alert definitions for the v3 pipeline
DEFAULT_ALERTS: list[MetricAlert] = [
    MetricAlert(
        name="schema_failure_rate_high",
        metric_name="schema_failures",
        condition="> 0.05",
        severity=AlertSeverity.CRITICAL,
        description="Schema validation failure rate exceeds 5%",
        threshold=0.05,
    ),
    MetricAlert(
        name="unsupported_claims_high",
        metric_name="unsupported_claim_rate",
        condition="> 0.10",
        severity=AlertSeverity.WARNING,
        description="Unsupported claim rate exceeds 10%",
        threshold=0.10,
    ),
    MetricAlert(
        name="calibration_drift",
        metric_name="calibration_ece",
        condition="> 0.08",
        severity=AlertSeverity.WARNING,
        description="Calibration ECE exceeds 8%",
        threshold=0.08,
    ),
    MetricAlert(
        name="queue_saturation",
        metric_name="queue_saturation_ratio",
        condition="> 0.90",
        severity=AlertSeverity.CRITICAL,
        description="Queue saturation exceeds 90%",
        threshold=0.90,
    ),
    MetricAlert(
        name="provider_probe_failure",
        metric_name="probe_failure_rate",
        condition="> 0.0",
        severity=AlertSeverity.CRITICAL,
        description="Provider capability probe failed",
        threshold=0.0,
    ),
    MetricAlert(
        name="gpu_memory_high",
        metric_name="gpu_memory_utilization",
        condition="> 0.85",
        severity=AlertSeverity.WARNING,
        description="GPU memory utilization exceeds 85%",
        threshold=0.85,
    ),
]


@dataclass
class MetricsCollector:
    """Collects and aggregates metrics across pipeline stages.

    In production, this would export to Prometheus/Grafana.
    This implementation provides the collection logic for testing.
    """

    _stages: dict[str, StageMetrics] = field(default_factory=dict)
    _alerts: list[MetricAlert] = field(default_factory=list)
    _counters: dict[str, float] = field(default_factory=dict)
    _fired_alerts: list[tuple[MetricAlert, float, datetime]] = field(
        default_factory=list
    )

    def __post_init__(self) -> None:
        if not self._alerts:
            self._alerts = list(DEFAULT_ALERTS)

    def get_stage(self, stage_name: str) -> StageMetrics:
        """Get or create metrics for a stage."""
        if stage_name not in self._stages:
            self._stages[stage_name] = StageMetrics(stage_name=stage_name)
        return self._stages[stage_name]

    def record_stage(
        self,
        stage_name: str,
        latency_ms: float,
        tokens_in: int = 0,
        tokens_out: int = 0,
        error: bool = False,
        batch_size: int = 1,
        gpu_seconds: float = 0.0,
        gpu_memory_mb: float = 0.0,
        gpu_utilization: float = 0.0,
    ) -> None:
        """Record a stage invocation."""
        stage = self.get_stage(stage_name)
        stage.record_invocation(
            latency_ms=latency_ms,
            tokens_in=tokens_in,
            tokens_out=tokens_out,
            error=error,
            batch_size=batch_size,
            gpu_seconds=gpu_seconds,
            gpu_memory_mb=gpu_memory_mb,
            gpu_utilization=gpu_utilization,
        )

    def increment_counter(self, name: str, value: float = 1.0) -> None:
        """Increment a named counter."""
        self._counters[name] = self._counters.get(name, 0.0) + value

    def get_counter(self, name: str) -> float:
        """Get current counter value."""
        return self._counters.get(name, 0.0)

    def check_alerts(self) -> list[tuple[MetricAlert, float]]:
        """Evaluate all alert conditions. Returns (alert, value) for fired alerts."""
        fired: list[tuple[MetricAlert, float]] = []
        for alert in self._alerts:
            value = self._counters.get(alert.metric_name, 0.0)
            if alert.evaluate(value):
                fired.append((alert, value))
                self._fired_alerts.append(
                    (alert, value, datetime.now(timezone.utc))
                )
        return fired

    @property
    def stage_names(self) -> list[str]:
        return list(self._stages.keys())

    def summary(self) -> dict[str, Any]:
        """Generate a metrics summary for dashboard display."""
        return {
            "stages": {
                name: stage.to_dict() for name, stage in self._stages.items()
            },
            "counters": dict(self._counters),
            "fired_alerts": len(self._fired_alerts),
        }