"""Snapshot selection and filesystem writers for runtime collection histories."""

from __future__ import annotations

import json
import logging
from pathlib import Path
from typing import Any, TypeVar

import numpy as np

from python_app.models.dataset_model import ResultCollection, SweepCollection
from python_app.storage.npz.paths import collection_dir_name, sanitize_path_component
from python_app.storage.npz.serialize import serialize_result_collection, serialize_trace_collection

logger = logging.getLogger(__name__)

TCollection = TypeVar("TCollection")


def select_aligned_histories(
    raw_history: list[SweepCollection],
    preprocessed_history: list[SweepCollection],
    result_history: list[ResultCollection],
    last_n: int,
) -> tuple[list[SweepCollection], list[SweepCollection], list[ResultCollection], dict[str, Any]]:
    """Select aligned tails with `results` as preferred anchor stage."""
    raw_index, raw_pos = _index_by_collection_sequence(raw_history)
    pre_index, pre_pos = _index_by_collection_sequence(preprocessed_history)
    result_index, result_pos = _index_by_collection_sequence(result_history)

    if result_index:
        ordered_result_keys = sorted(result_index, key=lambda key: result_pos[key])
        selected_keys = ordered_result_keys[-last_n:]
        selected_raw = [raw_index[key] for key in selected_keys if key in raw_index]
        selected_pre = [pre_index[key] for key in selected_keys if key in pre_index]
        selected_results = [result_index[key] for key in selected_keys]
        raw_missing = len(selected_keys) - len(selected_raw)
        pre_missing = len(selected_keys) - len(selected_pre)
        return (
            selected_raw,
            selected_pre,
            selected_results,
            {
                "selection_mode": "result_tail_with_optional_alignment",
                "anchor_stage": "results",
                "selected_collection_ids": [int(key[0]) for key in selected_keys],
                "aligned_key_count": len(selected_keys),
                "raw_missing_count": raw_missing,
                "preprocessed_missing_count": pre_missing,
            },
        )

    if raw_index:
        ordered_raw_keys = sorted(raw_index, key=lambda key: raw_pos[key])
        selected_keys = ordered_raw_keys[-last_n:]
        selected_raw = [raw_index[key] for key in selected_keys]
        selected_pre = [pre_index[key] for key in selected_keys if key in pre_index]
        selected_results = [result_index[key] for key in selected_keys if key in result_index]
        pre_missing = len(selected_keys) - len(selected_pre)
        result_missing = len(selected_keys) - len(selected_results)
        return (
            selected_raw,
            selected_pre,
            selected_results,
            {
                "selection_mode": "raw_tail_with_optional_alignment",
                "anchor_stage": "raw",
                "selected_collection_ids": [int(key[0]) for key in selected_keys],
                "aligned_key_count": len(selected_keys),
                "preprocessed_missing_count": pre_missing,
                "result_missing_count": result_missing,
            },
        )

    if pre_index:
        ordered_pre_keys = sorted(pre_index, key=lambda key: pre_pos[key])
        selected_keys = ordered_pre_keys[-last_n:]
        selected_raw = [raw_index[key] for key in selected_keys if key in raw_index]
        selected_pre = [pre_index[key] for key in selected_keys]
        selected_results = [result_index[key] for key in selected_keys if key in result_index]
        raw_missing = len(selected_keys) - len(selected_raw)
        result_missing = len(selected_keys) - len(selected_results)
        return (
            selected_raw,
            selected_pre,
            selected_results,
            {
                "selection_mode": "preprocessed_tail_with_optional_alignment",
                "anchor_stage": "preprocessed",
                "selected_collection_ids": [int(key[0]) for key in selected_keys],
                "aligned_key_count": len(selected_keys),
                "raw_missing_count": raw_missing,
                "result_missing_count": result_missing,
            },
        )

    return (
        raw_history[-last_n:],
        preprocessed_history[-last_n:],
        result_history[-last_n:],
        {
            "selection_mode": "independent_tail",
            "anchor_stage": "none",
            "selected_collection_ids": [],
            "aligned_key_count": 0,
        },
    )


