"""Scheduling metadata and receiver de-duplication for command copies.

Lab037 deliberately keeps the Lab033 :class:`~protocol.link_packet.LinkPacket`
wire representation unchanged.  A repeated command is therefore the exact
same encoded packet; only ``available_time_us`` and ``copy_index`` live in the
local scheduler metadata defined here.
"""

from __future__ import annotations

from dataclasses import dataclass
from enum import IntEnum
from typing import Iterable

from protocol.link_packet import LinkPacket, TrafficClass


class RepetitionMode(IntEnum):
    """Command protection alternatives compared by Lab037."""

    NONE = 1
    EMERGENCY_ONLY = 2
    EMERGENCY_AND_CONTROL = 3


@dataclass(frozen=True)
class CommandCopy:
    """One scheduling instance of an otherwise unchanged link packet."""

    packet: LinkPacket
    available_time_us: int
    copy_index: int
    arrival_order: int

    @property
    def is_repeat(self) -> bool:
        return self.copy_index > 0

    @property
    def command_key(self) -> tuple[int, int]:
        return self.packet.stream_id, self.packet.sequence_number


def copy_offsets_ms(
    traffic_class: TrafficClass,
    mode: RepetitionMode,
    control_repeat_delay_ms: float = 22.5,
) -> tuple[float, ...]:
    """Return copy offsets for a command class and protection mode."""

    traffic_class = TrafficClass(traffic_class)
    mode = RepetitionMode(mode)
    if traffic_class is TrafficClass.EMERGENCY:
        return (0.0, 15.0, 30.0) if mode >= RepetitionMode.EMERGENCY_ONLY else (0.0,)
    if traffic_class is TrafficClass.CONTROL:
        return (0.0, control_repeat_delay_ms) if mode is RepetitionMode.EMERGENCY_AND_CONTROL else (0.0,)
    raise ValueError("repetition is defined only for emergency and control commands")


def build_command_copies(
    packets: Iterable[LinkPacket],
    mode: RepetitionMode,
    *,
    first_arrival_order: int = 0,
    control_repeat_delay_ms: float = 22.5,
) -> tuple[CommandCopy, ...]:
    """Expand original commands into stable, time-ordered scheduler copies."""

    copies: list[CommandCopy] = []
    order = first_arrival_order
    for packet in packets:
        offsets = copy_offsets_ms(packet.traffic_class, mode, control_repeat_delay_ms)
        for copy_index, offset_ms in enumerate(offsets):
            copies.append(
                CommandCopy(
                    packet=packet,
                    available_time_us=packet.generation_time_us + int(round(offset_ms * 1000.0)),
                    copy_index=copy_index,
                    arrival_order=order,
                )
            )
            order += 1
    return tuple(
        sorted(
            copies,
            key=lambda item: (
                item.available_time_us,
                int(item.packet.traffic_class),
                item.arrival_order,
            ),
        )
    )


def cancel_superseded_control_copies(
    ready: Iterable[CommandCopy],
    newest: CommandCopy,
) -> tuple[tuple[CommandCopy, ...], tuple[CommandCopy, ...]]:
    """Remove only queued repeats belonging to older state commands.

    A packet already selected for transmission is intentionally outside this
    function, which makes packet serialization non-preemptive by construction.
    Primary copies and emergency commands are never removed here.
    """

    if newest.packet.traffic_class is not TrafficClass.CONTROL:
        return tuple(ready), ()
    retained: list[CommandCopy] = []
    cancelled: list[CommandCopy] = []
    for item in ready:
        if (
            item.packet.traffic_class is TrafficClass.CONTROL
            and item.packet.stream_id == newest.packet.stream_id
            and item.packet.sequence_number < newest.packet.sequence_number
            and item.is_repeat
        ):
            cancelled.append(item)
        else:
            retained.append(item)
    return tuple(retained), tuple(cancelled)


class CommandReceiver:
    """Apply only strictly newer sequence numbers independently per stream."""

    def __init__(self) -> None:
        self._last_sequence: dict[int, int] = {}
        self.accepted = 0
        self.suppressed = 0

    def accept(self, packet: LinkPacket) -> bool:
        """Return ``True`` only when this command advances its stream state."""

        if packet.traffic_class not in (TrafficClass.EMERGENCY, TrafficClass.CONTROL):
            raise ValueError("CommandReceiver accepts command traffic only")
        previous = self._last_sequence.get(packet.stream_id, -1)
        if packet.sequence_number <= previous:
            self.suppressed += 1
            return False
        self._last_sequence[packet.stream_id] = packet.sequence_number
        self.accepted += 1
        return True

    def last_sequence(self, stream_id: int) -> int | None:
        return self._last_sequence.get(stream_id)