"""
Regression tests for HDF5 snapshot decoding via AmunH5.

Compares a gzipped HDF5 snapshot (tests/data/snapshot.h5.gz) against the
AmunXML baseline snapshot (tests/data/snapshot-none.zip). The two snapshots
are stated to be identical, so attributes and data arrays should match.
"""
from __future__ import annotations

import gzip
import pathlib
import zipfile

import numpy as np
import pytest


H5_GZ = "snapshot.h5.gz"
XML_BASELINE_ZIP = "snapshot-none.zip"


def _load_amunxml(snapshot_dir: pathlib.Path):
    from amunpy.amunxml import AmunXML
    return AmunXML(str(snapshot_dir))


def _load_amunh5(h5_path: pathlib.Path):
    # Adjust module path if your class lives elsewhere, e.g. amunpy.amunh5
    from amunpy.amunh5 import AmunH5
    return AmunH5(str(h5_path))


def _is_scalar_like(x) -> bool:
    return isinstance(x, (str, bytes, int, float, bool, np.integer, np.floating))


def _normalize_scalar(x):
    # Normalize numpy scalars/bytes for stable comparisons
    if isinstance(x, np.generic):
        x = x.item()
    if isinstance(x, bytes):
        try:
            return x.decode("utf-8")
        except Exception:
            return x
    return x


def _compare_arrays(a: np.ndarray, b: np.ndarray, label: str):
    # If dtype matches, require exact equality; otherwise still require exact numeric equality.
    if a.dtype == b.dtype:
        np.testing.assert_array_equal(a, b, err_msg=label)
    else:
        np.testing.assert_allclose(a, b, rtol=0.0, atol=0.0, err_msg=label)


@pytest.fixture(scope="session")
def extracted_amunxml_baseline(tmp_path_factory) -> pathlib.Path:
    """
    Extract tests/data/snapshot-none.zip into a temp dir and return snapshot/ path.
    """
    root = pathlib.Path("tests/data")
    zpath = root / XML_BASELINE_ZIP
    if not zpath.exists():
        pytest.skip(f"Missing baseline zip: {zpath}")

    out_base = tmp_path_factory.mktemp("amunxml_baseline")
    outdir = out_base / "snapshot-none"
    outdir.mkdir(parents=True, exist_ok=True)

    with zipfile.ZipFile(zpath, "r") as zf:
        zf.extractall(outdir)

    snapdir = outdir / "snapshot"
    if not snapdir.exists() or not (snapdir / "metadata.xml").exists():
        pytest.skip("Baseline zip does not contain expected snapshot/metadata.xml structure")

    return snapdir


@pytest.fixture(scope="session")
def extracted_hdf5_snapshot(tmp_path_factory) -> pathlib.Path:
    """
    Gunzip tests/data/snapshot.h5.gz into a temp dir and return the .h5 path.
    """
    root = pathlib.Path("tests/data")
    gzpath = root / H5_GZ
    if not gzpath.exists():
        pytest.skip(f"Missing HDF5 snapshot: {gzpath}")

    out_base = tmp_path_factory.mktemp("hdf5_snapshot")
    h5path = out_base / "p000000_00000.h5"

    with gzip.open(gzpath, "rb") as f_in, open(h5path, "wb") as f_out:
        f_out.write(f_in.read())

    return h5path


@pytest.fixture(scope="session")
def xml_snapshot(extracted_amunxml_baseline):
    return _load_amunxml(extracted_amunxml_baseline)


@pytest.fixture(scope="session")
def h5_snapshot(extracted_hdf5_snapshot):
    return _load_amunh5(extracted_hdf5_snapshot)


def test_hdf5_attributes_match_amunxml(xml_snapshot, h5_snapshot):
    """
    Compare scalar attributes between HDF5 (AmunH5.attributes) and AmunXML.attributes.

    Uses intersection of keys to avoid failures due to format-specific metadata.
    """
    xml_attrs = {
        str(k): _normalize_scalar(v)
        for k, v in (xml_snapshot.attributes or {}).items()
        if _is_scalar_like(v) or isinstance(v, np.generic)
    }
    h5_attrs = {
        str(k): _normalize_scalar(v)
        for k, v in (h5_snapshot.attributes or {}).items()
        if _is_scalar_like(v) or isinstance(v, np.generic)
    }

    common = sorted(set(xml_attrs) & set(h5_attrs))
    if not common:
        pytest.skip("No common scalar attribute keys between AmunXML and AmunH5")

    for k in common:
        assert xml_attrs[k] == h5_attrs[k], f"Attribute mismatch for key '{k}'"


def test_hdf5_variable_set_covers_core_fields(xml_snapshot, h5_snapshot):
    """
    Sanity check: ensure core physics fields exist in both snapshots.
    """
    xml_vars = set((xml_snapshot.variables or {}).keys())
    h5_vars = set((h5_snapshot.variables or {}).keys())

    core = {
        "density",
        "x-velocity",
        "y-velocity",
        "x-magnetic field",
        "y-magnetic field",
    }

    missing_xml = sorted(core - xml_vars)
    missing_h5 = sorted(core - h5_vars)

    if missing_xml:
        pytest.skip(f"Baseline AmunXML snapshot does not contain required core fields: {missing_xml}")
    if missing_h5:
        pytest.skip(f"AmunH5 snapshot does not contain required core fields: {missing_h5}")


@pytest.mark.parametrize(
    "field",
    [
        "density",
        "x-velocity",
        "y-velocity",
        "x-magnetic field",
        "y-magnetic field",
    ],
)
def test_hdf5_core_fields_identical_to_amunxml(xml_snapshot, h5_snapshot, field):
    """
    Compare uniform-grid arrays produced by dataset(field) for both readers.
    """
    if field not in (xml_snapshot.variables or {}):
        pytest.skip(f"Baseline AmunXML snapshot missing field '{field}'")
    if field not in (h5_snapshot.variables or {}):
        pytest.skip(f"AmunH5 snapshot missing field '{field}'")

    xml_arr = np.asarray(xml_snapshot.dataset(field))
    h5_arr = np.asarray(h5_snapshot.dataset(field))

    _compare_arrays(h5_arr, xml_arr, label=f"Data mismatch for '{field}' (AmunH5 vs AmunXML)")


def test_hdf5_compare_all_common_fields(xml_snapshot, h5_snapshot):
    """
    Broader check: compare every variable name that exists on both sides.
    """
    xml_vars = set((xml_snapshot.variables or {}).keys())
    h5_vars = set((h5_snapshot.variables or {}).keys())

    common = sorted(xml_vars & h5_vars)
    if not common:
        pytest.skip("No common variable names between AmunXML and AmunH5")

    for name in common:
        xml_arr = np.asarray(xml_snapshot.dataset(name))
        h5_arr = np.asarray(h5_snapshot.dataset(name))
        _compare_arrays(h5_arr, xml_arr, label=f"Data mismatch for '{name}'")
