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.device_plugins.xgrids_k1.mqtt import CaptureFormatError, iter_capture_frames
from k1link.device_plugins.xgrids_k1.mqtt.capture import RAW_MAGIC
from k1link.device_plugins.xgrids_k1.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",
            )
