"""Convert radar_system runtime snapshots into vna_system sweep-history JSON."""

from __future__ import annotations

import argparse
import json
from collections import defaultdict
from dataclasses import dataclass
from datetime import datetime, timezone
from pathlib import Path
from typing import Any

import numpy as np


@dataclass(frozen=True)
class TraceRecord:
    """One raw/preprocessed trace extracted from a snapshot collection directory."""

    stage: str
    collection_id: int
    monotonic_ns: int
    stage_index: int
    frequency_hz: np.ndarray
    s21: np.ndarray


@dataclass(frozen=True)
class CollectionRef:
    """Minimal collection identity descriptor for stage-alignment diagnostics."""

    stage_index: int
    collection_id: int
    monotonic_ns: int


def _load_json(path: Path) -> dict[str, Any]:
    payload = json.loads(path.read_text(encoding="utf-8"))
    if not isinstance(payload, dict):
        raise ValueError(f"JSON root must be object: {path}")
    return payload


def _collection_dirs(stage_dir: Path) -> list[Path]:
    if not stage_dir.exists():
        return []
    return sorted([path for path in stage_dir.iterdir() if path.is_dir()], key=lambda path: path.name)


def _parse_stage_index(name: str, fallback: int) -> int:
    prefix = name.split("_", 1)[0]
    return int(prefix) if prefix.isdigit() else fallback


def _pick_trace_meta(meta: dict[str, Any], input_index: int, output_index: int) -> dict[str, Any] | None:
    traces = meta.get("traces", [])
    if not isinstance(traces, list):
        return None
    for trace in traces:
        if not isinstance(trace, dict):
            continue
        if int(trace.get("input", -1)) == input_index and int(trace.get("output", -1)) == output_index:
            return trace
    return None


def _load_stage_records(
    snapshot_dir: Path,
    stage: str,
    *,
    input_index: int,
    output_index: int,
) -> list[TraceRecord]:
    stage_dir = snapshot_dir / stage
    records: list[TraceRecord] = []

    for fallback_idx, collection_dir in enumerate(_collection_dirs(stage_dir)):
        meta_path = collection_dir / "meta.json"
        if not meta_path.exists():
            continue

        meta = _load_json(meta_path)
        trace_meta = _pick_trace_meta(meta, input_index, output_index)
        if trace_meta is None:
            continue

        freq_file = str(trace_meta.get("freq_file", ""))
        s21_file = str(trace_meta.get("s21_file", ""))
        if not freq_file or not s21_file:
            continue

        frequency_hz = np.asarray(np.load(collection_dir / freq_file), dtype=np.float64).reshape(-1)
        s21 = np.asarray(np.load(collection_dir / s21_file), dtype=np.complex128).reshape(-1)
        if frequency_hz.shape != s21.shape:
            raise ValueError(f"Shape mismatch in {collection_dir}: freq{frequency_hz.shape} vs s21{s21.shape}")
        if frequency_hz.size == 0:
            continue
        if not (
            np.isfinite(frequency_hz).all()
            and np.isfinite(np.real(s21)).all()
            and np.isfinite(np.imag(s21)).all()
        ):
            raise ValueError(f"Non-finite values in {collection_dir}")

        records.append(
            TraceRecord(
                stage=stage,
                collection_id=int(meta.get("collection_id", -1)),
                monotonic_ns=int(meta.get("monotonic_ns", 0)),
                stage_index=_parse_stage_index(collection_dir.name, fallback_idx),
                frequency_hz=frequency_hz,
                s21=s21,
            )
        )

    return records


def _load_stage_refs(snapshot_dir: Path, stage: str) -> list[CollectionRef]:
    """Load `(index, collection_id, monotonic_ns)` for one snapshot stage."""
    stage_dir = snapshot_dir / stage
    refs: list[CollectionRef] = []

    for fallback_idx, collection_dir in enumerate(_collection_dirs(stage_dir)):
        meta_path = collection_dir / "meta.json"
        if not meta_path.exists():
            continue
        meta = _load_json(meta_path)
        refs.append(
            CollectionRef(
                stage_index=_parse_stage_index(collection_dir.name, fallback_idx),
                collection_id=int(meta.get("collection_id", -1)),
                monotonic_ns=int(meta.get("monotonic_ns", 0)),
            )
        )
    return refs


