SDR-Rover/protocol/priority_scheduler.py

"""Deterministic non-preemptive queue schedulers for Lab033."""

from __future__ import annotations

from dataclasses import dataclass
from enum import Enum
from typing import Iterable

from protocol.link_packet import (
    HEADER_SIZE,
    LinkPacket,
    TrafficClass,
    encode_link_packet,
)


TIME_EPSILON_SECONDS = 1e-12


class SchedulerMode(str, Enum):
    FIFO = "fifo"
    STRICT_PRIORITY = "strict_priority"
    LATEST_STATE = "latest_state"


@dataclass(frozen=True)
class QueuedPacket:
    """A packet plus a stable global arrival-order tie breaker."""

    packet: LinkPacket
    arrival_order: int

    @property
    def generation_time_seconds(self) -> float:
        return self.packet.generation_time_us / 1_000_000.0

    @property
    def wire_size_bytes(self) -> int:
        return HEADER_SIZE + len(self.packet.payload)


@dataclass(frozen=True)
class Replacement:
    removed: QueuedPacket
    replacement: QueuedPacket


@dataclass(frozen=True)
class ScheduledPacket:
    queued: QueuedPacket
    start_seconds: float
    end_seconds: float
    blocked_by: QueuedPacket | None
    blocking_delay_seconds: float
    wire_packet: bytes


@dataclass(frozen=True)
class QueueSample:
    time_seconds: float
    packets: int
    bytes: int


@dataclass(frozen=True)
class ScheduleResult:
    mode: SchedulerMode
    channel_bitrate_bps: float
    transmitted: tuple[ScheduledPacket, ...]
    replacements: tuple[Replacement, ...]
    queue_samples: tuple[QueueSample, ...]


def transmission_duration_seconds(
    packet_size_bytes: int,
    channel_bitrate_bps: float,
) -> float:
    """Return exact serialization duration for a complete packet."""

    if packet_size_bytes < 0:
        raise ValueError("packet size must not be negative")
    if channel_bitrate_bps <= 0.0:
        raise ValueError("channel bitrate must be positive")
    return packet_size_bytes * 8.0 / channel_bitrate_bps


def _selection_key(item: QueuedPacket, mode: SchedulerMode) -> tuple[int, int]:
    if mode is SchedulerMode.FIFO:
        return (0, item.arrival_order)
    return (int(item.packet.traffic_class), item.arrival_order)


def _enqueue(
    ready: list[QueuedPacket],
    item: QueuedPacket,
    mode: SchedulerMode,
    replacements: list[Replacement],
) -> None:
    if (
        mode is SchedulerMode.LATEST_STATE
        and item.packet.traffic_class
        in (TrafficClass.CONTROL, TrafficClass.TELEMETRY)
    ):
        retained = []
        for old in ready:
            if (
                old.packet.traffic_class == item.packet.traffic_class
                and old.packet.stream_id == item.packet.stream_id
            ):
                replacements.append(Replacement(old, item))
            else:
                retained.append(old)
        ready[:] = retained
    ready.append(item)


def schedule_packets(
    packets: Iterable[QueuedPacket],
    mode: SchedulerMode,
    channel_bitrate_bps: float,
) -> ScheduleResult:
    """Serialize all packets through one non-preemptive shared resource.

    The caller supplies a unique ``arrival_order``. Packets are admitted by
    ``(generation_time, arrival_order)``. FIFO selects the smallest arrival
    order; priority modes select the smallest ``TrafficClass`` and preserve
    arrival order within a class. Latest-state replacement only touches items
    still waiting in ``ready`` and therefore can never remove an active packet.
    """

    try:
        mode = SchedulerMode(mode)
    except ValueError as error:
        raise ValueError("unsupported scheduler mode") from error
    if channel_bitrate_bps <= 0.0:
        raise ValueError("channel bitrate must be positive")
    arrivals = sorted(
        tuple(packets),
        key=lambda item: (
            item.generation_time_seconds,
            item.arrival_order,
        ),
    )
    if len({item.arrival_order for item in arrivals}) != len(arrivals):
        raise ValueError("arrival_order values must be unique")

    ready: list[QueuedPacket] = []
    transmitted: list[ScheduledPacket] = []
    replacements: list[Replacement] = []
    samples: list[QueueSample] = []
    blocker_by_order: dict[int, tuple[QueuedPacket, float]] = {}
    cursor = 0.0
    arrival_index = 0

    def sample(time_seconds: float) -> None:
        samples.append(
            QueueSample(
                time_seconds=time_seconds,
                packets=len(ready),
                bytes=sum(item.wire_size_bytes for item in ready),
            )
        )

    while arrival_index < len(arrivals) or ready:
        if not ready:
            cursor = max(cursor, arrivals[arrival_index].generation_time_seconds)
        while (
            arrival_index < len(arrivals)
            and arrivals[arrival_index].generation_time_seconds
            <= cursor + TIME_EPSILON_SECONDS
        ):
            _enqueue(ready, arrivals[arrival_index], mode, replacements)
            arrival_index += 1
        sample(cursor)
        selected_index = min(
            range(len(ready)),
            key=lambda index: _selection_key(ready[index], mode),
        )
        selected = ready.pop(selected_index)
        start = max(cursor, selected.generation_time_seconds)
        wire_packet = encode_link_packet(selected.packet)
        end = start + transmission_duration_seconds(
            len(wire_packet), channel_bitrate_bps
        )

        while (
            arrival_index < len(arrivals)
            and arrivals[arrival_index].generation_time_seconds
            <= end + TIME_EPSILON_SECONDS
        ):
            arriving = arrivals[arrival_index]
            if arriving.generation_time_seconds > start + TIME_EPSILON_SECONDS:
                blocker_by_order[arriving.arrival_order] = (
                    selected,
                    max(0.0, end - arriving.generation_time_seconds),
                )
            _enqueue(ready, arriving, mode, replacements)
            arrival_index += 1

        blocker, blocking_delay = blocker_by_order.get(
            selected.arrival_order, (None, 0.0)
        )
        transmitted.append(
            ScheduledPacket(
                queued=selected,
                start_seconds=start,
                end_seconds=end,
                blocked_by=blocker,
                blocking_delay_seconds=blocking_delay,
                wire_packet=wire_packet,
            )
        )
        cursor = end
        sample(cursor)

    return ScheduleResult(
        mode=mode,
        channel_bitrate_bps=channel_bitrate_bps,
        transmitted=tuple(transmitted),
        replacements=tuple(replacements),
        queue_samples=tuple(samples),
    )