"""
Metrics — aggregate statistics for the Nexus runtime.

Design notes:
  - Collects counters, gauges, and histograms
  - Tracks: tasks completed/failed, tokens consumed, cost, latency distributions
  - Per-agent and per-task-type breakdowns
  - Optional export to OpenTelemetry metrics
"""

from __future__ import annotations

import logging
import time
from collections import defaultdict
from dataclasses import dataclass, field
from typing import Any

logger = logging.getLogger("nexus.observability.metrics")


@dataclass
class Counter:
    """A monotonically increasing counter."""
    name: str
    value: float = 0.0
    labels: dict[str, str] = field(default_factory=dict)

    def increment(self, amount: float = 1.0) -> None:
        self.value += amount


@dataclass
class Gauge:
    """A value that can go up and down."""
    name: str
    value: float = 0.0

    def set(self, value: float) -> None:
        self.value = value

    def increment(self, amount: float = 1.0) -> None:
        self.value += amount

    def decrement(self, amount: float = 1.0) -> None:
        self.value -= amount


@dataclass
class Histogram:
    """Tracks distribution of values (e.g., latency)."""
    name: str
    values: list[float] = field(default_factory=list)

    def observe(self, value: float) -> None:
        self.values.append(value)

    @property
    def count(self) -> int:
        return len(self.values)

    @property
    def mean(self) -> float:
        return sum(self.values) / len(self.values) if self.values else 0.0

    @property
    def p50(self) -> float:
        return self._percentile(50)

    @property
    def p95(self) -> float:
        return self._percentile(95)

    @property
    def p99(self) -> float:
        return self._percentile(99)

    def _percentile(self, pct: float) -> float:
        if not self.values:
            return 0.0
        sorted_vals = sorted(self.values)
        idx = int(len(sorted_vals) * pct / 100)
        idx = min(idx, len(sorted_vals) - 1)
        return sorted_vals[idx]


class MetricsCollector:
    """
    Centralized metrics collection for the Nexus runtime.

    Usage:
        metrics = MetricsCollector()

        # Track task completion
        metrics.task_completed("research", "agent-123", latency=5.2, tokens=1500)

        # Track task failure
        metrics.task_failed("coding", "agent-456", error="timeout")

        # Get summary
        print(metrics.summary())
    """

    def __init__(self) -> None:
        # Counters
        self.tasks_total = Counter("nexus_tasks_total")
        self.tasks_completed = Counter("nexus_tasks_completed")
        self.tasks_failed = Counter("nexus_tasks_failed")
        self.tokens_total = Counter("nexus_tokens_total")
        self.cost_total = Counter("nexus_cost_usd_total")
        self.llm_calls_total = Counter("nexus_llm_calls_total")
        self.tool_calls_total = Counter("nexus_tool_calls_total")
        self.handoffs_total = Counter("nexus_handoffs_total")

        # Gauges
        self.active_tasks = Gauge("nexus_active_tasks")
        self.active_agents = Gauge("nexus_active_agents")

        # Histograms
        self.task_latency = Histogram("nexus_task_latency_seconds")
        self.llm_latency = Histogram("nexus_llm_latency_seconds")
        self.tokens_per_task = Histogram("nexus_tokens_per_task")

        # Per-agent breakdowns
        self._agent_tasks: dict[str, int] = defaultdict(int)
        self._agent_tokens: dict[str, int] = defaultdict(int)
        self._agent_failures: dict[str, int] = defaultdict(int)

        # Per-task-type breakdowns
        self._type_counts: dict[str, int] = defaultdict(int)
        self._type_latencies: dict[str, list[float]] = defaultdict(list)

        self._start_time = time.time()

    def task_started(self, task_type: str = "general") -> None:
        """Record that a task has started."""
        self.tasks_total.increment()
        self.active_tasks.increment()
        self._type_counts[task_type] += 1

    def task_completed(
        self,
        task_type: str = "general",
        agent_id: str = "",
        latency: float = 0.0,
        tokens: int = 0,
        cost: float = 0.0,
    ) -> None:
        """Record a successful task completion."""
        self.tasks_completed.increment()
        self.active_tasks.decrement()
        self.task_latency.observe(latency)
        self.tokens_per_task.observe(tokens)
        self.tokens_total.increment(tokens)
        self.cost_total.increment(cost)
        self._type_latencies[task_type].append(latency)

        if agent_id:
            self._agent_tasks[agent_id] += 1
            self._agent_tokens[agent_id] += tokens

    def task_failed(
        self,
        task_type: str = "general",
        agent_id: str = "",
        error: str = "",
    ) -> None:
        """Record a task failure."""
        self.tasks_failed.increment()
        self.active_tasks.decrement()

        if agent_id:
            self._agent_failures[agent_id] += 1

    def llm_call(self, latency: float = 0.0, tokens: int = 0) -> None:
        """Record an LLM call."""
        self.llm_calls_total.increment()
        self.llm_latency.observe(latency)
        self.tokens_total.increment(tokens)

    def tool_call(self) -> None:
        """Record a tool call."""
        self.tool_calls_total.increment()

    def handoff(self) -> None:
        """Record a handoff."""
        self.handoffs_total.increment()

    def summary(self) -> dict[str, Any]:
        """Generate a comprehensive metrics summary."""
        uptime = time.time() - self._start_time

        return {
            "uptime_seconds": round(uptime, 1),
            "tasks": {
                "total": int(self.tasks_total.value),
                "completed": int(self.tasks_completed.value),
                "failed": int(self.tasks_failed.value),
                "active": int(self.active_tasks.value),
                "success_rate": (
                    round(self.tasks_completed.value / self.tasks_total.value, 3)
                    if self.tasks_total.value > 0
                    else 0.0
                ),
            },
            "tokens": {
                "total": int(self.tokens_total.value),
                "per_task_avg": round(self.tokens_per_task.mean, 0),
                "per_task_p95": round(self.tokens_per_task.p95, 0),
            },
            "cost_usd": round(self.cost_total.value, 4),
            "latency": {
                "mean_seconds": round(self.task_latency.mean, 2),
                "p50_seconds": round(self.task_latency.p50, 2),
                "p95_seconds": round(self.task_latency.p95, 2),
                "p99_seconds": round(self.task_latency.p99, 2),
            },
            "llm_calls": int(self.llm_calls_total.value),
            "tool_calls": int(self.tool_calls_total.value),
            "handoffs": int(self.handoffs_total.value),
            "agent_breakdown": {
                agent_id: {
                    "tasks": count,
                    "tokens": self._agent_tokens.get(agent_id, 0),
                    "failures": self._agent_failures.get(agent_id, 0),
                }
                for agent_id, count in self._agent_tasks.items()
            },
        }

    def reset(self) -> None:
        """Reset all metrics."""
        self.__init__()  # type: ignore
