"""Whole-frame video scheduling and predictive admission for Lab035."""

from __future__ import annotations

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

from protocol.link_packet import TrafficClass, encode_link_packet
from protocol.priority_scheduler import transmission_duration_seconds
from protocol.video_age_policy import AgePolicyPacket


TIME_EPSILON_SECONDS = 1e-12


class FramePolicy(str, Enum):
    NO_DROP = "no_drop"
    LATEST_ONLY = "latest_only"
    TWO_WAITING = "two_waiting"
    PREDICT_1000MS = "predict_1000ms"
    PREDICT_500MS = "predict_500ms"

    @property
    def deadline_seconds(self) -> float | None:
        if self is FramePolicy.PREDICT_1000MS:
            return 1.0
        if self is FramePolicy.PREDICT_500MS:
            return 0.5
        return None


@dataclass(frozen=True)
class VideoFrameGroup:
    composite_frame_id: int
    generation_time_us: int
    packets: tuple[AgePolicyPacket, ...]

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

    @property
    def wire_size_bytes(self) -> int:
        return sum(packet.wire_size_bytes for packet in self.packets)


@dataclass(frozen=True)
class FrameDrop:
    frame: VideoFrameGroup
    drop_time_seconds: float
    reason: str


@dataclass(frozen=True)
class StateReplacement:
    removed: AgePolicyPacket
    replacement: AgePolicyPacket
    time_seconds: float


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


@dataclass(frozen=True)
class FrameAdmission:
    composite_frame_id: int
    predicted_completion_seconds: float
    actual_completion_seconds: float

    @property
    def prediction_error_seconds(self) -> float:
        return self.actual_completion_seconds - self.predicted_completion_seconds


@dataclass(frozen=True)
class FrameScheduleResult:
    policy: FramePolicy
    channel_bitrate_bps: float
    transmitted: tuple[FrameScheduledPacket, ...]
    dropped_frames: tuple[FrameDrop, ...]
    replacements: tuple[StateReplacement, ...]
    started_frame_ids: tuple[int, ...]
    completed_frame_ids: tuple[int, ...]
    admissions: tuple[FrameAdmission, ...]