def _index_by_collection_occurrence(records: list[TraceRecord]) -> tuple[dict[tuple[int, int], TraceRecord], list[tuple[int, int]]]:
    counters: defaultdict[int, int] = defaultdict(int)
    record_map: dict[tuple[int, int], TraceRecord] = {}
    order: list[tuple[int, int]] = []
    for record in records:
        occurrence = counters[record.collection_id]
        counters[record.collection_id] += 1
        key = (record.collection_id, occurrence)
        record_map[key] = record
        order.append(key)
    return record_map, order


def _complex_to_points(values: np.ndarray) -> list[list[float]]:
    return [[float(v.real), float(v.imag)] for v in values]


def _build_sweep_history(
    raw_records: list[TraceRecord],
    preprocessed_records: list[TraceRecord],
    *,
    primary_stage: str,
) -> list[dict[str, Any]]:
    raw_map, raw_order = _index_by_collection_occurrence(raw_records)
    pre_map, pre_order = _index_by_collection_occurrence(preprocessed_records)

    if primary_stage == "preprocessed":
        primary_order = pre_order or raw_order
    else:
        primary_order = raw_order or pre_order

    history: list[dict[str, Any]] = []
    for fallback_index, key in enumerate(primary_order):
        raw = raw_map.get(key)
        pre = pre_map.get(key)
        base = pre or raw
        if base is None:
            continue

        # vna_system uses calibrated_data if present, otherwise sweep_data.
        sweep_source = raw or pre
        calibrated_source = pre or raw
        if sweep_source is None or calibrated_source is None:
            continue

        start_freq_hz = float(base.frequency_hz[0])
        stop_freq_hz = float(base.frequency_hz[-1])
        timestamp_sec = float(base.monotonic_ns) / 1_000_000_000.0 if base.monotonic_ns > 0 else float(fallback_index)

        history.append(
            {
                "timestamp": timestamp_sec,
                "sweep_points": _complex_to_points(sweep_source.s21),
                "calibrated_points": _complex_to_points(calibrated_source.s21),
                "reference_points": [],
                "vna_config": {
                    "mode": "s11",
                    "start_freq": start_freq_hz,
                    "stop_freq": stop_freq_hz,
                    "points": int(base.frequency_hz.size),
                },
            }
        )

    return history


def _stage_alignment_warning(pre_refs: list[CollectionRef], result_refs: list[CollectionRef]) -> str | None:
    """Return warning text when preprocessed/results stages are not identity-aligned."""
    if not pre_refs or not result_refs:
        return None

    pre_by_index = {ref.stage_index: ref for ref in pre_refs}
    result_by_index = {ref.stage_index: ref for ref in result_refs}
    common_indices = sorted(set(pre_by_index) & set(result_by_index))
    if not common_indices:
        return None

    mismatches = 0
    first_mismatch: tuple[int, CollectionRef, CollectionRef] | None = None
    for index in common_indices:
        pre = pre_by_index[index]
        result = result_by_index[index]
        if pre.collection_id != result.collection_id or pre.monotonic_ns != result.monotonic_ns:
            mismatches += 1
            if first_mismatch is None:
                first_mismatch = (index, pre, result)

    if mismatches == 0:
        return None

    assert first_mismatch is not None
    idx, pre, result = first_mismatch
    return (
        "WARNING: snapshot stages are not fully aligned (preprocessed vs results). "
        f"Mismatches={mismatches}/{len(common_indices)}. "
        f"First mismatch at index={idx}: "
        f"pre=(id={pre.collection_id},ns={pre.monotonic_ns}) vs "
        f"results=(id={result.collection_id},ns={result.monotonic_ns}). "
        "Export uses preprocessed traces; loaded view in vna_system may differ from "
        "radar_system on-screen replayed results."
    )


