"""Write HDF5 and NetCDF-4 test files with h5py/netCDF4, and what libhdf5
reads back from them, for the clawhdf5-wasm tests.

    python make_fixture.py OUT_DIR

writes OUT_DIR/fixture.h5, OUT_DIR/fixture.nc and OUT_DIR/expected.json.
Both the Rust test (crates/clawhdf5-wasm/tests/h5py_interop.rs, native) and
the Node test (test.mjs, the built wasm package) compare against the same
expected.json, so the two check the same values.

Every expected value comes from h5py reading the file back (numpy slicing
for hyperslabs), never from the arrays that were written. Integers are
encoded as strings so JSON.parse keeps 64-bit values exact.
"""

import json
import sys
import warnings
from pathlib import Path

import h5py
import netCDF4
import numpy as np

# netCDF4 1.7 trips numpy 2.5's shape-setting deprecation on assignment.
warnings.filterwarnings("ignore", category=DeprecationWarning)

out = Path(sys.argv[1])
out.mkdir(parents=True, exist_ok=True)
h5 = out / "fixture.h5"
nc = out / "fixture.nc"

rng = np.random.default_rng(7)
with h5py.File(h5, "w") as f:
    f.attrs["title"] = "wasm fixture"
    f.attrs["version"] = np.int64(3)
    f.attrs["scale"] = np.array([0.5, 2.0])
    f.attrs["big"] = np.uint64(2**63 + 5)
    f.attrs.create("vlen_note", "héllo", dtype=h5py.string_dtype())
    f.create_dataset(
        "grid", data=np.arange(60, dtype="<f8").reshape(6, 10) / 4,
        chunks=(4, 3), compression="gzip", shuffle=True,
    )
    f.create_dataset("f32_be", data=rng.standard_normal(7).astype(">f4"))
    f.create_dataset("f16", data=np.array([0.5, -1.25, 65504], dtype="<f2"))
    f.create_dataset("i8", data=np.array([-128, -1, 0, 127], dtype="i1"))
    f.create_dataset("i16_be", data=np.array([-32768, 5, 32767], dtype=">i2"))
    f.create_dataset("u16", data=np.array([0, 40000, 65535], dtype="<u2"))
    f.create_dataset("u32", data=np.array([0, 4_000_000_000], dtype="<u4"))
    f.create_dataset("i64", data=np.array([-(2**63), 2**53 + 1, 7], dtype="<i8"))
    f.create_dataset("u64", data=np.array([2**64 - 1, 1], dtype="<u8"))
    f.create_dataset("scalar", data=np.float64(2.5))
    f.create_dataset("fixed_str", data=np.array([b"alpha", b"be", b""], dtype="S5"))
    f.create_dataset(
        "vlen_str", data=["one", "двa", ""], dtype=h5py.string_dtype()
    )
    f.create_dataset("flags", data=np.array([True, False, True]))
    # An array datatype: four elements, each an i4[2].
    pairs = f.create_dataset("pairs", shape=(4,), dtype=np.dtype(("<i4", (2,))))
    pairs[...] = np.arange(8, dtype="<i4").reshape(4, 2)
    f.create_dataset(
        "cube", data=np.arange(2 * 5 * 6, dtype="<i4").reshape(2, 5, 6),
        chunks=(1, 2, 3), compression="gzip",
    )
    comp = np.zeros(2, dtype=[("x", "<f8"), ("n", "<i4")])
    f.create_dataset("table", data=comp)
    g = f.create_group("sensors")
    g.attrs["location"] = "lab"
    g.create_dataset("temp", data=np.array([21.5, 22.0, 22.25], dtype="<f4"))
    g.create_group("empty")
    f["alias"] = h5py.SoftLink("/sensors/temp")

with netCDF4.Dataset(nc, "w") as d:
    d.title = "nc fixture"
    d.createDimension("time", None)
    d.createDimension("x", 4)
    t = d.createVariable("time", "f8", ("time",))
    t.units = "days since 2000-01-01"
    v = d.createVariable("temp", "f4", ("time", "x"), zlib=True)
    t[:] = np.arange(3)
    v[:] = np.arange(12, dtype="f4").reshape(3, 4) + 0.5