def save_trace_history_binary(stage_dir: Path, history: list[SweepCollection], magic: int) -> None:
    """Write binary trace history with lightweight metadata sidecars."""
    stage_dir.mkdir(parents=True, exist_ok=True)
    logger.debug("Writing %d binary trace collection(s) to %s", len(history), stage_dir)
    for index, collection in enumerate(history):
        binary_path = stage_dir / f"{index:04d}.bin"
        metadata_path = stage_dir / f"{index:04d}.json"
        binary_path.write_bytes(serialize_trace_collection(collection, magic))
        metadata_path.write_text(
            json.dumps(
                {
                    "collection_id": collection.collection_id,
                    "monotonic_ns": collection.monotonic_ns,
                    "capture_start_ns": int(collection.capture_start_ns),
                    "capture_end_ns": int(collection.capture_end_ns),
                    "trace_count": len(collection.traces),
                },
                indent=2,
            ),
            encoding="utf-8",
        )


def save_result_history_binary(stage_dir: Path, history: list[ResultCollection]) -> None:
    """Write binary processed-result history with metadata sidecars."""
    stage_dir.mkdir(parents=True, exist_ok=True)
    logger.debug("Writing %d binary result collection(s) to %s", len(history), stage_dir)
    for index, collection in enumerate(history):
        binary_path = stage_dir / f"{index:04d}.bin"
        metadata_path = stage_dir / f"{index:04d}.json"
        binary_path.write_bytes(serialize_result_collection(collection))
        metadata_path.write_text(
            json.dumps(
                {
                    "collection_id": collection.collection_id,
                    "monotonic_ns": collection.monotonic_ns,
                    "collection_payload_count": len(collection.collection_payloads),
                    "block_count": len(collection.blocks),
                },
                indent=2,
            ),
            encoding="utf-8",
        )


def save_trace_history_numpy(stage_dir: Path, history: list[SweepCollection]) -> None:
    """Write raw/preprocessed collections as NumPy directory tree."""
    stage_dir.mkdir(parents=True, exist_ok=True)
    logger.debug("Writing %d NumPy trace collection(s) to %s", len(history), stage_dir)
    for index, collection in enumerate(history):
        collection_dir = stage_dir / collection_dir_name(index, collection.collection_id, collection.monotonic_ns)
        collection_dir.mkdir(parents=True, exist_ok=False)

        traces_meta: list[dict[str, int | str]] = []
        for trace in collection.traces:
            tag = f"i{trace.combo.input}_o{trace.combo.output}"
            freq = np.asarray(trace.frequency_hz, dtype=np.float32)
            s11 = np.asarray(trace.s11, dtype=np.complex64)
            s21 = np.asarray(trace.s21, dtype=np.complex64)
            np.save(collection_dir / f"{tag}_freq.npy", freq)
            np.save(collection_dir / f"{tag}_s11.npy", s11)
            np.save(collection_dir / f"{tag}_s21.npy", s21)
            traces_meta.append(
                {
                    "input": int(trace.combo.input),
                    "output": int(trace.combo.output),
                    "points": int(freq.size),
                    "freq_file": f"{tag}_freq.npy",
                    "s11_file": f"{tag}_s11.npy",
                    "s21_file": f"{tag}_s21.npy",
                }
            )

        (collection_dir / "meta.json").write_text(
            json.dumps(
                {
                    "collection_id": int(collection.collection_id),
                    "monotonic_ns": int(collection.monotonic_ns),
                    "capture_start_ns": int(collection.capture_start_ns),
                    "capture_end_ns": int(collection.capture_end_ns),
                    "trace_count": len(collection.traces),
                    "traces": traces_meta,
                },
                indent=2,
            ),
            encoding="utf-8",
        )


