"""Frozen null bundle serialization."""

from __future__ import annotations

import json
from dataclasses import asdict, dataclass
from pathlib import Path
from typing import Any

from babappa.model.energy import GaussianEnergyModel
from babappa.utils.fs import write_json


@dataclass
class FrozenNullBundle:
    bundle_version: str
    taxa_order: list[str]
    representation_spec: dict[str, Any]
    smoothing_bandwidth: float
    null_model: dict[str, Any]
    energy_model: dict[str, Any]
    training: dict[str, Any]

    def to_dict(self) -> dict[str, Any]:
        return asdict(self)


def make_bundle(
    taxa_order: list[str],
    representation_spec: dict[str, Any],
    smoothing_bandwidth: float,
    null_model: dict[str, Any],
    energy_model: GaussianEnergyModel,
    training: dict[str, Any],
) -> FrozenNullBundle:
    return FrozenNullBundle(
        bundle_version="1.0",
        taxa_order=taxa_order,
        representation_spec=representation_spec,
        smoothing_bandwidth=float(smoothing_bandwidth),
        null_model=null_model,
        energy_model=energy_model.to_dict(),
        training=training,
    )


def write_bundle(path: str | Path, bundle: FrozenNullBundle) -> None:
    write_json(path, bundle.to_dict())


def read_bundle(path: str | Path) -> FrozenNullBundle:
    in_path = Path(path)
    payload = json.loads(in_path.read_text(encoding="utf-8"))
    return FrozenNullBundle(**payload)
