"""Swarm metrics collector -- tracks task, provider, and worker performance.

Stdlib-only (ADR-021). Provides the data model that external exporters
(OTel, Prometheus) can consume.
"""

from __future__ import annotations

import statistics
import threading
import time
from collections import defaultdict
from dataclasses import dataclass, field
from enum import Enum


class TaskStatus(Enum):
    """Status of a task in the swarm pipeline."""

    PENDING = "pending"
    RUNNING = "running"
    COMPLETED = "completed"
    FAILED = "failed"


@dataclass
class TaskMetric:
    """Metrics for a single task execution."""

    task_id: str
    task_type: str
    provider: str
    worker: str = ""
    status: TaskStatus = TaskStatus.PENDING
    start_time: float = 0.0
    end_time: float = 0.0
    tokens_in: int = 0
    tokens_out: int = 0
    cost_usd: float = 0.0
    error: str = ""

    @property
    def duration_ms(self) -> float:
        """Duration in milliseconds, or 0 if timestamps are missing."""
        if self.end_time <= 0 or self.start_time <= 0:
            return 0.0
        return (self.end_time - self.start_time) * 1000

    def to_dict(self) -> dict:
        """Serialize to a plain dict for JSON export."""
        return {
            "task_id": self.task_id,
            "task_type": self.task_type,
            "provider": self.provider,
            "worker": self.worker,
            "status": self.status.value,
            "duration_ms": round(self.duration_ms, 2),
            "tokens_in": self.tokens_in,
            "tokens_out": self.tokens_out,
            "cost_usd": round(self.cost_usd, 6),
            "error": self.error,
        }


@dataclass
class ProviderStats:
    """Aggregated stats for a provider."""

    name: str
    total_tasks: int = 0
    completed: int = 0
    failed: int = 0
    total_tokens_in: int = 0
    total_tokens_out: int = 0
    total_cost_usd: float = 0.0
    latencies_ms: list[float] = field(default_factory=list)

    @property
    def success_rate(self) -> float:
        """Fraction of tasks that completed successfully."""
        if self.total_tasks == 0:
            return 0.0
        return self.completed / self.total_tasks

    @property
    def avg_latency_ms(self) -> float:
        """Mean latency across completed tasks."""
        if not self.latencies_ms:
            return 0.0
        return statistics.mean(self.latencies_ms)

    @property
    def p95_latency_ms(self) -> float:
        """95th-percentile latency."""
        if len(self.latencies_ms) < 2:
            return self.avg_latency_ms
        sorted_lat = sorted(self.latencies_ms)
        idx = int(len(sorted_lat) * 0.95)
        return sorted_lat[min(idx, len(sorted_lat) - 1)]

    @property
    def cost_per_1k_tokens(self) -> float:
        """Cost normalized to 1,000 tokens."""
        total_tokens = self.total_tokens_in + self.total_tokens_out
        if total_tokens == 0:
            return 0.0
        return (self.total_cost_usd / total_tokens) * 1000

    def to_dict(self) -> dict:
        """Serialize to a plain dict for JSON export."""
        return {
            "name": self.name,
            "total_tasks": self.total_tasks,
            "completed": self.completed,
            "failed": self.failed,
            "success_rate": round(self.success_rate, 4),
            "avg_latency_ms": round(self.avg_latency_ms, 2),
            "p95_latency_ms": round(self.p95_latency_ms, 2),
            "total_tokens_in": self.total_tokens_in,
            "total_tokens_out": self.total_tokens_out,
            "total_cost_usd": round(self.total_cost_usd, 6),
            "cost_per_1k_tokens": round(self.cost_per_1k_tokens, 6),
        }


@dataclass
class WorkerStats:
    """Aggregated stats for a worker node."""

    name: str
    total_tasks: int = 0
    completed: int = 0
    failed: int = 0
    task_types: dict[str, int] = field(default_factory=lambda: defaultdict(int))
    latencies_ms: list[float] = field(default_factory=list)

    @property
    def avg_duration_ms(self) -> float:
        """Mean task duration for this worker."""
        if not self.latencies_ms:
            return 0.0
        return statistics.mean(self.latencies_ms)

    @property
    def specialization(self) -> str:
        """Most common task type for this worker."""
        if not self.task_types:
            return "none"
        return max(self.task_types, key=self.task_types.get)  # type: ignore[arg-type]

    def to_dict(self) -> dict:
        """Serialize to a plain dict for JSON export."""
        return {
            "name": self.name,
            "total_tasks": self.total_tasks,
            "completed": self.completed,
            "failed": self.failed,
            "avg_duration_ms": round(self.avg_duration_ms, 2),
            "specialization": self.specialization,
            "task_type_distribution": dict(self.task_types),
        }