def kind(dt):
    """The typed-array kind clawhdf5-wasm returns for a numpy dtype."""
    if dt.kind == "b" or h5py.check_enum_dtype(dt) is not None:
        return "strings"
    if h5py.check_string_dtype(dt) is not None or dt.kind == "S":
        return "strings"
    if dt.subdtype is not None:
        return kind(dt.subdtype[0])
    if dt.kind == "f":
        return "f64" if dt.itemsize == 8 else "f32"
    if dt.kind in "iu":
        return f"{dt.kind}{dt.itemsize * 8}"
    raise ValueError(dt)


def flat(a, k):
    a = np.asarray(a)
    if k == "strings":
        if a.dtype.kind == "b":
            return ["TRUE" if x else "FALSE" for x in a.ravel()]
        return [x.decode() if isinstance(x, bytes) else str(x) for x in a.ravel()]
    if k.startswith(("i", "u")):
        return [str(int(x)) for x in a.ravel()]
    return [float(x) for x in a.ravel()]


def entry(ds, slab=None):
    k = kind(ds.dtype)
    data = ds[()]
    e = {"kind": k, "shape": list(np.shape(data)), "values": flat(data, k)}
    if slab:
        start, count, stride = slab
        idx = tuple(slice(s, s + (c - 1) * st + 1, st)
                    for s, c, st in zip(start, count, stride))
        part = ds[idx]
        e["slab"] = {"start": start, "count": count, "stride": stride,
                     "shape": list(part.shape), "values": flat(part, k)}
    return e


def attr(v):
    v = np.asarray(v) if not isinstance(v, (str, bytes)) else v
    if isinstance(v, bytes):
        return {"string": v.decode()}
    if isinstance(v, str):
        return {"string": v}
    if v.dtype.kind in "iu":
        return {"int": [str(int(x)) for x in v.ravel()], "scalar": v.ndim == 0}
    if v.dtype.kind == "f":
        return {"float": [float(x) for x in v.ravel()], "scalar": v.ndim == 0}
    if v.dtype.kind in "OSU":
        items = [x.decode() if isinstance(x, bytes) else str(x) for x in v.ravel()]
        return {"string": items[0]} if v.ndim == 0 else {"strings": items}
    raise ValueError(v.dtype)


# Attributes the reader returns but the comparison leaves out: netCDF-4's
# internal ones (a leading underscore), and dimension-scale bookkeeping.
SKIP_ATTRS = ["DIMENSION_LIST", "REFERENCE_LIST", "CLASS", "NAME"]


def describe(path):
    expected = {"datasets": {}, "errors": {}, "attrs": {}, "lists": {},
                "skip_attrs": SKIP_ATTRS}
    with h5py.File(path, "r") as f:
        def walk(key, obj):
            expected["attrs"][key] = {
                k: attr(obj.attrs[k]) for k in obj.attrs
                if not k.startswith("_") and k not in SKIP_ATTRS}
            if isinstance(obj, h5py.Group):
                members = {n: obj.get(n) for n in obj}
                groups = [n for n, o in members.items() if isinstance(o, h5py.Group)]
                sets = [n for n, o in members.items() if isinstance(o, h5py.Dataset)]
                expected["lists"][key] = {"groups": sorted(groups),
                                          "datasets": sorted(sets)}
                for n, o in members.items():
                    walk(key.rstrip("/") + "/" + n, o)
            elif obj.dtype.names:
                expected["errors"][key] = "compound"
            else:
                expected["datasets"][key] = entry(obj, slab_for(obj))

        walk("/", f)
    return expected


def slab_for(obj):
    """A strided hyperslab inside the dataset's extent, or None."""
    if obj.ndim == 0 or obj.shape[0] < 2 or 0 in obj.shape:
        return None
    start = [1] + [0] * (obj.ndim - 1)
    count = [max(1, (obj.shape[0] - 1) // 2)] + [
        max(1, (n + 1) // 2) for n in obj.shape[1:]]
    stride = [2 if obj.shape[0] > 2 else 1] + [2] * (obj.ndim - 1)
    count = [min(c, (n - s - 1) // st + 1)
             for c, s, st, n in zip(count, start, stride, obj.shape)]
    return (start, count, stride)


json.dump({"fixture.h5": describe(h5), "fixture.nc": describe(nc)},
          open(out / "expected.json", "w"), indent=1, ensure_ascii=False)