"""
Extended Metrics system for Phase-2 Validation & Analysis.

New metrics added:
- Route stability (route_changes, parent_history, cost_history)
- Hop distribution (hop_histogram)
- Channel utilization (busy_time, idle_time, collision_time)
- Loss breakdown (LOSS_COLLISION, LOSS_NO_ROUTE, LOSS_RETRY_EXCEEDED, LOSS_CHANNEL_BUSY)
"""

from typing import Dict, List, Set, Tuple
from dataclasses import dataclass, field
from enum import Enum

from sim import config


class LossType(Enum):
    """Packet loss types."""

    LOSS_COLLISION = "collision"
    LOSS_NO_ROUTE = "no_route"
    LOSS_RETRY_EXCEEDED = "retry_exceeded"
    LOSS_CHANNEL_BUSY = "channel_busy"


@dataclass
class SimulationMetrics:
    """Metrics for the entire simulation."""

    # Basic packet counts
    total_sent: int = 0  # Data packets generated (all nodes)
    total_received: int = 0  # Data packets received at sink
    total_forwarded: int = 0  # Data packets forwarded by nodes
    total_dropped: int = 0  # Packets dropped (all reasons)

    # Routing
    convergence_time: float = 0.0
    route_updates: int = 0

    # MAC
    retries: int = 0
    acks_received: int = 0

    # Channel
    collisions: int = 0
    channel_busy_time: float = 0.0
    channel_idle_time: float = 0.0

    # Hop statistics
    hop_counts: List[int] = field(default_factory=list)

    # Per-node stats
    node_stats: Dict[int, dict] = field(default_factory=dict)

    # Track unique packets received at sink
    received_packet_ids: Set[tuple] = field(default_factory=set)

    # ========================================================================
    # NEW: Route Stability Metrics (3.1)
    # ========================================================================
    route_changes: Dict[int, int] = field(default_factory=dict)
    parent_history: Dict[int, List[Tuple[float, int]]] = field(default_factory=dict)
    cost_history: Dict[int, List[Tuple[float, int]]] = field(default_factory=dict)

    # ========================================================================
    # NEW: Hop Distribution (3.2)
    # ========================================================================
    hop_histogram: Dict[int, int] = field(default_factory=dict)

    # ========================================================================
    # NEW: Loss Breakdown (3.4)
    # ========================================================================
    loss_collision: int = 0
    loss_no_route: int = 0
    loss_retry_exceeded: int = 0
    loss_channel_busy: int = 0

    # =========================================================================
    # Route Stability Calculations
    # =========================================================================
    def calculate_route_change_rate(self, sim_time: float) -> float:
        """Calculate route change rate (changes per second)."""
        total_changes = sum(self.route_changes.values())
        if sim_time <= 0:
            return 0.0
        return total_changes / sim_time

    def get_parent_history(self, node_id: int) -> List[Tuple[float, int]]:
        """Get parent history for a node."""
        return self.parent_history.get(node_id, [])

    def get_cost_history(self, node_id: int) -> List[Tuple[float, int]]:
        """Get cost history for a node."""
        return self.cost_history.get(node_id, [])

    # =========================================================================
    # Hop Distribution
    # =========================================================================
    def calculate_hop_histogram(self) -> Dict[int, int]:
        """Calculate hop distribution histogram."""
        histogram = {}
        for hop in self.hop_counts:
            histogram[hop] = histogram.get(hop, 0) + 1
        self.hop_histogram = histogram
        return histogram

    def get_max_hop(self) -> int:
        """Get maximum hop count."""
        return max(self.hop_counts) if self.hop_counts else 0

    def get_avg_hop(self) -> float:
        """Calculate average hop count."""
        if not self.hop_counts:
            return 0.0
        return sum(self.hop_counts) / len(self.hop_counts)

    # =========================================================================
    # Channel Utilization (3.3)
    # =========================================================================
    def calculate_channel_utilization(self, sim_time: float) -> float:
        """Calculate channel utilization ratio."""
        if sim_time <= 0:
            return 0.0
        total_time = self.channel_busy_time + self.channel_idle_time
        if total_time <= 0:
            return 0.0
        return self.channel_busy_time / total_time

    def get_channel_stats(self) -> dict:
        """Get channel statistics."""
        total_time = self.channel_busy_time + self.channel_idle_time
        utilization = self.channel_busy_time / total_time if total_time > 0 else 0

        return {
            "busy_time": self.channel_busy_time,
            "idle_time": self.channel_idle_time,
            "total_time": total_time,
            "utilization": utilization,
            "collision_count": self.collisions,
        }

    # =========================================================================
    # Loss Breakdown (3.4)
    # =========================================================================
    def calculate_loss_rates(self) -> Dict[str, float]:
        """Calculate loss rates by type."""
        total_loss = (
            self.loss_collision
            + self.loss_no_route
            + self.loss_retry_exceeded
            + self.loss_channel_busy
        )
        if total_loss == 0:
            return {}

        return {
            "collision": round(self.loss_collision / total_loss * 100, 2)
            if self.loss_collision > 0
            else 0,
            "no_route": round(self.loss_no_route / total_loss * 100, 2)
            if self.loss_no_route > 0
            else 0,
            "retry_exceeded": round(self.loss_retry_exceeded / total_loss * 100, 2)
            if self.loss_retry_exceeded > 0
            else 0,
            "channel_busy": round(self.loss_channel_busy / total_loss * 100, 2)
            if self.loss_channel_busy > 0
            else 0,
        }

    # =========================================================================
    # Standard Metrics
    # =========================================================================
    def calculate_pdr(self) -> float:
        """Calculate Packet Delivery Ratio."""
        unique_received = len(self.received_packet_ids)
        if self.total_sent == 0:
            return 0.0
        return unique_received / self.total_sent

    def calculate_avg_retries(self) -> float:
        """Calculate average retries per packet."""
        if self.total_sent == 0:
            return 0.0
        return self.retries / self.total_sent

    def get_summary(self) -> dict:
        """Get metrics summary."""
        # Calculate derived metrics
        hop_hist = self.calculate_hop_histogram()
        max_hop = self.get_max_hop()
        avg_hop = self.get_avg_hop()
        loss_rates = self.calculate_loss_rates()
        unique_received = len(self.received_packet_ids)

        return {
            # Basic
            "total_sent": self.total_sent,
            "total_received": unique_received,
            "total_forwarded": self.total_forwarded,
            "total_dropped": self.total_dropped,
            "pdr": round(self.calculate_pdr() * 100, 2),
            # Hop distribution
            "max_hop": max_hop,
            "avg_hop": round(avg_hop, 2),
            "hop_histogram": hop_hist,
            "MULTIHOP_FORMED": max_hop >= 2,
            # Route stability
            "route_changes": sum(self.route_changes.values()),
            "route_change_rate": round(
                self.calculate_route_change_rate(config.SIM_TIME), 4
            ),
            # Channel utilization
            "collisions": self.collisions,
            "channel_utilization": round(
                self.calculate_channel_utilization(config.SIM_TIME) * 100, 2
            ),
            # Loss breakdown
            "loss_collision": self.loss_collision,
            "loss_no_route": self.loss_no_route,
            "loss_retry_exceeded": self.loss_retry_exceeded,
            "loss_channel_busy": self.loss_channel_busy,
            "loss_rates": loss_rates,
            # Timing
            "convergence_time": round(self.convergence_time, 2),
            # Legacy
            "avg_retries": round(self.calculate_avg_retries(), 2),
        }


