"""
Metrics system for simulation evaluation.

Collects and reports:
- sent_packets, received_packets
- delivery_ratio
- avg_delay
- avg_hop
- retransmissions
- collisions
- convergence_time
"""

from typing import Dict, List, Set
from dataclasses import dataclass, field

from sim import config


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

    # 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 due to collision

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

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

    # Channel
    collisions: int = 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)

    def calculate_pdr(self) -> float:
        """Calculate Packet Delivery Ratio (unique packets at sink / sent)."""
        unique_received = len(self.received_packet_ids)
        if self.total_sent == 0:
            return 0.0
        return unique_received / self.total_sent

    def calculate_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)

    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."""
        unique_received = len(self.received_packet_ids)
        return {
            "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),
            "avg_hop": round(self.calculate_avg_hop(), 2),
            "avg_retries": round(self.calculate_avg_retries(), 2),
            "convergence_time": round(self.convergence_time, 2),
            "collisions": self.collisions,
            "route_updates": self.route_updates,
        }


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

    def __init__(self):
        self.metrics = SimulationMetrics()
        self.start_time = 0.0

    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

    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

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

        if is_sink:
            # For sink, data_received is actual unique packets received
            # Track unique (src, seq) pairs
            pass  # Will handle sink specially below
        else:
            self.metrics.total_sent += node_stats.get("data_sent", 0)

        self.metrics.total_forwarded += node_stats.get("data_forwarded", 0)
        self.metrics.total_dropped += node_stats.get("packets_dropped", 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 (unique packet delivery tracking)."""
        self.metrics.node_stats[node_id] = stats

        node_stats = stats.get("stats", {})
        # Count how many unique packets the sink received
        # This is the actual delivery count for PDR
        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)

        # Update rest of stats
        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 received packet."""
        self.metrics.hop_counts.append(hops)

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