"""Raw producer that replays a recorded snapshot dataset (no hardware).

Selected by the supervisor instead of the model's normal acquisition producer
when ``radar.replay.dataset_dir`` is non-empty. The dataset is a directory of
snapshot folders named ``NNNN_idK_ns<timestamp>`` (the VNA_AlGo
preprocessed-capture layout), each holding ``i{input}_o{output}_s21.npy`` and
``i{input}_o{output}_freq.npy`` files for every run combo.

The recorded data is already background-subtracted, so the profile must pair
this producer with NEUTRAL preprocess sets (calibration = 1, reference = 0,
no notch) — then the C++ preprocessor passes the traces through bit-exactly
and the processors (e.g. sar3d) see the recorded scene.

Timeline: snapshot timestamps are re-anchored to this process's monotonic
clock, preserving the recorded inter-frame gaps. On loop, the seam advances by
the median inter-frame gap, so motion-model consumers (sar3d aperture offsets)
never see time jump backwards. ``collection_id`` increases forever.
"""

from __future__ import annotations

import argparse
import json
import logging
from pathlib import Path
import re
import signal
import statistics
import threading
import time

import numpy as np

from python_app.logging_setup import coerce_level
from python_app.models.dataset_model import ComboKey, SweepCollection, TraceData
from python_app.models.run_config_model import RunConfigModel
from python_app.orchestration.shm import ShmRingWriter
from python_app.storage.npz.serialize import RAW_MAGIC, serialize_trace_collection

logger = logging.getLogger(__name__)

_SNAPSHOT_NAME = re.compile(r"^\d+_id\d+_ns(\d+)$")


class ReplayDataset:
    """Recorded snapshots preloaded into memory, with their recorded timeline."""

    def __init__(self, snapshots: list[list[TraceData]], timestamps_ns: list[int]) -> None:
        self.snapshots = snapshots
        self.timestamps_ns = timestamps_ns
        gaps = [b - a for a, b in zip(timestamps_ns, timestamps_ns[1:])]
        # Median gap doubles as the loop-seam step and the capture-window span;
        # a single-snapshot dataset falls back to an arbitrary sane cadence.
        self.median_gap_ns = int(statistics.median(gaps)) if gaps else 160_000_000

    def __len__(self) -> int:
        return len(self.snapshots)


def _load_dataset(dataset_dir: Path, combos: list[ComboKey], expected_points: int) -> ReplayDataset:
    """Preload every snapshot's traces; fail loudly on a malformed dataset."""
    if not dataset_dir.is_dir():
        raise RuntimeError(f"radar.replay.dataset_dir is not a directory: {dataset_dir}")

    named: list[tuple[str, int]] = []
    for entry in sorted(dataset_dir.iterdir()):
        match = _SNAPSHOT_NAME.match(entry.name)
        if entry.is_dir() and match:
            named.append((entry.name, int(match.group(1))))
    if not named:
        raise RuntimeError(
            f"No snapshot folders (NNNN_idK_ns<ts>) found in {dataset_dir}"
        )

    snapshots: list[list[TraceData]] = []
    timestamps_ns: list[int] = []
    for name, timestamp_ns in named:
        folder = dataset_dir / name
        # Recorded per-trace capture windows from the snapshot's meta.json: the
        # sar3d segmented timing model reconstructs the acquisition segments
        # from them. Missing/partial meta leaves the zero "unmeasured" sentinel
        # and the consumer falls back to its uniform-segment model.
        windows: dict[tuple[int, int], tuple[int, int]] = {}
        meta_path = folder / "meta.json"
        if meta_path.is_file():
            try:
                meta = json.loads(meta_path.read_text())
                for entry in meta.get("traces", []):
                    windows[(int(entry["input"]), int(entry["output"]))] = (
                        int(entry["capture_start_ns"]),
                        int(entry["capture_end_ns"]),
                    )
            except (OSError, ValueError, KeyError, TypeError):
                windows = {}
        traces: list[TraceData] = []
        for combo in combos:
            s21_path = folder / f"i{combo.input}_o{combo.output}_s21.npy"
            freq_path = folder / f"i{combo.input}_o{combo.output}_freq.npy"
            if not s21_path.is_file() or not freq_path.is_file():
                raise RuntimeError(
                    f"Snapshot {name} is missing files for combo "
                    f"input={combo.input} output={combo.output}: expected "
                    f"{s21_path.name} and {freq_path.name}. Check run.combos vs the dataset."
                )
            s21 = np.load(s21_path).astype(np.complex64, copy=False)
            freq = np.load(freq_path).astype(np.float32, copy=False)
            if s21.size != freq.size:
                raise RuntimeError(f"{s21_path}: s21 size {s21.size} != freq size {freq.size}")
            if expected_points and s21.size != expected_points:
                raise RuntimeError(
                    f"{s21_path}: {s21.size} points, but radar.sweep.points = "
                    f"{expected_points}. Align the profile's sweep with the dataset."
                )
            window = windows.get((combo.input, combo.output), (0, 0))
            traces.append(
                TraceData(
                    combo=combo,
                    frequency_hz=freq,
                    s11=np.zeros_like(s21),
                    s21=s21,
                    capture_start_ns=window[0],
                    capture_end_ns=window[1],
                )
            )
        snapshots.append(traces)
        timestamps_ns.append(timestamp_ns)

    return ReplayDataset(snapshots, timestamps_ns)


