"""Controlled Zarr store-key versus inner-chunk acceptance experiment."""
import argparse
import hashlib
import importlib.metadata
import json
from pathlib import Path
import platform
import struct
import tempfile

import numpy as np
import zarr
from zarr.codecs import BytesCodec, Crc32cCodec, ShardingCodec

VALUES = list(range(10, 18))
FILL = -999
LAYOUTS = ("plain", "shard-start", "shard-end")
BASE_STATES = ("complete", "sparse-tail", "deleted", "fill-elided", "fill-stored", "corrupt-payload")
SELECTORS = {"full": slice(None), "prefix": slice(0, 4), "tail": slice(4, 8)}
PACKAGES = ("zarr", "numpy", "numcodecs", "google-crc32c", "donfig", "PyYAML", "packaging", "typing_extensions", "msgspec")


def require(condition, message):
    if not condition:
        raise ValueError(message)


def crc32c(data):
    crc = 0xFFFFFFFF
    for byte in data:
        crc ^= byte
        for _ in range(8):
            crc = (crc >> 1) ^ (0x82F63B78 if crc & 1 else 0)
    return crc ^ 0xFFFFFFFF


def inventory(root):
    return {str(p.relative_to(root)): hashlib.sha256(p.read_bytes()).hexdigest()
            for p in sorted(root.rglob("*")) if p.is_file()}


def states(layout):
    return BASE_STATES + (("swapped",) if layout == "plain" else ("corrupt-index", "stale-sparse-shard"))


def create(root, layout, write_empty=False, strict=False):
    inner = [BytesCodec(endian="little"), Crc32cCodec()]
    codecs = inner if layout == "plain" else [ShardingCodec(
        chunk_shape=(4,), codecs=inner,
        index_codecs=[BytesCodec(endian="little"), Crc32cCodec()],
        index_location=layout.split("-")[1])]
    return zarr.create(shape=(8,), chunks=(4,) if layout == "plain" else (8,),
                       dtype="int32", fill_value=FILL, store=str(root),
                       zarr_format=3, codecs=codecs,
                       config={"write_empty_chunks": write_empty, "read_missing_chunks": not strict})


def parse_shard(data, layout):
    size = 2 * 2 * 8 + 4
    raw = data[:size] if layout == "shard-start" else data[-size:]
    require(len(raw) == size, "truncated index")
    pairs = [list(struct.unpack_from("<QQ", raw, i * 16)) for i in range(2)]
    missing = 2**64 - 1
    valid = []
    for offset, length in pairs:
        if offset == missing and length == missing:
            valid.append(None)
        else:
            require(offset + length <= len(data), "index outside container")
            chunk = data[offset:offset + length]
            require(len(chunk) >= 4, "truncated payload")
            valid.append(crc32c(chunk[:-4]) == int.from_bytes(chunk[-4:], "little"))
    return {"entries": pairs,
            "index_crc_valid": crc32c(raw[:-4]) == int.from_bytes(raw[-4:], "little"),
            "payload_crc_valid": valid}


def build(root, layout, state):
    arr = create(root, layout, write_empty=state == "fill-stored")
    # Retain an actual earlier valid generation for the rollback control.
    arr[:4] = np.asarray(VALUES[:4], dtype="int32")
    old_shard = (root / "c/0").read_bytes() if layout != "plain" else None
    if state in ("fill-elided", "fill-stored"):
        arr[4:] = FILL
    elif state != "sparse-tail":
        arr[4:] = np.asarray(VALUES[4:], dtype="int32")
    before = inventory(root)
    key = "c/1" if layout == "plain" else "c/0"
    target = root / key
    if state == "deleted":
        target.unlink()
    elif state == "corrupt-payload":
        data = bytearray(target.read_bytes())
        offset = 0 if layout == "plain" else parse_shard(data, layout)["entries"][1][0]
        data[offset] ^= 1
        target.write_bytes(data)
    elif state == "corrupt-index":
        data = bytearray(target.read_bytes())
        index_start = 0 if layout == "shard-start" else len(data) - 36
        data[index_start + 32] ^= 1
        target.write_bytes(data)
    elif state == "stale-sparse-shard":
        target.write_bytes(old_shard)
    elif state == "swapped":
        first = (root / "c/0").read_bytes()
        second = (root / "c/1").read_bytes()
        (root / "c/0").write_bytes(second)
        (root / "c/1").write_bytes(first)
    after = inventory(root)
    meta = json.loads((root / "zarr.json").read_text())
    info = {"layout": layout, "state": state, "before": before, "after": after,
            "changed_keys": sorted(k for k in set(before) | set(after) if before.get(k) != after.get(k)),
            "metadata": meta}
    if layout != "plain" and target.exists():
        info["shard"] = parse_shard(target.read_bytes(), layout)
    return info


def read(root, strict, selection):
    array = zarr.open_array(str(root), mode="r").with_config({"read_missing_chunks": not strict})
    require(array.config.read_missing_chunks is not strict, "read policy was not applied")
    try:
        values = array[SELECTORS[selection]]
        return {"status": "ok", "values": values.tolist(),
                "sha256_int32_le": hashlib.sha256(values.astype("<i4").tobytes()).hexdigest(),
                "mean": float(values.mean())}
    except Exception as error:
        # Preserve exact class and a path-independent message; the oracle rejects unexpected errors.
        return {"status": "error", "error_type": type(error).__name__,
                "message": str(error).replace(str(root), "<store>")}