class MetricsCollector:
    """Collects metrics from simulation."""

    def __init__(self):
        self.metrics = SimulationMetrics()
        self.start_time = 0.0
        self._last_sample_time = 0.0
        self._time_series_data: List[dict] = []

    def set_start_time(self, time: float):
        """Set simulation start time."""
        self.start_time = time

    def set_convergence_time(self, time: float):
        """Set convergence time."""
        self.metrics.convergence_time = time - self.start_time

    # =========================================================================
    # Route Stability Tracking
    # =========================================================================
    def record_route_change(self, node_id: int, new_parent: int, time: float):
        """Record a route change event."""
        if node_id not in self.metrics.route_changes:
            self.metrics.route_changes[node_id] = 0
        self.metrics.route_changes[node_id] += 1

        # Record history
        if node_id not in self.metrics.parent_history:
            self.metrics.parent_history[node_id] = []
        self.metrics.parent_history[node_id].append((time, new_parent))

    def record_cost_change(self, node_id: int, new_cost: int, time: float):
        """Record a cost change event."""
        if node_id not in self.metrics.cost_history:
            self.metrics.cost_history[node_id] = []
        self.metrics.cost_history[node_id].append((time, new_cost))

    # =========================================================================
    # Hop Distribution Tracking
    # =========================================================================
    def record_hop_count(self, hops: int):
        """Record hop count for a packet."""
        self.metrics.hop_counts.append(hops)

    # =========================================================================
    # Channel Utilization Tracking
    # =========================================================================
    def record_channel_busy(self, duration: float):
        """Record channel busy time."""
        self.metrics.channel_busy_time += duration

    def record_channel_idle(self, duration: float):
        """Record channel idle time."""
        self.metrics.channel_idle_time += duration

    # =========================================================================
    # Loss Breakdown Tracking
    # =========================================================================
    def record_collision_loss(self):
        """Record collision loss."""
        self.metrics.loss_collision += 1
        self.metrics.total_dropped += 1

    def record_no_route_loss(self):
        """Record loss due to no route."""
        self.metrics.loss_no_route += 1
        self.metrics.total_dropped += 1

    def record_retry_exceeded_loss(self):
        """Record loss due to max retries exceeded."""
        self.metrics.loss_retry_exceeded += 1
        self.metrics.total_dropped += 1

    def record_channel_busy_loss(self):
        """Record loss due to channel busy."""
        self.metrics.loss_channel_busy += 1
        self.metrics.total_dropped += 1

    # =========================================================================
    # Standard Stats Collection
    # =========================================================================
    def add_node_stats(self, node_id: int, stats: dict, is_sink: bool = False):
        """Add per-node statistics."""
        self.metrics.node_stats[node_id] = stats

        node_stats = stats.get("stats", {})

        if not is_sink:
            self.metrics.total_sent += node_stats.get("data_sent", 0)

        self.metrics.total_forwarded += node_stats.get("data_forwarded", 0)
        self.metrics.route_updates += node_stats.get("route_updates", 0)

        # MAC stats
        mac_stats = stats.get("mac", {})
        self.metrics.retries += mac_stats.get("retries", 0)
        self.metrics.acks_received += mac_stats.get("received_acks", 0)

    def add_sink_stats(self, node_id: int, stats: dict):
        """Add sink-specific statistics."""
        self.metrics.node_stats[node_id] = stats

        node_stats = stats.get("stats", {})
        received = node_stats.get("data_received", 0)
        self.metrics.total_received = received

        # Also track all packets sent
        for nid, nstats in self.metrics.node_stats.items():
            if nid != node_id:
                self.metrics.total_sent += nstats.get("stats", {}).get("data_sent", 0)

        self.metrics.total_dropped += node_stats.get("packets_dropped", 0)

        mac_stats = stats.get("mac", {})
        self.metrics.retries += mac_stats.get("retries", 0)
        self.metrics.acks_received += mac_stats.get("received_acks", 0)

    def add_collision(self, count: int = 1):
        """Add collision count."""
        self.metrics.collisions += count

    def add_hop_count(self, hops: int):
        """Add hop count for a packet."""
        self.metrics.hop_counts.append(hops)

    def get_metrics(self) -> SimulationMetrics:
        """Get collected metrics."""
        return self.metrics

    # =========================================================================
    # Time Series Sampling
    # =========================================================================
    def should_sample(self, current_time: float, sample_interval: float = 1.0) -> bool:
        """Check if it's time to take a sample."""
        if current_time - self._last_sample_time >= sample_interval:
            self._last_sample_time = current_time
            return True
        return False

    def record_time_series_sample(self, current_time: float):
        """Record a time series sample."""
        sample = {
            "time": round(current_time, 2),
            "avg_cost": 0,  # Would need per-node tracking
            "route_changes": sum(self.metrics.route_changes.values()),
            "channel_utilization": self.metrics.channel_busy_time / current_time
            if current_time > 0
            else 0,
            "pdr": len(self.metrics.received_packet_ids) / self.metrics.total_sent
            if self.metrics.total_sent > 0
            else 0,
        }
        self._time_series_data.append(sample)

    def get_time_series_data(self) -> List[dict]:
        """Get recorded time series data."""
        return self._time_series_data