"""Frame-aware stale-video policy for the Lab034 shared link queue.

The scheduler remains strict-priority, state-replacing, and non-preemptive.
Unlike Lab033, packet availability is separate from the timestamp stored in
the common header.  Every aligned video packet therefore carries the source
composite-frame generation time while FEC construction metadata remains local
to the scheduler.
"""

from __future__ import annotations

from dataclasses import dataclass
from typing import Iterable

from protocol.link_packet import HEADER_SIZE, LinkPacket, TrafficClass, encode_link_packet
from protocol.priority_scheduler import transmission_duration_seconds


TIME_EPSILON_SECONDS = 1e-12


@dataclass(frozen=True)
class AgePolicyPacket:
    """A common packet plus non-serialized availability/frame metadata."""

    packet: LinkPacket
    arrival_order: int
    available_time_us: int
    composite_frame_id: int | None = None

    @property
    def available_time_seconds(self) -> float:
        return self.available_time_us / 1_000_000.0

    @property
    def frame_generation_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 StateReplacement:
    removed: AgePolicyPacket
    replacement: AgePolicyPacket
    time_seconds: float


@dataclass(frozen=True)
class DroppedVideoPacket:
    item: AgePolicyPacket
    drop_time_seconds: float


@dataclass(frozen=True)
class AgeScheduledPacket:
    item: AgePolicyPacket
    start_seconds: float
    end_seconds: float
    blocked_by: AgePolicyPacket | None
    blocking_delay_seconds: float
    wire_packet: bytes


@dataclass(frozen=True)
class AgeScheduleResult:
    channel_bitrate_bps: float
    max_video_age_seconds: float | None
    transmitted: tuple[AgeScheduledPacket, ...]
    dropped_video: tuple[DroppedVideoPacket, ...]
    replacements: tuple[StateReplacement, ...]


def _enqueue(
    ready: list[AgePolicyPacket],
    item: AgePolicyPacket,
    replacements: list[StateReplacement],
) -> None:
    if 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(
                    StateReplacement(old, item, item.available_time_seconds)
                )
            else:
                retained.append(old)
        ready[:] = retained
    ready.append(item)


def schedule_with_video_age(
    packets: Iterable[AgePolicyPacket],
    channel_bitrate_bps: float,
    max_video_age_seconds: float | None,
) -> AgeScheduleResult:
    """Run a strict-priority non-preemptive scheduler with frame drops.

    At every packet-selection instant all queued packets belonging to a video
    frame older than ``max_video_age_seconds`` are removed together.  An active
    packet has already left the queue and is never interrupted.
    """

    if channel_bitrate_bps <= 0.0:
        raise ValueError("channel bitrate must be positive")
    if max_video_age_seconds is not None and max_video_age_seconds <= 0.0:
        raise ValueError("maximum video age must be positive")
    arrivals = sorted(
        tuple(packets),
        key=lambda item: (item.available_time_seconds, item.arrival_order),
    )
    if len({item.arrival_order for item in arrivals}) != len(arrivals):
        raise ValueError("arrival_order values must be unique")
    for item in arrivals:
        if item.packet.traffic_class is TrafficClass.VIDEO:
            if item.composite_frame_id is None:
                raise ValueError("video packet lacks composite_frame_id")
        elif item.composite_frame_id is not None:
            raise ValueError("non-video packet has composite_frame_id")

    ready: list[AgePolicyPacket] = []
    transmitted: list[AgeScheduledPacket] = []
    dropped: list[DroppedVideoPacket] = []
    replacements: list[StateReplacement] = []
    blocker_by_order: dict[int, tuple[AgePolicyPacket, float]] = {}
    dropped_frames: set[int] = set()
    arrival_index = 0
    cursor = 0.0

    def admit_until(limit: float, active: AgePolicyPacket | None = None, active_end: float = 0.0) -> None:
        nonlocal arrival_index
        while (
            arrival_index < len(arrivals)
            and arrivals[arrival_index].available_time_seconds
            <= limit + TIME_EPSILON_SECONDS
        ):
            item = arrivals[arrival_index]
            arrival_index += 1
            if (
                active is not None
                and item.available_time_seconds > cursor + TIME_EPSILON_SECONDS
            ):
                blocker_by_order[item.arrival_order] = (
                    active,
                    max(0.0, active_end - item.available_time_seconds),
                )
            if item.composite_frame_id in dropped_frames:
                dropped.append(DroppedVideoPacket(item, limit))
            else:
                _enqueue(ready, item, replacements)

    def expire_video(now: float) -> None:
        if max_video_age_seconds is None:
            return
        expired_frames = {
            item.composite_frame_id
            for item in ready
            if (
                item.packet.traffic_class is TrafficClass.VIDEO
                and now - item.frame_generation_seconds
                > max_video_age_seconds + TIME_EPSILON_SECONDS
            )
        }
        if not expired_frames:
            return
        retained = []
        for item in ready:
            if item.composite_frame_id in expired_frames:
                dropped.append(DroppedVideoPacket(item, now))
                assert item.composite_frame_id is not None
                dropped_frames.add(item.composite_frame_id)
            else:
                retained.append(item)
        ready[:] = retained

    while arrival_index < len(arrivals) or ready:
        if not ready:
            cursor = max(cursor, arrivals[arrival_index].available_time_seconds)
        admit_until(cursor)
        expire_video(cursor)
        if not ready:
            continue
        selected_index = min(
            range(len(ready)),
            key=lambda index: (
                int(ready[index].packet.traffic_class),
                ready[index].arrival_order,
            ),
        )
        selected = ready.pop(selected_index)
        start = max(cursor, selected.available_time_seconds)
        wire_packet = encode_link_packet(selected.packet)
        end = start + transmission_duration_seconds(
            len(wire_packet), channel_bitrate_bps
        )
        cursor = start
        admit_until(end, selected, end)
        blocker, blocking_delay = blocker_by_order.get(
            selected.arrival_order, (None, 0.0)
        )
        transmitted.append(
            AgeScheduledPacket(
                item=selected,
                start_seconds=start,
                end_seconds=end,
                blocked_by=blocker,
                blocking_delay_seconds=blocking_delay,
                wire_packet=wire_packet,
            )
        )
        cursor = end

    return AgeScheduleResult(
        channel_bitrate_bps=channel_bitrate_bps,
        max_video_age_seconds=max_video_age_seconds,
        transmitted=tuple(transmitted),
        dropped_video=tuple(dropped),
        replacements=tuple(replacements),
    )