"""
Reliable MAC layer with ACK and retransmission.

Implements:
- Send queue management
- Random backoff before transmission
- ACK等待 and retransmission
- Maximum retry limit
"""

import random
from typing import Optional, Dict
from dataclasses import dataclass, field
from collections import deque

import simpy

from sim.core.packet import Packet, PacketType
from sim.radio import airtime as airtime_calc
from sim import config


@dataclass
class PendingAck:
    """Tracks a packet waiting for ACK."""

    packet: Packet
    dst: int
    retry_count: int = 0
    send_time: float = 0.0


class ReliableMAC:
    """
    Reliable MAC layer with CSMA-like backoff and ACK.

    Send flow:
    1. Enqueue packet
    2. Wait for random backoff
    3. Transmit
    4. Wait for ACK
    5. Retry or success
    """

    def __init__(self, env: simpy.Environment, node_id: int):
        """
        Initialize MAC layer.

        Args:
            env: SimPy environment
            node_id: This node's ID
        """
        self.env = env
        self.node_id = node_id

        # Send queue
        self.queue: deque = deque()

        # Pending ACKs {seq: PendingAck}
        self.pending_acks: Dict[int, PendingAck] = {}

        # Statistics
        self.sent_packets = 0
        self.received_acks = 0
        self.retries = 0

        # Channel access (set by node)
        self.channel = None

    def enqueue(self, packet: Packet, dst: int):
        """
        Add packet to send queue.

        Args:
            packet: Packet to send
            dst: Destination node ID
        """
        self.queue.append((packet, dst))

    def dequeue(self) -> Optional[tuple]:
        """
        Get next packet from queue.

        Returns:
            Tuple of (packet, dst) or None if queue empty
        """
        if self.queue:
            return self.queue.popleft()
        return None

    def has_pending(self) -> bool:
        """Check if there are packets to send."""
        return len(self.queue) > 0

    def calculate_backoff(self) -> float:
        """
        Calculate random backoff time.

        Returns:
            Backoff time in seconds
        """
        return random.uniform(config.BACKOFF_MIN, config.BACKOFF_MAX)

    def calculate_ack_timeout(self, packet: Packet) -> float:
        """
        Calculate ACK timeout based on packet airtime.

        Args:
            packet: The packet waiting for ACK

        Returns:
            Timeout in seconds
        """
        if packet.is_data:
            ack_time = airtime_calc.get_ack_airtime()
        else:
            ack_time = airtime_calc.get_hello_airtime()

        return ack_time * config.ACK_TIMEOUT_FACTOR

    def start_pending_ack(self, packet: Packet, dst: int):
        """
        Start tracking a packet waiting for ACK.

        Args:
            packet: The sent packet
            dst: Destination node ID
        """
        self.pending_acks[packet.seq] = PendingAck(
            packet=packet, dst=dst, retry_count=0, send_time=self.env.now
        )

    def ack_received(self, seq: int) -> bool:
        """
        Handle ACK received for a packet.

        Args:
            seq: Sequence number of acknowledged packet

        Returns:
            True if ACK was pending (success)
        """
        if seq in self.pending_acks:
            del self.pending_acks[seq]
            self.received_acks += 1
            return True
        return False

    def should_retry(self, seq: int) -> bool:
        """
        Check if a packet should be retried.

        Args:
            seq: Sequence number

        Returns:
            True if should retry
        """
        if seq not in self.pending_acks:
            return False

        pending = self.pending_acks[seq]
        if pending.retry_count >= config.MAX_RETRY:
            # Max retries reached, remove from pending
            del self.pending_acks[seq]
            return False

        return True

    def increment_retry(self, seq: int):
        """
        Increment retry count for a packet.

        Args:
            seq: Sequence number
        """
        if seq in self.pending_acks:
            self.pending_acks[seq].retry_count += 1
            self.retries += 1

    def get_retry_packet(self, seq: int) -> Optional[Packet]:
        """
        Get packet for retry.

        Args:
            seq: Sequence number

        Returns:
            Packet to retry, or None if max retries reached
        """
        if seq in self.pending_acks:
            pending = self.pending_acks[seq]
            if pending.retry_count < config.MAX_RETRY:
                return pending.packet
        return None

    def get_pending_count(self) -> int:
        """Get number of packets waiting for ACK."""
        return len(self.pending_acks)

    def get_queue_length(self) -> int:
        """Get send queue length."""
        return len(self.queue)

    def reset_stats(self):
        """Reset MAC statistics."""
        self.sent_packets = 0
        self.received_acks = 0
        self.retries = 0

    def record_send(self):
        """Record a packet send."""
        self.sent_packets += 1

    def get_stats(self) -> dict:
        """Get MAC statistics."""
        return {
            "sent_packets": self.sent_packets,
            "received_acks": self.received_acks,
            "retries": self.retries,
            "pending_acks": len(self.pending_acks),
            "queue_length": len(self.queue),
        }