"""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
from contextlib import ExitStack
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[int, str, int]] = []
    for entry in dataset_dir.iterdir():
        match = _SNAPSHOT_NAME.match(entry.name)
        if entry.is_dir() and match:
            named.append((int(match.group(1)), entry.name, int(match.group(2))))
    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 sorted(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, allow_pickle=False).astype(np.complex64, copy=False)
            freq = np.load(freq_path, allow_pickle=False).astype(np.float32, copy=False)
            s11_path = folder / f"i{combo.input}_o{combo.output}_s11.npy"
            s11 = np.load(s11_path, allow_pickle=False).astype(np.complex64, copy=False) if s11_path.is_file() else np.zeros_like(s21)
            if freq.ndim != 1 or s21.shape != freq.shape or s11.shape != freq.shape:
                raise RuntimeError(f"Snapshot {name}: frequency, S11 and S21 must be equally sized 1D arrays")
            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=s11,
                    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,
    )

    cleanup = ExitStack()
    try:
        raw_writer = ShmRingWriter(
            config.rings.raw.name,
            config.rings.raw.capacity,
            config.rings.raw.slot_size_bytes,
        )
        cleanup.callback(raw_writer.close)
        raw_tap_writer = ShmRingWriter(
            config.rings.raw_tap.name,
            config.rings.raw_tap.capacity,
            config.rings.raw_tap.slot_size_bytes,
        )
        cleanup.callback(raw_tap_writer.close)

        payload_size = len(serialize_trace_collection(
            _build_collection(dataset, 0, 1, time.monotonic_ns()), RAW_MAGIC
        ))
        for ring in (config.rings.raw, config.rings.raw_tap):
            if ring.slot_size_bytes < payload_size:
                raise ValueError(f"Replay requires {payload_size} bytes per slot; ring {ring.name} has {ring.slot_size_bytes}")

        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
        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 not raw_writer.push(payload):
                    raise RuntimeError(f"Replay raw ring rejected collection {collection_id}")
                raw_tap_writer.push(payload)
                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:
        cleanup.close()

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


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