"""
runtime/agent_monitoring.py — Module 10: Agent Monitoring
═══════════════════════════════════════════════════════════════════════════════
Observability dashboard for the multi-agent runtime.
Like /proc, top, and htop combined — real-time view of agent health.

Extends existing observability layer with multi-agent metrics:
  • Per-agent latency, throughput, success rate
  • Resource utilization across agents
  • Communication patterns on the message bus
  • Anomaly detection: spike in failures, latency degradation
  • Runtime health score
"""

from __future__ import annotations
import time
from collections import defaultdict
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional, Tuple
from infrarely.observability import logger


@dataclass
class AgentMetrics:
    """Per-agent performance metrics."""

    agent_id: str
    tasks_completed: int = 0
    tasks_failed: int = 0
    total_latency_ms: float = 0.0
    tokens_consumed: int = 0
    messages_sent: int = 0
    messages_received: int = 0
    errors: List[Dict[str, Any]] = field(default_factory=list)
    latency_samples: List[float] = field(default_factory=list)

    @property
    def avg_latency_ms(self) -> float:
        if not self.latency_samples:
            return 0.0
        return sum(self.latency_samples) / len(self.latency_samples)

    @property
    def p95_latency_ms(self) -> float:
        if not self.latency_samples:
            return 0.0
        sorted_samples = sorted(self.latency_samples)
        idx = int(len(sorted_samples) * 0.95)
        return sorted_samples[min(idx, len(sorted_samples) - 1)]

    @property
    def success_rate(self) -> float:
        total = self.tasks_completed + self.tasks_failed
        if total == 0:
            return 1.0
        return self.tasks_completed / total

    @property
    def throughput(self) -> float:
        """Tasks per minute (rough estimate from recent latency samples)."""
        total = self.tasks_completed + self.tasks_failed
        return total  # simplified — enhanced with time windows in practice

    def record_task(self, success: bool, latency_ms: float, tokens: int = 0):
        if success:
            self.tasks_completed += 1
        else:
            self.tasks_failed += 1
        self.total_latency_ms += latency_ms
        self.tokens_consumed += tokens
        self.latency_samples.append(latency_ms)
        if len(self.latency_samples) > 200:
            self.latency_samples = self.latency_samples[-100:]

    def record_error(self, error: str, context: str = ""):
        self.errors.append(
            {
                "error": error,
                "context": context,
                "timestamp": time.time(),
            }
        )
        if len(self.errors) > 50:
            self.errors = self.errors[-25:]


@dataclass
class RuntimeAlert:
    """An alert generated by the monitoring system."""

    alert_id: str
    severity: str  # "info", "warning", "critical"
    agent_id: str
    alert_type: str
    message: str
    timestamp: float = field(default_factory=time.time)
    acknowledged: bool = False


