"""Write the variable-length data files the h5rs VL tests run on.

usage: gen_vl_files.py OUTDIR

For 8-byte (`vl8`) and 4-byte (`vl4`) offsets, writes OUTDIR/vl8.h5 and
OUTDIR/vl4.h5, which libhdf5 reads in full, and OUTDIR/bad8.h5 and
OUTDIR/bad4.h5, whose `bad` and `badseq` elements 0 have a length that
disagrees with their global heap object (libhdf5: "Expected global heap
object size does not match"), and whose `undef` element 1 has length 0 and
the undefined heap address (libhdf5: "addr undefined"). h5py cannot write a VL string with a NUL in
it or a null element in a contiguous dataset, so those are patched in.

Prints one JSON object: for each file, each dataset's values as h5py reads
them one element at a time (strings as text, sequences as lists, a compound
as a list of its fields), with null for an element h5py cannot read; the
root attribute `va`; and the addresses of the `bad` elements' collections.
"""

import json
import os
import struct
import sys

import h5py
import numpy as np

out = sys.argv[1]
S = h5py.string_dtype("utf-8")
I4 = h5py.vlen_dtype(np.dtype("<i4"))


def create(path, sizes):
    if sizes is None:
        return h5py.File(path, "w")
    fcpl = h5py.h5p.create(h5py.h5p.FILE_CREATE)
    fcpl.set_sizes(*sizes)
    return h5py.File(h5py.h5f.create(path.encode(), h5py.h5f.ACC_TRUNC, fcpl=fcpl))


def element(length, addr, index, os_):
    return struct.pack("<I", length) + addr.to_bytes(os_, "little") + struct.pack("<I", index)


def good(path, sizes):
    os_ = 8 if sizes is None else sizes[0]
    with create(path, sizes) as f:
        f.create_dataset(
            "d", data=np.array(["aXb", "", "ok", "zz", "hello"], dtype=object), dtype=S
        )
        u = f.create_dataset("u", shape=(4,), dtype=S, chunks=(1,))
        u[1] = "w"
        s = f.create_dataset("seq", shape=(3,), dtype=I4)
        s[0] = [1, 2, 3]
        s[1] = []
        s[2] = [-5]
        s = f.create_dataset("sequ", shape=(3,), dtype=I4, chunks=(1,))
        s[0] = [7, 8]
        ct = np.dtype([("id", "<i4"), ("name", S)])
        arr = np.zeros(3, dtype=ct)
        arr["id"] = [1, 2, 3]
        arr["name"] = ["one", "", "three"]
        f.create_dataset("cmp", data=arr)
        f.attrs.create("va", np.array(["p", "", "q"], dtype=object), dtype=S)
        off = f["d"].id.get_offset()
    b = bytearray(open(path, "rb").read())
    i = b.index(b"aXb")
    b[i + 1] = 0  # "a\0b"
    es = 8 + os_
    b[off + 2 * es : off + 3 * es] = element(2, 0, 1, os_)  # "ok" -> null
    open(path, "wb").write(bytes(b))


def bad(path, sizes):
    os_ = 8 if sizes is None else sizes[0]
    with create(path, sizes) as f:
        f.create_dataset("bad", data=np.array(["cdefgh", "ok"], dtype=object), dtype=S)
        s = f.create_dataset("badseq", shape=(2,), dtype=I4)
        s[0] = [1, 2, 3]
        s[1] = [4]
        f.create_dataset("undef", data=np.array(["x", "", "yz"], dtype=object), dtype=S)
        off, soff = f["bad"].id.get_offset(), f["badseq"].id.get_offset()
        uoff = f["undef"].id.get_offset()
    b = bytearray(open(path, "rb").read())
    gcol = int.from_bytes(b[off + 4 : off + 4 + os_], "little")
    struct.pack_into("<I", b, off, 3)  # "cdefgh": length 6 -> 3
    struct.pack_into("<I", b, soff, 2)  # [1, 2, 3]: length 3 -> 2
    # "": length 0 at the undefined address (all 0xff), which libhdf5 fails
    # to read ("addr undefined"); it writes a null element as address 0.
    es = 8 + os_
    b[uoff + es : uoff + 2 * es] = element(0, (1 << (8 * os_)) - 1, 1, os_)
    open(path, "wb").write(bytes(b))
    return gcol


def value(v):
    if isinstance(v, bytes):
        return v.decode()
    if isinstance(v, str):
        return v
    if isinstance(v, np.void):
        return [value(x) for x in v]
    if isinstance(v, np.ndarray):
        return [value(x) for x in v]
    return v.item() if hasattr(v, "item") else v


def read(ds):
    got = []
    for i in range(ds.shape[0]):
        try:
            got.append(value(ds[i]))
        except OSError:
            got.append(None)
    return got


result = {}
for tag, sizes in (("8", None), ("4", (4, 4))):
    g, x = os.path.join(out, f"vl{tag}.h5"), os.path.join(out, f"bad{tag}.h5")
    good(g, sizes)
    gcol = bad(x, sizes)
    with h5py.File(g, "r") as f:
        result[f"vl{tag}"] = {n: read(f[n]) for n in ("d", "u", "seq", "sequ", "cmp")}
        result[f"vl{tag}"]["va"] = [value(s) for s in f.attrs["va"]]
    with h5py.File(x, "r") as f:
        result[f"bad{tag}"] = {n: read(f[n]) for n in ("bad", "badseq", "undef")}
        result[f"bad{tag}"]["gcol"] = gcol
json.dump(result, sys.stdout)