"""Builders for vna_system-compatible sweep-history JSON payloads."""

from __future__ import annotations

from collections import defaultdict
from dataclasses import dataclass
from datetime import datetime, timezone
from typing import Any

import numpy as np

from python_app.models.dataset_model import ResultCollection, SweepCollection, TraceData


@dataclass(frozen=True)
class TraceRecord:
    """One raw/preprocessed trace selected for vna_history export."""

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


def _normalize_channel(channel: str) -> str:
    normalized = str(channel).strip().lower()
    if normalized not in {"s21", "s11"}:
        raise ValueError("channel must be either 's21' or 's11'")
    return normalized


def _select_trace_samples(trace: TraceData, channel: str) -> np.ndarray:
    if channel == "s11":
        return np.asarray(trace.s11, dtype=np.complex128).reshape(-1)
    return np.asarray(trace.s21, dtype=np.complex128).reshape(-1)


def _pick_trace(collection: SweepCollection, input_index: int, output_index: int) -> TraceData | None:
    for trace in collection.traces:
        if int(trace.combo.input_pos) == int(input_index) and int(trace.combo.output_pos) == int(output_index):
            return trace
    return None


def _build_stage_records(
    stage: str,
    history: list[SweepCollection],
    *,
    input_index: int,
    output_index: int,
    channel: str,
) -> list[TraceRecord]:
    records: list[TraceRecord] = []
    for stage_index, collection in enumerate(history):
        trace = _pick_trace(collection, input_index, output_index)
        if trace is None:
            continue

        frequency_hz = np.asarray(trace.frequency_hz, dtype=np.float64).reshape(-1)
        samples = _select_trace_samples(trace, channel)
        if frequency_hz.shape != samples.shape:
            raise ValueError(
                f"Shape mismatch in {stage} stage for collection_id={collection.collection_id}: "
                f"freq{frequency_hz.shape} vs {channel}{samples.shape}"
            )
        if frequency_hz.size == 0:
            continue
        if not (
            np.isfinite(frequency_hz).all()
            and np.isfinite(np.real(samples)).all()
            and np.isfinite(np.imag(samples)).all()
        ):
            raise ValueError(f"Non-finite values in {stage} stage for collection_id={collection.collection_id}")

        records.append(
            TraceRecord(
                stage=stage,
                collection_id=int(collection.collection_id),
                monotonic_ns=int(collection.monotonic_ns),
                stage_index=int(stage_index),
                frequency_hz=frequency_hz,
                samples=samples,
            )
        )
    return records


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(value.real), float(value.imag)] for value in values]


def _build_sweep_history(
    raw_records: list[TraceRecord],
    preprocessed_records: list[TraceRecord],
    *,
    channel: str,
    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

        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.samples),
                "calibrated_points": _complex_to_points(calibrated_source.samples),
                "reference_points": [],
                "vna_config": {
                    "mode": channel,
                    "start_freq": start_freq_hz,
                    "stop_freq": stop_freq_hz,
                    "points": int(base.frequency_hz.size),
                },
            }
        )
    return history


def _stage_alignment_warning(pre_history: list[SweepCollection], result_history: list[ResultCollection]) -> str | None:
    if not pre_history or not result_history:
        return None

    common_size = min(len(pre_history), len(result_history))
    if common_size <= 0:
        return None

    mismatches = 0
    first_mismatch: tuple[int, SweepCollection, ResultCollection] | None = None
    for index in range(common_size):
        pre = pre_history[index]
        result = result_history[index]
        if int(pre.collection_id) != int(result.collection_id) or int(pre.monotonic_ns) != int(result.monotonic_ns):
            mismatches += 1
            if first_mismatch is None:
                first_mismatch = (index, pre, result)

    if mismatches == 0 or first_mismatch is None:
        return None

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


def build_vna_history_payload(
    raw_history: list[SweepCollection],
    preprocessed_history: list[SweepCollection],
    result_history: list[ResultCollection],
    *,
    input_index: int,
    output_index: int,
    channel: str = "s21",
    primary_stage: str = "preprocessed",
) -> dict[str, Any]:
    """Build vna_system-compatible history JSON payload from runtime histories."""
    channel = _normalize_channel(channel)
    if primary_stage not in {"preprocessed", "raw"}:
        raise ValueError("primary_stage must be either 'preprocessed' or 'raw'")

    raw_records = _build_stage_records(
        "raw",
        raw_history,
        input_index=input_index,
        output_index=output_index,
        channel=channel,
    )
    preprocessed_records = _build_stage_records(
        "preprocessed",
        preprocessed_history,
        input_index=input_index,
        output_index=output_index,
        channel=channel,
    )
    if not raw_records and not preprocessed_records:
        raise ValueError(
            "No matching raw/preprocessed traces were found in runtime history "
            f"for input={input_index}, output={output_index}."
        )

    sweep_history = _build_sweep_history(
        raw_records,
        preprocessed_records,
        channel=channel,
        primary_stage=primary_stage,
    )
    if not sweep_history:
        raise ValueError("Conversion produced empty sweep_history.")

    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": "<runtime_history>",
        "input_index": int(input_index),
        "output_index": int(output_index),
        "channel": channel,
        "primary_stage": str(primary_stage),
        "raw_record_count": len(raw_records),
        "preprocessed_record_count": len(preprocessed_records),
        "sweep_history": sweep_history,
    }

    alignment_warning = _stage_alignment_warning(preprocessed_history, result_history)
    if alignment_warning is not None:
        payload["alignment_warning"] = alignment_warning

    return payload