def _build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        description=(
            "Convert radar_system snapshot (numpy-directory-v1) to a vna_system-compatible "
            "history JSON file with `sweep_history`."
        )
    )
    parser.add_argument("snapshot_dir", type=Path, help="Path to snapshot directory containing raw/preprocessed/results.")
    parser.add_argument(
        "-o",
        "--output",
        type=Path,
        default=None,
        help="Output JSON path (default: <snapshot_dir>/vna_bscan_history.json).",
    )
    parser.add_argument("--input", dest="input_index", type=int, default=0, help="Input switch index to export.")
    parser.add_argument("--output-index", dest="output_index", type=int, default=0, help="Output switch index to export.")
    parser.add_argument(
        "--primary-stage",
        choices=("preprocessed", "raw"),
        default="preprocessed",
        help="Stage order to drive collection selection/alignment.",
    )
    parser.add_argument(
        "--last-n",
        type=int,
        default=0,
        help="Keep only the last N sweeps in output (0 means all available).",
    )
    return parser


def main() -> None:
    parser = _build_parser()
    args = parser.parse_args()

    snapshot_dir = args.snapshot_dir.expanduser().resolve()
    if not snapshot_dir.exists():
        raise FileNotFoundError(f"Snapshot directory not found: {snapshot_dir}")

    output_path = (
        args.output.expanduser().resolve()
        if args.output is not None
        else snapshot_dir / "vna_bscan_history.json"
    )

    raw_records = _load_stage_records(
        snapshot_dir,
        "raw",
        input_index=args.input_index,
        output_index=args.output_index,
    )
    preprocessed_records = _load_stage_records(
        snapshot_dir,
        "preprocessed",
        input_index=args.input_index,
        output_index=args.output_index,
    )
    if not raw_records and not preprocessed_records:
        raise ValueError(
            "No matching raw/preprocessed traces were found in snapshot "
            f"for input={args.input_index}, output={args.output_index}."
        )

    sweep_history = _build_sweep_history(raw_records, preprocessed_records, primary_stage=args.primary_stage)
    if args.last_n > 0:
        sweep_history = sweep_history[-args.last_n :]
    if not sweep_history:
        raise ValueError("Conversion produced empty sweep_history.")

    pre_refs = _load_stage_refs(snapshot_dir, "preprocessed")
    result_refs = _load_stage_refs(snapshot_dir, "results")
    alignment_warning = _stage_alignment_warning(pre_refs, result_refs)

    manifest_path = snapshot_dir / "manifest.json"
    manifest = _load_json(manifest_path) if manifest_path.exists() else {}

    payload: dict[str, Any] = {
        "format": "vna-system-history-v1",
        "converter": "python_app/scripts/convert_snapshot_to_vna_history.py",
        "converted_at_utc": datetime.now(timezone.utc).isoformat(),
        "source_snapshot_dir": str(snapshot_dir),
        "input_index": int(args.input_index),
        "output_index": int(args.output_index),
        "primary_stage": args.primary_stage,
        "raw_record_count": len(raw_records),
        "preprocessed_record_count": len(preprocessed_records),
        "sweep_history": sweep_history,
    }
    if alignment_warning is not None:
        payload["alignment_warning"] = alignment_warning
    if manifest:
        payload["snapshot_manifest"] = manifest

    output_path.parent.mkdir(parents=True, exist_ok=True)
    output_path.write_text(json.dumps(payload, ensure_ascii=False, indent=2), encoding="utf-8")

    print(
        "Converted snapshot to vna_system history JSON:\n"
        f"  input snapshot: {snapshot_dir}\n"
        f"  output file:    {output_path}\n"
        f"  sweeps written: {len(sweep_history)}\n"
        f"  raw records:    {len(raw_records)}\n"
        f"  pre records:    {len(preprocessed_records)}"
    )
    if alignment_warning is not None:
        print(f"[convert-warning] {alignment_warning}")


if __name__ == "__main__":
    main()
