"""
Gradient-based routing protocol.

Implements:
- Cost-based routing (gradient routing)
- HELLO message handling for neighbor discovery
- Parent selection based on cost + link penalty
- Data forwarding to parent node
"""

from typing import Dict, Optional
from dataclasses import dataclass, field

from sim.core.packet import Packet, PacketType
from sim.radio import propagation
from sim import config


@dataclass
class NeighborInfo:
    """Information about a neighbor node."""

    node_id: int
    cost: int
    rssi: float
    last_hello_time: float


class GradientRouting:
    """
    Gradient routing protocol.

    Each node maintains:
    - cost: Distance to sink (in hops + penalty)
    - parent: Next hop towards sink
    - neighbors: Dict of known neighbors with their costs
    """

    def __init__(self, node_id: int, is_sink: bool = False):
        """
        Initialize routing.

        Args:
            node_id: This node's ID
            is_sink: Whether this node is the sink
        """
        self.node_id = node_id
        self.is_sink = is_sink

        # Routing state
        self.cost = 0 if is_sink else float("inf")
        self.parent: Optional[int] = None
        self.neighbors: Dict[int, NeighborInfo] = {}

        # Sequence number for HELLO messages
        self.hello_seq = 0

    def reset(self):
        """Reset routing state."""
        self.cost = 0 if self.is_sink else float("inf")
        self.parent = None
        self.neighbors.clear()
        self.hello_seq = 0

    def create_hello_packet(self) -> Packet:
        """
        Create a HELLO packet for neighbor discovery.

        Returns:
            HELLO packet with current cost
        """
        packet = Packet(
            type=PacketType.HELLO,
            src=self.node_id,
            dst=-1,  # Broadcast
            seq=self.hello_seq,
            hop=0,
            payload=str(int(self.cost)) if self.cost != float("inf") else "inf",
        )
        self.hello_seq += 1
        return packet

    def process_hello(self, packet: Packet, rssi: float) -> bool:
        """
        Process received HELLO packet.

        Args:
            packet: Received HELLO packet
            rssi: RSSI of received signal

        Returns:
            True if routing state changed (cost/parent updated)
        """
        # Parse cost from payload
        try:
            neighbor_cost = int(packet.payload) if packet.payload else 0
        except ValueError:
            neighbor_cost = 0

        # Calculate link penalty based on RSSI
        link_penalty = propagation.calculate_link_penalty(rssi)

        # Calculate new cost to sink through this neighbor
        new_cost = neighbor_cost + 1 + int(link_penalty)

        # Update neighbor info
        old_neighbor = self.neighbors.get(packet.src)
        self.neighbors[packet.src] = NeighborInfo(
            node_id=packet.src,
            cost=neighbor_cost,
            rssi=rssi,
            last_hello_time=rssi,  # Use rssi field to store time
        )

        # Check if we should update our route
        # Update condition: new_cost < cost - 1
        old_cost = self.cost
        if new_cost < self.cost - config.ROUTE_UPDATE_THRESHOLD:
            self.cost = new_cost
            self.parent = packet.src
            return True

        # Also update if we have no route yet
        if self.parent is None and not self.is_sink:
            if new_cost < float("inf"):
                self.cost = new_cost
                self.parent = packet.src
                return True

        return old_cost != self.cost

    def get_next_hop(self, packet: Packet = None) -> Optional[int]:
        """
        Get next hop towards sink.

        Args:
            packet: Optional packet (for compatibility)

        Returns:
            Parent node ID, or None if no route
        """
        return self.parent

    def is_route_valid(self) -> bool:
        """Check if current route is valid."""
        if self.is_sink:
            return True
        return self.parent is not None and self.cost < float("inf")

    def cleanup_stale_neighbors(self, current_time: float, timeout: float = 30.0):
        """Remove neighbors that haven't sent HELLO recently."""
        stale = [
            nid
            for nid, info in self.neighbors.items()
            if current_time - info.last_hello_time > timeout
        ]
        for nid in stale:
            del self.neighbors[nid]

        # If our parent is stale, we need to find a new one
        if self.parent in stale:
            self.parent = None
            self.cost = float("inf")
            # Try to find new parent
            for nid, info in self.neighbors.items():
                if info.cost < self.cost:
                    self.cost = info.cost + 1
                    self.parent = nid

    def get_routing_table(self) -> dict:
        """Get routing table for debugging/visualization."""
        return {
            "node_id": self.node_id,
            "is_sink": self.is_sink,
            "cost": int(self.cost) if self.cost != float("inf") else -1,
            "parent": self.parent,
            "neighbors": {
                nid: {"cost": info.cost, "rssi": round(info.rssi, 2)}
                for nid, info in self.neighbors.items()
            },
        }