"""
Regression tests for AmunXML snapshot decoding.

These tests load real AMUN snapshots stored as ZIP archives and verify
that metadata parsing, AMR structure, binary decoding, data filters, and
derived quantities behave identically across different compression and
encoding variants.
"""
import pathlib
import zipfile

import numpy as np
import pytest


# ----------------------------
# Snapshot discovery / loading
# ----------------------------

SNAPSHOT_ZIPS = [
    "snapshot-none.zip",
    "snapshot-lz4.zip",
    "snapshot-zstd.zip",
    "snapshot-lzma.zip",
    "snapshot-zstd-bytedelta.zip",
    "snapshot-zstd-shuffle.zip",
]


def _require_codec_for(name: str):
    """
    Skip tests gracefully if the environment lacks a dependency for a codec.
    (none uses no dependency; lzma is stdlib)
    """
    name = name.lower()
    if "zstd" in name:
        pytest.importorskip("zstandard")
    elif "lz4" in name:
        pytest.importorskip("lz4.frame")
    elif "lzma" in name:
        # stdlib
        pass


@pytest.fixture(scope="session")
def extracted_snapshots(tmp_path_factory):
    """
    Extract all snapshot zips from tests/data into per-zip temp dirs.
    Return a dict {zip_basename: snapshot_root_dir}.
    """
    root = pathlib.Path("tests/data")
    missing = [z for z in SNAPSHOT_ZIPS if not (root / z).exists()]
    if missing:
        pytest.skip(f"Missing snapshot zips in tests/data: {missing}")

    out_base = tmp_path_factory.mktemp("amun_snapshots")

    extracted = {}
    for zname in SNAPSHOT_ZIPS:
        _require_codec_for(zname)

        zpath = root / zname
        outdir = out_base / zname.replace(".zip", "")
        outdir.mkdir(parents=True, exist_ok=True)

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

        snapdir = outdir / "snapshot"
        if not snapdir.exists():
            pytest.skip(f"{zname} does not contain snapshot/ directory at expected path")
        if not (snapdir / "metadata.xml").exists():
            pytest.skip(f"{zname} snapshot/ missing metadata.xml")

        extracted[zname] = snapdir

    return extracted


def _load_amunxml(snapshot_dir: pathlib.Path):
    """
    Instantiate AmunXML against the extracted snapshot directory.
    """
    from amunpy.amunxml import AmunXML
    return AmunXML(str(snapshot_dir))


# ----------------------------
# Normalization helpers
# ----------------------------

def _normalize_attributes(attrs: dict) -> dict:
    """
    Remove attributes that may be environment- or path-dependent.
    Keep only primitive values.
    """
    drop = {
        # Add keys here if you find non-deterministic ones across variants
    }
    out = {}
    for k, v in attrs.items():
        if k in drop:
            continue
        # Keep only scalar-like (str/int/float/bool)
        if isinstance(v, (str, int, float, bool, np.integer, np.floating)):
            out[k] = v
    return out


def _compare_chunk_dicts(a_chunks: dict, b_chunks: dict):
    """
    Compare AMR chunk structure (levels/coords/bounds) between two snapshots.
    """
    assert set(a_chunks.keys()) == set(b_chunks.keys())

    for cn in a_chunks.keys():
        a = a_chunks[cn]
        b = b_chunks[cn]

        # dblocks should match
        assert a.get("dblocks") == b.get("dblocks")

        # levels might be list/array
        if "levels" in a and "levels" in b:
            np.testing.assert_array_equal(np.asarray(a["levels"]), np.asarray(b["levels"]))

        # coords/bounds should match if present
        if "coords" in a and "coords" in b:
            np.testing.assert_array_equal(np.asarray(a["coords"]), np.asarray(b["coords"]))

        if "bounds" in a and "bounds" in b:
            np.testing.assert_array_equal(np.asarray(a["bounds"]), np.asarray(b["bounds"]))


def _available_datasets(snap) -> set[str]:
    """
    Identify dataset names that AmunXML.dataset(...) should accept,
    based on snap.variables mapping.
    """
    return set(snap.variables.keys())


# ----------------------------
# The tests
# ----------------------------

def test_snapshots_parse_consistently(extracted_snapshots):
    """
    Ensures that XML parsing + metadata structures match across all compression variants.
    """
    base_name = "snapshot-none.zip"
    base = _load_amunxml(extracted_snapshots[base_name])

    base_attrs = _normalize_attributes(base.attributes)

    for name, path in extracted_snapshots.items():
        other = _load_amunxml(path)

        # Attributes: compare normalized subset
        other_attrs = _normalize_attributes(other.attributes)
        assert base_attrs == other_attrs, f"Attribute mismatch vs {base_name} for {name}"

        # Variables: identical dataset name set
        assert set(base.variables.keys()) == set(other.variables.keys()), f"Variable set mismatch for {name}"

        # Chunks: identical AMR structure
        _compare_chunk_dicts(base.chunks, other.chunks)


@pytest.mark.parametrize("field", ["density", "x-velocity", "y-velocity", "x-magnetic field", "y-magnetic field"])
def test_snapshots_data_fields_identical(extracted_snapshots, field):
    """
    Uses AmunXML.dataset() to exercise:
      - XML parsing of DataFile entries
      - compressor-specific decoding
      - optional encoder decoding
      - AMR assembly into uniform grid

    It compares the full uniform-grid array against snapshot-none baseline.
    """
    base_name = "snapshot-none.zip"
    base = _load_amunxml(extracted_snapshots[base_name])

    if field not in _available_datasets(base):
        pytest.skip(f"Baseline snapshot does not contain dataset '{field}'")

    base_arr = base.dataset(field)

    for name, path in extracted_snapshots.items():
        other = _load_amunxml(path)
        if field not in _available_datasets(other):
            pytest.skip(f"{name} does not contain dataset '{field}'")

        arr = other.dataset(field)

        # Prefer exact equality if dtype matches; otherwise fall back to allclose
        if arr.dtype == base_arr.dtype:
            np.testing.assert_array_equal(arr, base_arr, err_msg=f"Data mismatch for {field} in {name}")
        else:
            np.testing.assert_allclose(arr, base_arr, rtol=0.0, atol=0.0, err_msg=f"Data mismatch for {field} in {name}")


def test_snapshots_derived_velocity_identical(extracted_snapshots):
    """
    Optional but high-value: tests derived field computation path (velocity magnitude).
    This exercises __get_dataset__ computations and binary reads for velx/y/z.
    """
    base_name = "snapshot-none.zip"
    base = _load_amunxml(extracted_snapshots[base_name])

    if "velocity magnitude" not in _available_datasets(base):
        pytest.skip("Baseline snapshot does not define derived dataset 'velocity magnitude'")

    base_arr = base.dataset("velocity magnitude")

    for name, path in extracted_snapshots.items():
        other = _load_amunxml(path)
        if "velocity magnitude" not in _available_datasets(other):
            pytest.skip(f"{name} does not define derived dataset 'velocity magnitude'")

        arr = other.dataset("velocity magnitude")
        np.testing.assert_allclose(arr, base_arr, rtol=0.0, atol=0.0, err_msg=f"Derived 'velocity magnitude' mismatch in {name}")
