"""
Hybrid SimPy-Based Event-Driven Simulator for LEACH and LEACH-C Protocols.

This module implements a complete discrete event simulation using Simpy framework,
combining the best of both approaches:

FROM PAUL'S APPROACH:
- Full refactor to use SimPy properly with event-driven architecture
- Parallel node mobility processes
- Proper discrete event simulation model
- Sequential round events with all phases

FROM SORTI'S APPROACH:
- DRY (Don't Repeat Yourself) principles
- KISS (Keep It Simple) architecture
- Reusable helper methods
- Clean code structure with proper separation of concerns
- Comprehensive event logging
- Support for static/dynamic modes

Key Features:
- SimPy Environment for discrete event management
- Parallel node mobility as background processes
- Structured round phases: initialization → CH election → communication → cleanup
- Complete metrics collection
- Event logging for debugging and analysis
- Static/Dynamic network support via ENABLE_MOBILITY flag
"""

import simpy
import random
from typing import List, Dict, Optional, Tuple
from node import Node
from metrics import Metrics
from config import (
    FIELD_WIDTH, FIELD_HEIGHT, INITIAL_ENERGY, BS_POSITION,
    ENABLE_MOBILITY, DEBUG
)


class HybridSimPySimulator:
    """
    Hybrid event-driven simulator combining Paul's full refactor with Sorti's quality approach.

    Architecture:
    - Main process: Handles round execution (election, communication, metrics)
    - Background processes: Node mobility (runs in parallel)
    - Event model: Discrete events at each round boundary
    - State management: Proper environment tracking

    Args:
        protocol_name: "LEACH" or "LEACH-C"
        nodes: List of Node objects
        packet_size: Data packet size in bits
        probability_ch: Probability of becoming cluster head
        max_rounds: Maximum simulation rounds
    """

    def __init__(
        self,
        protocol_name: str,
        nodes: List[Node],
        packet_size: int,
        probability_ch: float,
        max_rounds: int
    ):
        self.env = simpy.Environment()
        self.protocol_name = protocol_name
        self.nodes = nodes
        self.packet_size = packet_size
        self.probability_ch = probability_ch
        self.max_rounds = max_rounds

        # State management
        self.round_num = 0
        self.cluster_heads = []
        self.clusters: Dict[int, List[int]] = {}
        self.events_log = []

        # Statistics
        self.total_packets_to_ch = 0
        self.total_packets_to_bs = 0
        self.muted_rounds = []

        # Metrics collector
        self.metrics = Metrics()

    # ========== DRY: Reusable Helper Methods ==========

    def _log_event(self, event_type: str, round_num: int = 0, **details) -> None:
        """
        DRY: Single method for all event logging.

        Args:
            event_type: Type of event (e.g., 'CH_ELECTED', 'COMMUNICATION', 'MOBILITY')
            round_num: Current round number
            details: Additional event details
        """
        self.events_log.append({
            'time': self.env.now,
            'event': event_type,
            'round': round_num,
            **details
        })

    def _get_alive_nodes(self) -> List[Node]:
        """DRY: Get list of currently alive nodes."""
        return [n for n in self.nodes if n.is_alive]

    def _find_closest_cluster_head(self, node: Node) -> Optional[int]:
        """
        DRY: Find the closest cluster head for a node.

        Args:
            node: Node looking for closest CH

        Returns:
            Cluster head node ID or None if no CHs exist
        """
        if not self.cluster_heads:
            return None

        closest_ch = None
        min_distance = float('inf')

        for ch_id in self.cluster_heads:
            ch_node = self.nodes[ch_id]
            distance = node.distance_to(ch_node.x, ch_node.y)
            if distance < min_distance:
                min_distance = distance
                closest_ch = ch_id

        return closest_ch

    # ========== CH Election: Protocol-Specific ==========

    def _elect_cluster_heads_leach(self) -> None:
        """
        LEACH: Distributed cluster head election.
        Each alive node has probability p of becoming a CH.
        """
        self.cluster_heads = []
        self.clusters = {}

        # Phase 1: Nodes decide if they become CH
        for node in self._get_alive_nodes():
            if random.random() < self.probability_ch:
                node.is_cluster_head = True
                self.cluster_heads.append(node.node_id)
                self.clusters[node.node_id] = [node.node_id]
                node.cluster_id = node.node_id

        # Phase 2: Non-CH nodes join closest CH
        for node in self._get_alive_nodes():
            if not node.is_cluster_head:
                closest_ch = self._find_closest_cluster_head(node)
                if closest_ch is not None:
                    node.cluster_id = closest_ch
                    if closest_ch not in self.clusters:
                        self.clusters[closest_ch] = []
                    self.clusters[closest_ch].append(node.node_id)

        self._log_event(
            'CH_ELECTED_LEACH',
            self.round_num,
            num_ch=len(self.cluster_heads),
            alive_nodes=len(self._get_alive_nodes())
        )

    def _elect_cluster_heads_leachc(self) -> None:
        """
        LEACH-C: Centralized cluster head election by base station.
        BS selects top 10% nodes by energy as CHs.
        """
        self.cluster_heads = []
        self.clusters = {}

        alive_nodes = self._get_alive_nodes()
        if not alive_nodes:
            return

        # Phase 1: BS collects node information
        # Each node sends ~32 bits (position + energy) to BS
        for node in alive_nodes:
            distance_to_bs = node.distance_to(*BS_POSITION)
            node.transmit(32, distance_to_bs)

        # Phase 2: BS selects CHs (top 10% by energy)
        num_expected_ch = max(1, int(len(alive_nodes) * 0.1))
        sorted_nodes = sorted(alive_nodes, key=lambda n: n.energy, reverse=True)
        selected_ch = sorted_nodes[:num_expected_ch]

        for node in selected_ch:
            node.is_cluster_head = True
            self.cluster_heads.append(node.node_id)
            self.clusters[node.node_id] = [node.node_id]
            node.cluster_id = node.node_id

        # Phase 3: BS broadcasts CH list to all nodes
        for node in alive_nodes:
            if not node.is_cluster_head:
                distance_to_bs = node.distance_to(*BS_POSITION)
                node.receive(len(self.cluster_heads) * 8)

        # Phase 4: Non-CH nodes join closest CH
        for node in alive_nodes:
            if not node.is_cluster_head:
                closest_ch = self._find_closest_cluster_head(node)
                if closest_ch is not None:
                    node.cluster_id = closest_ch
                    if closest_ch not in self.clusters:
                        self.clusters[closest_ch] = []
                    self.clusters[closest_ch].append(node.node_id)

        self._log_event(
            'CH_ELECTED_LEACHC',
            self.round_num,
            num_ch=len(self.cluster_heads),
            alive_nodes=len(alive_nodes)
        )

    # ========== Communication Phase ==========

    def _communication_phase(self) -> None:
        """
        Execute communication phase: transmission from nodes to CH to BS.
        """
        if not self.cluster_heads:
            # Muted round: no CHs available
            self.muted_rounds.append(self.round_num)
            self._log_event('MUTED_ROUND', self.round_num)
            return

        packets_this_round = {'to_ch': 0, 'to_bs': 0}

        # Phase 1: Non-CH nodes send to their CH
        for node in self._get_alive_nodes():
            if node.is_alive and not node.is_cluster_head:
                # Data transmission probability
                if random.random() < self.probability_ch:
                    ch_node = self.nodes[node.cluster_id] if node.cluster_id else None
                    if ch_node and ch_node.is_alive:
                        distance = node.distance_to(ch_node.x, ch_node.y)
                        node.transmit(self.packet_size, distance)
                        ch_node.receive(self.packet_size)
                        packets_this_round['to_ch'] += 1
                        self.total_packets_to_ch += 1

        # Phase 2: CHs aggregate and send to BS
        for ch_id in self.cluster_heads:
            ch_node = self.nodes[ch_id]
            if ch_node.is_alive:
                # Aggregation: nodes in cluster - 1 (excluding CH itself)
                num_packets = len(self.clusters.get(ch_id, [1])) - 1

                if num_packets > 0:
                    aggregated_data = self.packet_size
                    ch_node.aggregate(aggregated_data)

                    distance_to_bs = ch_node.distance_to(*BS_POSITION)
                    ch_node.transmit(aggregated_data, distance_to_bs)
                    packets_this_round['to_bs'] += 1
                    self.total_packets_to_bs += 1

        self._log_event(
            'COMMUNICATION',
            self.round_num,
            packets_to_ch=packets_this_round['to_ch'],
            packets_to_bs=packets_this_round['to_bs']
        )

    # ========== Mobility Phase ==========

    def _mobility_phase(self) -> None:
        """Execute mobility phase: update node positions (if enabled)."""
        if not ENABLE_MOBILITY:
            return

        moved_count = 0
        for node in self._get_alive_nodes():
            node.move()
            moved_count += 1

        if moved_count > 0:
            self._log_event('MOBILITY', self.round_num, nodes_moved=moved_count)

    # ========== Parallel Node Mobility Process ==========

    def _node_mobility_background_process(self, node: Node) -> None:
        """
        Background SimPy process for continuous node mobility.
        Runs in parallel with main simulation, independent of rounds.

        This is optional - for more realistic continuous movement.
        Currently, we use discrete mobility in each round.

        Args:
            node: Node to move
        """
        while node.is_alive and self.round_num < self.max_rounds:
            yield self.env.timeout(1.0)  # Wait for 1 time unit
            if ENABLE_MOBILITY and node.is_alive:
                node.move()

    # ========== Main Round Process ==========

    def _round_process(self) -> None:
        """
        Main SimPy process: Execute protocol rounds as discrete events.

        Each iteration:
        1. Advance time
        2. Reset node states
        3. Elect cluster heads
        4. Execute communication
        5. Update mobility
        6. Record metrics
        7. Check termination condition
        """
        while self.round_num < self.max_rounds:
            # Advance time by 1 round unit
            yield self.env.timeout(1.0)

            # Reset node states for this round
            for node in self.nodes:
                node.reset_for_round()

            # CH Election
            if self.protocol_name == "LEACH":
                self._elect_cluster_heads_leach()
            elif self.protocol_name == "LEACH-C":
                self._elect_cluster_heads_leachc()

            # Communication Phase
            self._communication_phase()

            # Mobility Phase
            self._mobility_phase()

            # Record metrics for this round
            alive_nodes = self._get_alive_nodes()
            self.metrics.record_round(
                round_num=self.round_num,
                nodes=self.nodes,
                ch_nodes=[self.nodes[ch_id] for ch_id in self.cluster_heads],
                packets_to_ch=self.total_packets_to_ch,
                packets_to_bs=self.total_packets_to_bs,
                muted=(len(self.cluster_heads) == 0)
            )

            # Update dead node tracking
            self.metrics.update_dead_nodes(self.nodes)

            # Debug output
            if DEBUG and self.round_num % 100 == 0:
                alive_count = len(alive_nodes)
                avg_energy = sum(n.energy for n in self.nodes) / len(self.nodes)
                print(
                    f"  Round {self.round_num}: {alive_count} alive, "
                    f"{len(self.cluster_heads)} CHs, avg_energy={avg_energy:.2e}"
                )

            self.round_num += 1

            # Termination: All nodes dead
            if not alive_nodes:
                if DEBUG:
                    print(f"  All nodes dead at round {self.round_num}")
                break

    # ========== Simulation Execution ==========

    def run(self) -> Dict:
        """
        Execute the complete simulation.

        Starts main round process and optional background mobility processes.
        Returns complete metrics after simulation ends.

        Returns:
            Dict with all metrics (FDN, FMR, DLBI, RSPI, etc.)
        """
        if DEBUG:
            print(f"\n{'='*60}")
            print(f"Hybrid SimPy Simulation")
            print(f"Protocol: {self.protocol_name}")
            print(f"Nodes: {len(self.nodes)}, Packet size: {self.packet_size}")
            print(f"Probability CH: {self.probability_ch}, Max rounds: {self.max_rounds}")
            print(f"{'='*60}")

        # Start background mobility processes (optional)
        if ENABLE_MOBILITY:
            for node in self.nodes:
                self.env.process(self._node_mobility_background_process(node))

        # Start main round process
        self.env.process(self._round_process())

        # Execute simulation
        self.env.run()

        # Calculate final metrics
        # FDN and FMR are tracked during execution
        fdn = self.metrics.first_dead_node_round
        fmr = self.metrics.first_muted_round
        dlbi = self.metrics.calculate_dlbi()
        rspi = self.metrics.calculate_rspi(self.max_rounds)

        if DEBUG:
            print(f"\n{self.protocol_name} Results:")
            print(f"  FDN: {fdn}, FMR: {fmr}")
            print(f"  DLBI: {dlbi:.4f}, RSPI: {rspi:.4f}")

        return {
            "fdn": fdn,
            "fmr": fmr,
            "dlbi": dlbi,
            "rspi": rspi,
            "metrics": self.metrics,
            "rounds_data": self.metrics.rounds_data,
            "events_log": self.events_log,
            "num_nodes": len(self.nodes),
            "num_rounds": self.round_num,
            "total_packets_to_ch": self.total_packets_to_ch,
            "total_packets_to_bs": self.total_packets_to_bs
        }