class MetricsCollector:
    """Thread-safe metrics collector for the swarm."""

    def __init__(self) -> None:
        self._tasks: list[TaskMetric] = []
        self._provider_stats: dict[str, ProviderStats] = {}
        self._worker_stats: dict[str, WorkerStats] = {}
        self._start_time = time.time()
        self._lock = threading.Lock()

    def record_task(self, metric: TaskMetric) -> None:
        """Record a completed task metric."""
        with self._lock:
            self._tasks.append(metric)
            self._update_provider_stats(metric)
            if metric.worker:
                self._update_worker_stats(metric)

    def _update_provider_stats(self, m: TaskMetric) -> None:
        if m.provider not in self._provider_stats:
            self._provider_stats[m.provider] = ProviderStats(name=m.provider)
        ps = self._provider_stats[m.provider]
        ps.total_tasks += 1
        ps.total_tokens_in += m.tokens_in
        ps.total_tokens_out += m.tokens_out
        ps.total_cost_usd += m.cost_usd
        if m.status == TaskStatus.COMPLETED:
            ps.completed += 1
            if m.duration_ms > 0:
                ps.latencies_ms.append(m.duration_ms)
        elif m.status == TaskStatus.FAILED:
            ps.failed += 1

    def _update_worker_stats(self, m: TaskMetric) -> None:
        if m.worker not in self._worker_stats:
            self._worker_stats[m.worker] = WorkerStats(name=m.worker)
        ws = self._worker_stats[m.worker]
        ws.total_tasks += 1
        ws.task_types[m.task_type] += 1
        if m.status == TaskStatus.COMPLETED:
            ws.completed += 1
            if m.duration_ms > 0:
                ws.latencies_ms.append(m.duration_ms)
        elif m.status == TaskStatus.FAILED:
            ws.failed += 1

    def tasks_per_minute(self) -> float:
        """Current throughput in tasks/minute."""
        elapsed = time.time() - self._start_time
        if elapsed <= 0:
            return 0.0
        with self._lock:
            return (len(self._tasks) / elapsed) * 60

    def total_cost(self) -> float:
        """Sum of cost_usd across all recorded tasks."""
        with self._lock:
            return sum(t.cost_usd for t in self._tasks)

    def error_rate(self) -> float:
        """Fraction of tasks that failed."""
        with self._lock:
            if not self._tasks:
                return 0.0
            failed = sum(1 for t in self._tasks if t.status == TaskStatus.FAILED)
            return failed / len(self._tasks)

    def provider_summary(self) -> dict[str, dict]:
        """Per-provider aggregated stats."""
        with self._lock:
            return {name: stats.to_dict() for name, stats in sorted(self._provider_stats.items())}

    def worker_summary(self) -> dict[str, dict]:
        """Per-worker aggregated stats."""
        with self._lock:
            return {name: stats.to_dict() for name, stats in sorted(self._worker_stats.items())}

    def dashboard_snapshot(self) -> dict:
        """Full metrics snapshot for Mission Control dashboard."""
        with self._lock:
            total = len(self._tasks)
            completed = sum(1 for t in self._tasks if t.status == TaskStatus.COMPLETED)
            failed = sum(1 for t in self._tasks if t.status == TaskStatus.FAILED)
        return {
            "uptime_seconds": round(time.time() - self._start_time, 1),
            "total_tasks": total,
            "completed_tasks": completed,
            "failed_tasks": failed,
            "tasks_per_minute": round(self.tasks_per_minute(), 2),
            "total_cost_usd": round(self.total_cost(), 6),
            "error_rate": round(self.error_rate(), 4),
            "providers": self.provider_summary(),
            "workers": self.worker_summary(),
            "timestamp": time.time(),
        }

    def recent_tasks(self, limit: int = 20) -> list[dict]:
        """Return the most recent *limit* task dicts."""
        with self._lock:
            return [t.to_dict() for t in self._tasks[-limit:]]

    def reset(self) -> None:
        """Reset all metrics (for testing)."""
        with self._lock:
            self._tasks.clear()
            self._provider_stats.clear()
            self._worker_stats.clear()
            self._start_time = time.time()