def _build_collection(
    dataset: ReplayDataset,
    index: int,
    collection_id: int,
    frame_ns: int,
) -> SweepCollection:
    """Wrap one preloaded snapshot in a SweepCollection at the given frame time.

    Recorded per-trace capture windows are re-anchored by the same shift as the
    frame time, preserving the real intra-frame segment layout. Traces without
    recorded windows get a synthetic window spanning the median frame gap.
    """
    span_ns = dataset.median_gap_ns
    start_ns = frame_ns - span_ns // 2
    end_ns = frame_ns + span_ns // 2
    shift_ns = frame_ns - dataset.timestamps_ns[index]
    traces = []
    for source in dataset.snapshots[index]:
        if source.capture_end_ns > source.capture_start_ns:
            capture = (source.capture_start_ns + shift_ns, source.capture_end_ns + shift_ns)
        else:
            capture = (start_ns, end_ns)
        traces.append(
            TraceData(
                combo=source.combo,
                frequency_hz=source.frequency_hz,
                s11=source.s11,
                s21=source.s21,
                capture_start_ns=capture[0],
                capture_end_ns=capture[1],
            )
        )
    return SweepCollection(
        collection_id=collection_id,
        monotonic_ns=frame_ns,
        traces=traces,
        capture_start_ns=min(t.capture_start_ns for t in traces),
        capture_end_ns=max(t.capture_end_ns for t in traces),
    )


def main() -> int:
    """Replay the configured dataset into the raw SHM rings until stopped."""
    parser = argparse.ArgumentParser(description="Replay a recorded dataset to raw SHM rings")
    parser.add_argument("--config", required=True, type=Path, help="Path to run_config.json")
    args = parser.parse_args()

    logging.basicConfig(level=logging.INFO, format="%(levelname)s %(name)s: %(message)s")
    stop_requested = threading.Event()

    def request_stop(_signum: int, _frame: object) -> None:
        stop_requested.set()

    signal.signal(signal.SIGINT, request_stop)
    signal.signal(signal.SIGTERM, request_stop)

    config = RunConfigModel.load_from_path(args.config)
    config.apply_device_model_constraints()
    logging.getLogger("python_app").setLevel(coerce_level(config.logging.level))

    replay = config.radar.replay
    if not replay.dataset_dir:
        raise RuntimeError("replay_raw_producer requires a non-empty radar.replay.dataset_dir")
    rate_scale = float(replay.rate_scale) if replay.rate_scale > 0 else 1.0
    combos = [ComboKey(input=combo.input, output=combo.output) for combo in config.combos]

    # expanduser: profiles are copied between machines, so "~/..." dataset paths
    # must resolve against the local home instead of failing as a literal "~".
    dataset = _load_dataset(
        Path(replay.dataset_dir).expanduser(), combos, int(config.radar.sweep.points)
    )
    logger.info(
        "Replay producer starting: dataset=%s, snapshots=%d, combos=%d, "
        "median_gap=%.3f s, loop=%s, rate_scale=%.2f",
        replay.dataset_dir, len(dataset), len(combos),
        dataset.median_gap_ns / 1e9, replay.loop, rate_scale,
    )

    raw_writer = ShmRingWriter(
        config.rings.raw.name,
        config.rings.raw.capacity,
        config.rings.raw.slot_size_bytes,
    )
    raw_tap_writer = ShmRingWriter(
        config.rings.raw_tap.name,
        config.rings.raw_tap.capacity,
        config.rings.raw_tap.slot_size_bytes,
    )

    try:
        base_ns = time.monotonic_ns()
        wall_start_s = time.monotonic()
        first_ts = dataset.timestamps_ns[0]
        cycle_span_ns = dataset.timestamps_ns[-1] - first_ts + dataset.median_gap_ns

        collection_id = 1
        publish_failures = 0
        cycle = 0
        while not stop_requested.is_set():
            for index in range(len(dataset)):
                if stop_requested.is_set():
                    break
                # Monotonic replay timeline: recorded offsets re-anchored at
                # base_ns, each loop shifted by one full cycle span.
                frame_ns = (
                    base_ns
                    + (dataset.timestamps_ns[index] - first_ts)
                    + cycle * cycle_span_ns
                )
                # Publish on the recorded cadence (scaled); waiting via the stop
                # event keeps shutdown prompt.
                deadline_s = wall_start_s + (frame_ns - base_ns) / 1e9 / rate_scale
                delay_s = deadline_s - time.monotonic()
                if delay_s > 0 and stop_requested.wait(delay_s):
                    break

                collection = _build_collection(dataset, index, collection_id, frame_ns)
                payload = serialize_trace_collection(collection, RAW_MAGIC)
                if raw_writer.push(payload):
                    raw_tap_writer.push(payload)  # best-effort GUI tap; never fatal
                else:
                    publish_failures += 1
                    if publish_failures == 1 or publish_failures % 100 == 0:
                        logger.error(
                            "Replay payload %d B exceeds ring slot %d B; dropping collection %d (drops=%d)",
                            len(payload), raw_writer.slot_size_bytes, collection_id, publish_failures,
                        )
                if collection_id == 1 or collection_id % 50 == 0:
                    logger.info(
                        "replay collection %d published (snapshot %d/%d, cycle %d)",
                        collection_id, index + 1, len(dataset), cycle,
                    )
                collection_id += 1
                if not config.runtime.continuous:
                    stop_requested.set()
                    break
            if not replay.loop:
                break
            cycle += 1
    finally:
        raw_tap_writer.close()
        raw_writer.close()

    logger.info("replay raw producer stopped")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())