class AgentMonitoring:
    """
    Central monitoring system for multi-agent runtime.

    Features:
      • Per-agent metrics collection
      • Anomaly detection (failure spikes, latency degradation)
      • Runtime health score (0–100)
      • Alert system for critical conditions
      • Communication pattern analysis
    """

    LATENCY_ANOMALY_FACTOR = 3.0  # 3× avg = anomaly
    FAILURE_RATE_ALERT = 0.3  # 30% failure rate triggers alert
    MIN_SAMPLES_FOR_ANOMALY = 5

    def __init__(self):
        self._metrics: Dict[str, AgentMetrics] = {}
        self._alerts: List[RuntimeAlert] = []
        self._alert_counter = 0
        self._runtime_start = time.time()
        self._global_tasks = 0
        self._global_failures = 0

    # ── Register agent ────────────────────────────────────────────────────────
    def register(self, agent_id: str):
        """Start tracking metrics for an agent."""
        if agent_id not in self._metrics:
            self._metrics[agent_id] = AgentMetrics(agent_id=agent_id)

    def unregister(self, agent_id: str):
        """Stop tracking (keeps data for historical analysis)."""
        pass  # keep metrics for dashboard — mark inactive if needed

    # ── Record events ─────────────────────────────────────────────────────────
    def record_task(
        self,
        agent_id: str,
        success: bool,
        latency_ms: float,
        tokens: int = 0,
    ):
        """Record a task outcome. Triggers anomaly checks."""
        metrics = self._metrics.get(agent_id)
        if not metrics:
            self.register(agent_id)
            metrics = self._metrics[agent_id]

        metrics.record_task(success, latency_ms, tokens)
        self._global_tasks += 1
        if not success:
            self._global_failures += 1

        # Anomaly detection
        self._check_anomalies(agent_id, metrics, latency_ms, success)

    def record_message(self, sender: str, recipient: str):
        """Track message bus communication patterns."""
        if sender in self._metrics:
            self._metrics[sender].messages_sent += 1
        if recipient in self._metrics:
            self._metrics[recipient].messages_received += 1

    def record_error(self, agent_id: str, error: str, context: str = ""):
        metrics = self._metrics.get(agent_id)
        if metrics:
            metrics.record_error(error, context)

    # ── Anomaly detection ─────────────────────────────────────────────────────
    def _check_anomalies(
        self,
        agent_id: str,
        metrics: AgentMetrics,
        latency_ms: float,
        success: bool,
    ):
        """Detect performance anomalies and generate alerts."""
        total = metrics.tasks_completed + metrics.tasks_failed

        # Latency anomaly
        if (
            total >= self.MIN_SAMPLES_FOR_ANOMALY
            and metrics.avg_latency_ms > 0
            and latency_ms > metrics.avg_latency_ms * self.LATENCY_ANOMALY_FACTOR
        ):
            self._create_alert(
                "warning",
                agent_id,
                "latency_spike",
                f"Latency spike: {latency_ms:.0f}ms "
                f"(avg={metrics.avg_latency_ms:.0f}ms)",
            )

        # High failure rate
        if total >= self.MIN_SAMPLES_FOR_ANOMALY:
            if metrics.success_rate < (1.0 - self.FAILURE_RATE_ALERT):
                self._create_alert(
                    "critical",
                    agent_id,
                    "high_failure_rate",
                    f"Failure rate: {1 - metrics.success_rate:.0%} "
                    f"({metrics.tasks_failed}/{total} tasks failed)",
                )

    def _create_alert(
        self,
        severity: str,
        agent_id: str,
        alert_type: str,
        message: str,
    ):
        self._alert_counter += 1
        alert = RuntimeAlert(
            alert_id=f"alert_{self._alert_counter:04d}",
            severity=severity,
            agent_id=agent_id,
            alert_type=alert_type,
            message=message,
        )
        self._alerts.append(alert)
        if len(self._alerts) > 200:
            self._alerts = self._alerts[-100:]

        log_fn = logger.error if severity == "critical" else logger.warn
        log_fn(f"Monitor ALERT [{severity}] {agent_id}: {message}")

    # ── Acknowledge alert ─────────────────────────────────────────────────────
    def acknowledge_alert(self, alert_id: str) -> bool:
        for alert in self._alerts:
            if alert.alert_id == alert_id:
                alert.acknowledged = True
                return True
        return False

    # ── Runtime health score ──────────────────────────────────────────────────
    def health_score(self) -> float:
        """
        Compute overall runtime health (0–100).
        Based on: avg success rate, alert count, resource usage.
        """
        if not self._metrics:
            return 100.0

        # Average success rate across agents
        total_success = sum(m.success_rate for m in self._metrics.values())
        avg_success = total_success / len(self._metrics) if self._metrics else 1.0

        # Active unacknowledged alerts penalty
        active_alerts = sum(1 for a in self._alerts if not a.acknowledged)
        critical_alerts = sum(
            1 for a in self._alerts if not a.acknowledged and a.severity == "critical"
        )

        # Calculate score
        score = 100.0
        score *= avg_success  # success rate factor
        score -= min(20, active_alerts * 2)  # alert penalty
        score -= min(30, critical_alerts * 10)  # critical alert penalty

        return max(0.0, min(100.0, score))

    # ── Agent ranking ─────────────────────────────────────────────────────────
    def top_agents(
        self, metric: str = "throughput", limit: int = 5
    ) -> List[Dict[str, Any]]:
        """Rank agents by a metric."""
        agents = list(self._metrics.values())
        if metric == "throughput":
            agents.sort(key=lambda m: m.tasks_completed + m.tasks_failed, reverse=True)
        elif metric == "success_rate":
            agents.sort(key=lambda m: m.success_rate, reverse=True)
        elif metric == "latency":
            agents.sort(key=lambda m: m.avg_latency_ms)
        elif metric == "tokens":
            agents.sort(key=lambda m: m.tokens_consumed, reverse=True)

        return [
            {
                "agent_id": m.agent_id,
                "tasks": m.tasks_completed + m.tasks_failed,
                "success_rate": f"{m.success_rate:.0%}",
                "avg_latency": f"{m.avg_latency_ms:.0f}ms",
                "p95_latency": f"{m.p95_latency_ms:.0f}ms",
                "tokens": m.tokens_consumed,
            }
            for m in agents[:limit]
        ]

    # ── Communication patterns ────────────────────────────────────────────────
    def communication_matrix(self) -> Dict[str, Dict[str, int]]:
        """Get message send/receive counts per agent."""
        return {
            aid: {
                "sent": m.messages_sent,
                "received": m.messages_received,
            }
            for aid, m in self._metrics.items()
        }

    # ── Query ─────────────────────────────────────────────────────────────────
    def get_metrics(self, agent_id: str) -> Optional[Dict[str, Any]]:
        m = self._metrics.get(agent_id)
        if not m:
            return None
        return {
            "tasks_completed": m.tasks_completed,
            "tasks_failed": m.tasks_failed,
            "success_rate": f"{m.success_rate:.0%}",
            "avg_latency_ms": round(m.avg_latency_ms, 1),
            "p95_latency_ms": round(m.p95_latency_ms, 1),
            "tokens_consumed": m.tokens_consumed,
            "messages_sent": m.messages_sent,
            "messages_received": m.messages_received,
            "recent_errors": m.errors[-5:],
        }

    def active_alerts(self) -> List[Dict[str, Any]]:
        return [
            {
                "id": a.alert_id,
                "severity": a.severity,
                "agent": a.agent_id,
                "type": a.alert_type,
                "message": a.message,
                "timestamp": a.timestamp,
            }
            for a in self._alerts
            if not a.acknowledged
        ]

    def snapshot(self) -> Dict[str, Any]:
        uptime_s = time.time() - self._runtime_start
        return {
            "uptime_seconds": round(uptime_s, 1),
            "health_score": round(self.health_score(), 1),
            "agents_monitored": len(self._metrics),
            "global_tasks": self._global_tasks,
            "global_failures": self._global_failures,
            "global_success_rate": (
                f"{1 - (self._global_failures / max(1, self._global_tasks)):.0%}"
            ),
            "active_alerts": len([a for a in self._alerts if not a.acknowledged]),
            "total_alerts": len(self._alerts),
            "top_agents": self.top_agents(limit=5),
            "alerts": self.active_alerts()[:5],
        }