def expected(layout, state, strict, selection):
    values = VALUES[:]
    if state in ("sparse-tail", "fill-elided", "fill-stored", "stale-sparse-shard"):
        values[4:] = [FILL] * 4
    if state == "deleted":
        values = VALUES[:4] + [FILL] * 4 if layout == "plain" else [FILL] * 8
    if state == "swapped":
        values = VALUES[4:] + VALUES[:4]
    missing = state == "deleted" or (layout == "plain" and state in ("sparse-tail", "fill-elided"))
    touches_missing = layout != "plain" or selection != "prefix"
    if strict and missing and touches_missing:
        return "ChunkNotFoundError", None
    if state == "corrupt-index" or (state == "corrupt-payload" and selection != "prefix"):
        return "ValueError", None
    return None, values[SELECTORS[selection]]


def check_row(row):
    error, values = expected(row["layout"], row["state"], row["strict"], row["selection"])
    result = row["result"]
    if error:
        require(result["status"] == "error" and result.get("error_type") == error,
                f"unexpected error outcome: {row}")
        if error == "ValueError":
            require("checksum do not match" in result["message"], "not a checksum rejection")
    else:
        require(result["status"] == "ok" and result.get("values") == values,
                f"unexpected values: {row}")
        require(result["mean"] == sum(values) / len(values), "aggregate mismatch")
        digest = hashlib.sha256(np.asarray(values, dtype="<i4").tobytes()).hexdigest()
        require(result["sha256_int32_le"] == digest, "value hash mismatch")


def validate(report):
    require(report["schema_version"] == 1, "wrong schema")
    require(report["environment"]["packages"]["zarr"] == "3.4.1", "untested version")
    combos = {(l, s, strict, selection) for l in LAYOUTS for s in states(l)
              for strict in (False, True) for selection in SELECTORS}
    rows = report["observations"]
    keys = [(r["layout"], r["state"], r["strict"], r["selection"]) for r in rows]
    require(len(keys) == len(set(keys)) and set(keys) == combos, "missing/duplicate/extra rows")
    for row in rows:
        check_row(row)
    fixtures = {(r["layout"], r["state"]): r for r in report["fixtures"]}
    require(len(fixtures) == len(report["fixtures"]) == 23, "fixture inventory mismatch")
    for layout in LAYOUTS:
        for state in states(layout):
            fixture = fixtures[layout, state]
            require(fixture["before"]["zarr.json"] == fixture["after"]["zarr.json"], "metadata mutated")
        sparse = fixtures[layout, "sparse-tail"]["after"]
        require(sparse == fixtures[layout, "fill-elided"]["after"], "elision differs")
        if layout == "plain":
            require(sparse == fixtures[layout, "deleted"]["after"], "different absence histories")
        else:
            require(sparse == fixtures[layout, "stale-sparse-shard"]["after"], "stale history differs")
            require(fixtures[layout, "stale-sparse-shard"]["shard"]["index_crc_valid"], "invalid rollback")
            require(fixtures[layout, "sparse-tail"]["shard"]["payload_crc_valid"] == [True, None], "sparse checksum")
            require(fixtures[layout, "fill-stored"]["shard"]["payload_crc_valid"] == [True, True], "stored fill")
            require(fixtures[layout, "corrupt-payload"]["shard"]["payload_crc_valid"] == [True, False], "bad corruption control")
            require(fixtures[layout, "corrupt-index"]["shard"]["index_crc_valid"] is False, "bad index control")
    controls = report["reopen_controls"]
    require(len(controls) == 1, "reopen control missing")
    require(controls[0] == {"created_strict_error": "ChunkNotFoundError",
                           "default_reopen_values": VALUES[:4] + [FILL] * 4,
                           "policy_in_metadata": False}, "policy persisted unexpectedly")


def run():
    fixtures, observations = [], []
    with tempfile.TemporaryDirectory(prefix="hesela-zarr-strict-") as temp:
        for layout in LAYOUTS:
            for state in states(layout):
                root = Path(temp) / layout / state
                fixtures.append(build(root, layout, state))
                for strict in (False, True):
                    for selection in SELECTORS:
                        row = {"layout": layout, "state": state, "strict": strict,
                               "selection": selection, "result": read(root, strict, selection)}
                        observations.append(row)
        root = Path(temp) / "runtime-policy"
        array = create(root, "plain", strict=True)
        array[:4] = np.asarray(VALUES[:4], dtype="int32")
        error_type = None
        try:
            array[:]
        except Exception as error:
            error_type = type(error).__name__
        control = {"created_strict_error": error_type,
                   "default_reopen_values": zarr.open_array(str(root), mode="r")[:].tolist(),
                   "policy_in_metadata": "read_missing_chunks" in (root / "zarr.json").read_text()}
    report = {"schema_version": 1, "experiment_date": "2026-10-23",
              "environment": {"python": platform.python_version(), "system": platform.system(),
                              "release": platform.release(), "machine": platform.machine(),
                              "packages": {p: importlib.metadata.version(p) for p in PACKAGES}},
              "fixtures": fixtures, "observations": observations, "reopen_controls": [control]}
    validate(report)
    return report


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--output", type=Path, required=True)
    args = parser.parse_args()
    report = run()
    args.output.write_text(json.dumps(report, indent=2, sort_keys=True) + "\n")
    print(f"validated {len(report['fixtures'])} fixtures, {len(report['observations'])} reads, one reopen control")