def save_result_history_numpy(stage_dir: Path, history: list[ResultCollection]) -> None:
    """Write processed result collections as NumPy directory tree."""
    stage_dir.mkdir(parents=True, exist_ok=True)
    logger.debug("Writing %d NumPy result collection(s) to %s", len(history), stage_dir)
    for index, collection in enumerate(history):
        collection_dir = stage_dir / collection_dir_name(index, collection.collection_id, collection.monotonic_ns)
        collection_dir.mkdir(parents=True, exist_ok=False)

        collection_payload_meta: list[dict[str, int | str | float]] = []
        for payload_index, payload in enumerate(collection.collection_payloads):
            safe_name = sanitize_path_component(payload.processing_name or "processor")
            base_name = f"collection_{payload_index:03d}_{safe_name}_kind{payload.kind}"
            if payload.kind == 1:
                freq = np.asarray(payload.frequency_hz, dtype=np.float32)
                trace = np.asarray(payload.trace, dtype=np.complex64)
                np.save(collection_dir / f"{base_name}_freq.npy", freq)
                np.save(collection_dir / f"{base_name}_trace.npy", trace)
                collection_payload_meta.append(
                    {
                        "kind": int(payload.kind),
                        "name": payload.processing_name,
                        "points": int(freq.size),
                        "freq_file": f"{base_name}_freq.npy",
                        "trace_file": f"{base_name}_trace.npy",
                    }
                )
            elif payload.kind == 2:
                scalar = np.asarray([float(payload.scalar_value)], dtype=np.float32)
                np.save(collection_dir / f"{base_name}_scalar.npy", scalar)
                collection_payload_meta.append(
                    {
                        "kind": int(payload.kind),
                        "name": payload.processing_name,
                        "scalar_file": f"{base_name}_scalar.npy",
                        "scalar_value": float(payload.scalar_value),
                    }
                )
            elif payload.kind == 3:
                image_x_axis = np.asarray(payload.image_x_axis, dtype=np.float32)
                image_y_axis = np.asarray(payload.image_y_axis, dtype=np.float32)
                image = np.asarray(payload.image, dtype=np.float32)
                np.save(collection_dir / f"{base_name}_x_axis.npy", image_x_axis)
                np.save(collection_dir / f"{base_name}_y_axis.npy", image_y_axis)
                np.save(collection_dir / f"{base_name}_image.npy", image)
                collection_payload_meta.append(
                    {
                        "kind": int(payload.kind),
                        "name": payload.processing_name,
                        "x_points": int(image_x_axis.size),
                        "y_points": int(image_y_axis.size),
                        "x_axis_file": f"{base_name}_x_axis.npy",
                        "y_axis_file": f"{base_name}_y_axis.npy",
                        "image_file": f"{base_name}_image.npy",
                    }
                )
            elif payload.kind == 4:
                table = np.asarray(payload.table, dtype=np.float32)
                np.save(collection_dir / f"{base_name}_table.npy", table)
                collection_payload_meta.append(
                    {
                        "kind": int(payload.kind),
                        "name": payload.processing_name,
                        "rows": int(table.shape[0]) if table.ndim == 2 else 0,
                        "columns": int(table.shape[1]) if table.ndim == 2 else 0,
                        "table_file": f"{base_name}_table.npy",
                    }
                )

        blocks_meta: list[dict[str, int | str | list[dict[str, int | str | float]]]] = []
        for block_index, block in enumerate(collection.blocks):
            block_dir = collection_dir / f"block_{block_index:03d}_i{block.combo.input}_o{block.combo.output}"
            block_dir.mkdir(parents=True, exist_ok=False)

            payload_meta: list[dict[str, int | str | float]] = []
            for payload_index, payload in enumerate(block.payloads):
                safe_name = sanitize_path_component(payload.processing_name or "processor")
                base_name = f"{payload_index:03d}_{safe_name}_kind{payload.kind}"
                if payload.kind == 1:
                    freq = np.asarray(payload.frequency_hz, dtype=np.float32)
                    trace = np.asarray(payload.trace, dtype=np.complex64)
                    np.save(block_dir / f"{base_name}_freq.npy", freq)
                    np.save(block_dir / f"{base_name}_trace.npy", trace)
                    payload_meta.append(
                        {
                            "kind": int(payload.kind),
                            "name": payload.processing_name,
                            "points": int(freq.size),
                            "freq_file": f"{base_name}_freq.npy",
                            "trace_file": f"{base_name}_trace.npy",
                        }
                    )
                elif payload.kind == 2:
                    scalar = np.asarray([float(payload.scalar_value)], dtype=np.float32)
                    np.save(block_dir / f"{base_name}_scalar.npy", scalar)
                    payload_meta.append(
                        {
                            "kind": int(payload.kind),
                            "name": payload.processing_name,
                            "scalar_file": f"{base_name}_scalar.npy",
                            "scalar_value": float(payload.scalar_value),
                        }
                    )
                elif payload.kind == 3:
                    image_x_axis = np.asarray(payload.image_x_axis, dtype=np.float32)
                    image_y_axis = np.asarray(payload.image_y_axis, dtype=np.float32)
                    image = np.asarray(payload.image, dtype=np.float32)
                    np.save(block_dir / f"{base_name}_x_axis.npy", image_x_axis)
                    np.save(block_dir / f"{base_name}_y_axis.npy", image_y_axis)
                    np.save(block_dir / f"{base_name}_image.npy", image)
                    payload_meta.append(
                        {
                            "kind": int(payload.kind),
                            "name": payload.processing_name,
                            "x_points": int(image_x_axis.size),
                            "y_points": int(image_y_axis.size),
                            "x_axis_file": f"{base_name}_x_axis.npy",
                            "y_axis_file": f"{base_name}_y_axis.npy",
                            "image_file": f"{base_name}_image.npy",
                        }
                    )
                elif payload.kind == 4:
                    table = np.asarray(payload.table, dtype=np.float32)
                    np.save(block_dir / f"{base_name}_table.npy", table)
                    payload_meta.append(
                        {
                            "kind": int(payload.kind),
                            "name": payload.processing_name,
                            "rows": int(table.shape[0]) if table.ndim == 2 else 0,
                            "columns": int(table.shape[1]) if table.ndim == 2 else 0,
                            "table_file": f"{base_name}_table.npy",
                        }
                    )

            blocks_meta.append(
                {
                    "input": int(block.combo.input),
                    "output": int(block.combo.output),
                    "payload_count": len(block.payloads),
                    "dir": block_dir.name,
                    "payloads": payload_meta,
                }
            )

        (collection_dir / "meta.json").write_text(
            json.dumps(
                {
                    "collection_id": int(collection.collection_id),
                    "monotonic_ns": int(collection.monotonic_ns),
                    "collection_payload_count": len(collection.collection_payloads),
                    "collection_payloads": collection_payload_meta,
                    "block_count": len(collection.blocks),
                    "blocks": blocks_meta,
                },
                indent=2,
            ),
            encoding="utf-8",
        )


def validate_snapshot_name(snapshot_name: str) -> str:
    """Normalize snapshot directory stem."""
    return sanitize_path_component(snapshot_name)


def _index_by_collection_sequence(
    history: list[TCollection],
) -> tuple[dict[tuple[int, int], TCollection], dict[tuple[int, int], int]]:
    """Index collections by latest-first `(collection_id, occurrence)` preserving history position.

    Occurrence is counted from the tail (newest item is occurrence `0` for its
    collection id). This avoids cross-run misalignment when collection ids are
    reused after restarts and one stage keeps longer history than another.
    """
    indexed: dict[tuple[int, int], TCollection] = {}
    positions: dict[tuple[int, int], int] = {}
    seen_count_from_tail_by_id: dict[int, int] = {}

    # Walk from newest to oldest so occurrence=0 always means the most recent
    # instance for the given collection id.
    for index in range(len(history) - 1, -1, -1):
        item = history[index]
        collection_id = int(getattr(item, "collection_id"))
        occurrence = seen_count_from_tail_by_id.get(collection_id, 0)
        key = (collection_id, occurrence)
        indexed[key] = item
        positions[key] = index
        seen_count_from_tail_by_id[collection_id] = occurrence + 1
    return indexed, positions