NODEDC_MISSION_CORE/src/k1link/viewer/replay.py

from __future__ import annotations

import json
import time
from collections.abc import Generator, Iterator
from decimal import Decimal, InvalidOperation
from pathlib import Path
from typing import IO, Literal

from k1link.mqtt import CaptureFormatError, iter_capture_frames
from k1link.mqtt.capture import RAW_MAGIC
from k1link.viewer.messages import StreamMessage

MAX_REPLAY_PAYLOAD_BYTES = 2 * 1024 * 1024
MAX_LEGACY_LINE_BYTES = MAX_REPLAY_PAYLOAD_BYTES * 2 + 64 * 1024
MAX_METADATA_LINE_CHARS = 1024 * 1024
ReplayFormat = Literal["k1mqtt", "legacy_tsv"]


class ReplayFormatError(ValueError):
    """A replay input is corrupt or outside the reviewed bounds."""


def detect_replay_format(path: Path) -> ReplayFormat:
    resolved = path.expanduser()
    with resolved.open("rb") as stream:
        prefix = stream.read(len(RAW_MAGIC))
    if prefix == RAW_MAGIC:
        return "k1mqtt"
    if resolved.suffix.casefold() == ".tsv":
        return "legacy_tsv"
    raise ReplayFormatError("input is neither a K1MQTT capture nor the reviewed legacy TSV")


def iter_replay_messages(path: Path) -> Generator[StreamMessage, None, None]:
    """Yield bounded messages from native evidence or the one reviewed TSV export."""
    resolved = path.expanduser().resolve()
    replay_format = detect_replay_format(resolved)
    if replay_format == "legacy_tsv":
        yield from _iter_legacy_tsv(resolved)
        return
    yield from _iter_native_capture(resolved)


def _iter_native_capture(path: Path) -> Iterator[StreamMessage]:
    metadata_path = path.with_name("mqtt.metadata.jsonl")
    metadata_stream: IO[str] | None = None
    if metadata_path.is_file():
        metadata_stream = metadata_path.open("r", encoding="utf-8")
    fallback_epoch_ns = time.time_ns()
    try:
        for frame in iter_capture_frames(path, max_payload_bytes=MAX_REPLAY_PAYLOAD_BYTES):
            epoch_ns = fallback_epoch_ns + frame.sequence - 1
            monotonic_ns: int | None = None
            if metadata_stream is not None:
                timing = _read_native_timing(
                    metadata_stream,
                    expected_sequence=frame.sequence,
                    fallback_epoch_ns=epoch_ns,
                )
                if timing is None:
                    # A crash may leave one raw frame ahead of the last fully
                    # committed metadata line. Only the aligned prefix has a
                    # trustworthy source timeline and is safe to replay.
                    return
                epoch_ns, monotonic_ns = timing
            yield StreamMessage(
                sequence=frame.sequence,
                topic=frame.topic,
                payload=frame.payload,
                received_at_epoch_ns=epoch_ns,
                received_monotonic_ns=monotonic_ns,
                source="k1mqtt",
            )
    except CaptureFormatError as exc:
        raise ReplayFormatError(str(exc)) from exc
    finally:
        if metadata_stream is not None:
            metadata_stream.close()


def _read_native_timing(
    stream: IO[str],
    *,
    expected_sequence: int,
    fallback_epoch_ns: int,
) -> tuple[int, int | None] | None:
    line = stream.readline(MAX_METADATA_LINE_CHARS + 1)
    if not line:
        return None
    if len(line) > MAX_METADATA_LINE_CHARS:
        raise ReplayFormatError("native metadata line exceeds the reviewed bound")
    try:
        record = json.loads(line)
    except json.JSONDecodeError as exc:
        if not line.endswith(("\n", "\r")):
            # A non-newline final tail is the only tolerated corruption: the
            # writer may have crashed between raw and metadata group commits.
            return None
        raise ReplayFormatError(
            f"native metadata line {expected_sequence} is not valid JSON"
        ) from exc
    if record.get("record_type") != "message" or record.get("sequence") != expected_sequence:
        raise ReplayFormatError(f"native metadata is not aligned at message {expected_sequence}")
    epoch_ns = record.get("received_at_epoch_ns", fallback_epoch_ns)
    monotonic_ns = record.get("received_monotonic_ns")
    if not isinstance(epoch_ns, int) or epoch_ns < 0:
        raise ReplayFormatError("native metadata received_at_epoch_ns is invalid")
    if monotonic_ns is not None and (not isinstance(monotonic_ns, int) or monotonic_ns < 0):
        raise ReplayFormatError("native metadata received_monotonic_ns is invalid")
    return epoch_ns, monotonic_ns


def _iter_legacy_tsv(path: Path) -> Iterator[StreamMessage]:
    with path.open("rb") as stream:
        line_number = 0
        while True:
            raw_line = stream.readline(MAX_LEGACY_LINE_BYTES + 1)
            if not raw_line:
                return
            line_number += 1
            if len(raw_line) > MAX_LEGACY_LINE_BYTES:
                raise ReplayFormatError(
                    f"legacy TSV line {line_number} exceeds {MAX_LEGACY_LINE_BYTES} bytes"
                )
            stripped = raw_line.rstrip(b"\r\n")
            if not stripped:
                continue
            columns = stripped.split(b"\t")
            if len(columns) != 4:
                raise ReplayFormatError(
                    f"legacy TSV line {line_number} must contain exactly four columns"
                )
            timestamp_raw, topic_raw, length_raw, payload_hex = columns
            try:
                timestamp = Decimal(timestamp_raw.decode("ascii"))
                declared_length = int(length_raw.decode("ascii"), 10)
                topic = topic_raw.decode("utf-8")
            except (InvalidOperation, UnicodeDecodeError, ValueError) as exc:
                raise ReplayFormatError(
                    f"legacy TSV line {line_number} has invalid timestamp/topic/length"
                ) from exc
            if not timestamp.is_finite() or timestamp < 0:
                raise ReplayFormatError(
                    f"legacy TSV line {line_number} timestamp is outside bounds"
                )
            if not topic:
                raise ReplayFormatError(f"legacy TSV line {line_number} has an empty topic")
            if declared_length < 0 or declared_length > MAX_REPLAY_PAYLOAD_BYTES:
                raise ReplayFormatError(
                    f"legacy TSV line {line_number} payload length is outside bounds"
                )
            if len(payload_hex) != declared_length * 2:
                raise ReplayFormatError(
                    f"legacy TSV line {line_number} declared payload length does not match hex"
                )
            try:
                payload = bytes.fromhex(payload_hex.decode("ascii"))
            except (UnicodeDecodeError, ValueError) as exc:
                raise ReplayFormatError(
                    f"legacy TSV line {line_number} payload is not valid hex"
                ) from exc
            epoch_ns = int(timestamp * Decimal(1_000_000_000))
            yield StreamMessage(
                sequence=line_number,
                topic=topic,
                payload=payload,
                received_at_epoch_ns=epoch_ns,
                source="legacy_tsv",
            )