if __name__ == "__main__":
    """
    Demo: Hybrid SimPy Simulator combining best of both approaches.
    Shows full discrete event simulation with parallel processes.
    """
    from config import get_num_rounds_for_scenario

    print("=" * 70)
    print("HYBRID SIMPY SIMULATOR DEMONSTRATION")
    print("Combines Paul's full refactor with Sorti's quality approach")
    print("=" * 70)

    # Create test scenario
    random.seed(42)
    num_nodes = 50
    packet_size = 2000
    probability_ch = 0.05
    max_rounds = get_num_rounds_for_scenario(num_nodes)

    # Create nodes
    test_nodes = []
    for i in range(num_nodes):
        x = random.uniform(0, FIELD_WIDTH)
        y = random.uniform(0, FIELD_HEIGHT)
        test_nodes.append(Node(i, x, y, INITIAL_ENERGY))

    # Run LEACH simulation
    print(f"\nRunning LEACH with {num_nodes} nodes, {max_rounds} rounds...")
    sim_leach = HybridSimPySimulator(
        protocol_name="LEACH",
        nodes=test_nodes,
        packet_size=packet_size,
        probability_ch=probability_ch,
        max_rounds=max_rounds
    )

    leach_results = sim_leach.run()

    print(f"\n✓ LEACH Simulation Complete")
    print(f"  Events logged: {len(leach_results['events_log'])}")
    print(f"  Rounds executed: {leach_results['num_rounds']}")
    print(f"  Final metrics:")
    print(f"    FDN: {leach_results['fdn']}")
    print(f"    FMR: {leach_results['fmr']}")
    print(f"    DLBI: {leach_results['dlbi']:.4f}")
    print(f"    RSPI: {leach_results['rspi']:.4f}")