def _replace_state(
    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 _serve_high_priority(
    cursor: float,
    ready: list[AgePolicyPacket],
    future: tuple[AgePolicyPacket, ...],
    future_index: int,
    channel_bitrate_bps: float,
) -> tuple[float, int, list[AgePolicyPacket]]:
    """Pure local helper used only by the completion predictor."""

    copied_ready = list(ready)
    index = future_index
    while index < len(future) and future[index].available_time_seconds <= cursor + TIME_EPSILON_SECONDS:
        item = future[index]
        index += 1
        copied_ready = [
            old for old in copied_ready
            if not (
                old.packet.traffic_class == item.packet.traffic_class
                and old.packet.stream_id == item.packet.stream_id
                and item.packet.traffic_class in (TrafficClass.CONTROL, TrafficClass.TELEMETRY)
            )
        ]
        copied_ready.append(item)
    while copied_ready:
        selected = min(
            copied_ready,
            key=lambda item: (int(item.packet.traffic_class), item.arrival_order),
        )
        copied_ready.remove(selected)
        cursor += transmission_duration_seconds(selected.wire_size_bytes, channel_bitrate_bps)
        while index < len(future) and future[index].available_time_seconds <= cursor + TIME_EPSILON_SECONDS:
            item = future[index]
            index += 1
            copied_ready = [
                old for old in copied_ready
                if not (
                    old.packet.traffic_class == item.packet.traffic_class
                    and old.packet.stream_id == item.packet.stream_id
                    and item.packet.traffic_class in (TrafficClass.CONTROL, TrafficClass.TELEMETRY)
                )
            ]
            copied_ready.append(item)
    return cursor, index, copied_ready


def predict_frame_completion(
    current_time_seconds: float,
    frame: VideoFrameGroup,
    channel_bitrate_bps: float,
    ready_high_priority: Iterable[AgePolicyPacket],
    future_high_priority: Iterable[AgePolicyPacket],
    active_packet_remaining_seconds: float = 0.0,
) -> float:
    """Predict full-frame completion without mutating any caller collection.

    Known periodic commands, telemetry, and the scheduled emergency event are
    simulated exactly at packet boundaries. Unknown future discrete events are
    outside the Lab035 traffic model and therefore cannot be included.
    """

    if channel_bitrate_bps <= 0.0:
        raise ValueError("channel bitrate must be positive")
    if active_packet_remaining_seconds < 0.0:
        raise ValueError("active packet remainder must not be negative")
    cursor = current_time_seconds + active_packet_remaining_seconds
    ready = list(ready_high_priority)
    future = tuple(sorted(future_high_priority, key=lambda item: (item.available_time_seconds, item.arrival_order)))
    future_index = 0
    for video_packet in frame.packets:
        cursor, future_index, ready = _serve_high_priority(
            cursor, ready, future, future_index, channel_bitrate_bps
        )
        cursor += transmission_duration_seconds(video_packet.wire_size_bytes, channel_bitrate_bps)
        while future_index < len(future) and future[future_index].available_time_seconds <= cursor + TIME_EPSILON_SECONDS:
            ready.append(future[future_index])
            future_index += 1
    return cursor


def schedule_video_frames(
    frames: Iterable[VideoFrameGroup],
    high_priority_packets: Iterable[AgePolicyPacket],
    policy: FramePolicy,
    channel_bitrate_bps: float,
) -> FrameScheduleResult:
    """Run strict-priority service while keeping video frames contiguous."""

    policy = FramePolicy(policy)
    if channel_bitrate_bps <= 0.0:
        raise ValueError("channel bitrate must be positive")
    frame_arrivals = tuple(sorted(frames, key=lambda frame: (frame.generation_time_seconds, frame.composite_frame_id)))
    high_arrivals = tuple(sorted(high_priority_packets, key=lambda item: (item.available_time_seconds, item.arrival_order)))
    ready_high: list[AgePolicyPacket] = []
    pending_frames: list[VideoFrameGroup] = []
    transmitted: list[FrameScheduledPacket] = []
    dropped: list[FrameDrop] = []
    replacements: list[StateReplacement] = []
    started: list[int] = []
    completed: list[int] = []
    predicted_by_frame: dict[int, float] = {}
    actual_by_frame: dict[int, float] = {}
    active: VideoFrameGroup | None = None
    active_index = 0
    cursor = 0.0
    frame_index = high_index = 0
    blocker_by_order: dict[int, tuple[AgePolicyPacket, float]] = {}

    def drop_frame(frame: VideoFrameGroup, reason: str) -> None:
        dropped.append(FrameDrop(frame, cursor, reason))

    def admit(now: float, active_packet: AgePolicyPacket | None = None, active_end: float = 0.0) -> None:
        nonlocal frame_index, high_index
        while high_index < len(high_arrivals) and high_arrivals[high_index].available_time_seconds <= now + TIME_EPSILON_SECONDS:
            item = high_arrivals[high_index]
            high_index += 1
            if active_packet is not None and item.available_time_seconds > cursor + TIME_EPSILON_SECONDS:
                blocker_by_order[item.arrival_order] = (
                    active_packet,
                    max(0.0, active_end - item.available_time_seconds),
                )
            _replace_state(ready_high, item, replacements)
        while frame_index < len(frame_arrivals) and frame_arrivals[frame_index].generation_time_seconds <= now + TIME_EPSILON_SECONDS:
            frame = frame_arrivals[frame_index]
            frame_index += 1
            if policy is FramePolicy.LATEST_ONLY:
                for old in pending_frames:
                    drop_frame(old, "replaced_by_newest")
                pending_frames[:] = [frame]
            elif policy is FramePolicy.TWO_WAITING:
                pending_frames.append(frame)
                while len(pending_frames) > 2:
                    drop_frame(pending_frames.pop(0), "waiting_limit")
            else:
                pending_frames.append(frame)

    while (
        frame_index < len(frame_arrivals)
        or high_index < len(high_arrivals)
        or ready_high
        or pending_frames
        or active is not None
    ):
        if not ready_high and not pending_frames and active is None:
            next_times = []
            if frame_index < len(frame_arrivals):
                next_times.append(frame_arrivals[frame_index].generation_time_seconds)
            if high_index < len(high_arrivals):
                next_times.append(high_arrivals[high_index].available_time_seconds)
            cursor = max(cursor, min(next_times))
        admit(cursor)

        selected: AgePolicyPacket | None = None
        if ready_high:
            selected = min(
                ready_high,
                key=lambda item: (int(item.packet.traffic_class), item.arrival_order),
            )
            ready_high.remove(selected)
        else:
            if active is None:
                while pending_frames and active is None:
                    candidate = pending_frames.pop(0)
                    deadline = policy.deadline_seconds
                    if deadline is not None:
                        future_high = high_arrivals[high_index:]
                        predicted = predict_frame_completion(
                            cursor,
                            candidate,
                            channel_bitrate_bps,
                            tuple(ready_high),
                            future_high,
                        )
                        if predicted - candidate.generation_time_seconds > deadline + TIME_EPSILON_SECONDS:
                            drop_frame(candidate, "prediction_reject")
                            continue
                        predicted_by_frame[candidate.composite_frame_id] = predicted
                    active = candidate
                    active_index = 0
                    started.append(active.composite_frame_id)
            if active is not None:
                selected = active.packets[active_index]

        if selected is None:
            continue
        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(end, selected, end)
        blocker, blocking_delay = blocker_by_order.get(selected.arrival_order, (None, 0.0))
        transmitted.append(
            FrameScheduledPacket(selected, start, end, blocker, blocking_delay, wire_packet)
        )
        cursor = end
        if selected.packet.traffic_class is TrafficClass.VIDEO:
            assert active is not None
            active_index += 1
            if active_index == len(active.packets):
                completed.append(active.composite_frame_id)
                actual_by_frame[active.composite_frame_id] = end
                active = None
                active_index = 0

    admissions = tuple(
        FrameAdmission(frame_id, predicted, actual_by_frame[frame_id])
        for frame_id, predicted in sorted(predicted_by_frame.items())
    )
    return FrameScheduleResult(
        policy=policy,
        channel_bitrate_bps=channel_bitrate_bps,
        transmitted=tuple(transmitted),
        dropped_frames=tuple(dropped),
        replacements=tuple(replacements),
        started_frame_ids=tuple(started),
        completed_frame_ids=tuple(completed),
        admissions=admissions,
    )