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", )