# SPDX-License-Identifier: CC0-1.0
"""Observe HDF5 file-driver reads, not physical disk I/O or elapsed-time speed."""
import argparse
import hashlib
import io
import json
import platform
import sys
import zlib
from datetime import datetime, timezone
from pathlib import Path

import h5py
import numpy as np

CHUNK = 32
SHAPE = (64, 64)
CHUNK_BYTES = CHUNK * CHUNK * 4
COORDS = ((0, 0), (0, 32), (32, 0), (32, 32))
CASES = (
    ("one-slot", 1, 4 * CHUNK_BYTES, "round-robin"),
    ("more-bytes-one-slot", 1, 16 * CHUNK_BYTES, "round-robin"),
    ("four-chunks-fit", 17, 4 * CHUNK_BYTES, "round-robin"),
    ("compressed-only-fits", 17, CHUNK_BYTES // 2, "round-robin"),
    ("one-chunk-fits", 17, CHUNK_BYTES, "round-robin"),
    ("grouped-one-slot", 1, CHUNK_BYTES, "grouped"),
)


class ReadTrace(io.BytesIO):
    def __init__(self, image):
        super().__init__(image)
        self.events = []

    def readinto(self, buffer):
        offset = self.tell()
        count = super().readinto(buffer)
        self.events.append({"offset": offset, "bytes": count})
        return count

    def read(self, size=-1):
        offset = self.tell()
        result = super().read(size)
        self.events.append({"offset": offset, "bytes": len(result)})
        return result


def raw_chunk_reads(events, chunks):
    """Reject ambiguous overlaps instead of silently calling them cache misses."""
    hits = []
    for event in events:
        offset, size = event["offset"], event["bytes"]
        if type(offset) is not int or type(size) is not int or offset < 0 or size < 0:
            raise ValueError("Invalid read extent")
        for index, chunk in enumerate(chunks):
            start, length = chunk["offset"], chunk["storedBytes"]
            if offset < start + length and start < offset + size:
                if offset != start or size != length:
                    raise ValueError("Read overlaps a chunk without matching its exact extent")
                hits.append(index)
    return hits


def fixture():
    payload = np.arange(SHAPE[0] * SHAPE[1], dtype="<u4").reshape(SHAPE)
    with io.BytesIO() as stream:
        with h5py.File(stream, "w", libver=("v110", "v110")) as handle:
            dataset = handle.create_dataset("values", data=payload, chunks=(CHUNK, CHUNK),
                                            compression="gzip", compression_opts=6, track_times=False)
            handle.flush()
            chunks = []
            for coord in COORDS:
                info = dataset.id.get_chunk_info_by_coord(coord)
                if info.filter_mask != 0 or not 0 < info.size < CHUNK_BYTES // 2:
                    raise RuntimeError("Fixture no longer has the required compressed chunk size")
                chunks.append({"coordinate": list(coord), "offset": info.byte_offset,
                               "storedBytes": info.size, "filterMask": info.filter_mask})
        image = stream.getvalue()
    return image, payload, chunks


def access_trace(order):
    if order not in ("round-robin", "grouped"):
        raise ValueError("Unknown access order")
    pairs = [(chunk, sample) for sample in range(8) for chunk in range(4)]
    if order == "grouped":
        pairs.sort()
    return [(row, col + sample) for chunk, sample in pairs for row, col in [COORDS[chunk]]]


def observe(image, payload, chunks, case):
    label, slots, capacity, order = case
    steps = []
    with ReadTrace(image) as stream:
        with h5py.File(stream, "r", rdcc_nslots=slots, rdcc_nbytes=capacity,
                       rdcc_w0=0.0) as handle:
            dataset = handle["values"]
            actual = dataset.id.get_access_plist().get_chunk_cache()
            if actual != (slots, capacity, 0.0):
                raise RuntimeError("Requested and effective cache parameters differ")
            for coord in access_trace(order):
                stream.events.clear()
                value = int(dataset[coord])
                if value != int(payload[coord]):
                    raise RuntimeError("Selected payload mismatch")
                events = list(stream.events)
                hits = raw_chunk_reads(events, chunks)
                expected_chunk = (coord[0] // CHUNK) * 2 + coord[1] // CHUNK
                if any(hit != expected_chunk for hit in hits) or len(hits) > 1:
                    raise RuntimeError("Unexpected chunk read for scalar selection")
                steps.append({"coordinate": list(coord), "value": value,
                              "rawChunkReads": hits, "driverReads": events})
    reads = sum(len(step["rawChunkReads"]) for step in steps)
    raw_bytes = sum(chunks[index]["storedBytes"] for step in steps
                    for index in step["rawChunkReads"])
    return {"case": label, "order": order, "rdccNslots": slots, "rdccNbytes": capacity,
            "rdccW0": 0.0, "scalarSelections": len(steps), "selectedBytes": len(steps) * 4,
            "rawChunkReads": reads, "rawChunkBytes": raw_bytes,
            "chunkDecodeBytesInferred": reads * CHUNK_BYTES,
            "steps": steps}


def run(repetitions=3):
    if type(repetitions) is not int or not 1 <= repetitions <= 20:
        raise ValueError("repetitions must be an integer from 1 to 20")
    image, payload, chunks = fixture()
    rows = []
    for repetition in range(1, repetitions + 1):
        for case in CASES:
            rows.append({"repetition": repetition, **observe(image, payload, chunks, case)})
    return {
        "schemaVersion": 1, "generatedAt": datetime.now(timezone.utc).isoformat(),
        "evidenceType": "measured-library-reads", "license": "CC0-1.0",
        "environment": {"python": platform.python_version(), "h5py": h5py.__version__,
                        "hdf5": h5py.version.hdf5_version, "numpy": np.__version__,
                        "zlib": zlib.ZLIB_RUNTIME_VERSION, "system": platform.system(),
                        "machine": platform.machine(), "driver": "fileobj over BytesIO",
                        "formatBounds": ["v110", "v110"]},
        "fixture": {"shape": list(SHAPE), "chunkShape": [CHUNK, CHUNK],
                    "dtype": "<u4", "chunkBytes": CHUNK_BYTES, "compression": "gzip",
                    "compressionLevel": 6, "chunks": chunks,
                    "payloadSha256": hashlib.sha256(payload.tobytes()).hexdigest()},
        "repetitions": repetitions, "observations": rows,
        "sourceSha256": hashlib.sha256(Path(__file__).read_bytes()).hexdigest(),
        "requirementsSha256": hashlib.sha256(Path(__file__).with_name(
            "hdf5-chunk-cache-requirements.txt").read_bytes()).hexdigest(),
        "method": "Fresh HDF5 open per case over the same synthetic in-memory file. Trace read/readinto calls during 32 scalar selections; match exact stored chunk extents; retain metadata reads separately in each step.",
        "limits": [
            "One library build and synthetic trace, not a workload sample or paper replication.",
            "No disk, OS page-cache, latency, throughput, power-loss or device-failure measurement.",
            "Decoded bytes are inferred as raw chunk reads times uncompressed chunk size; filter callbacks were not instrumented.",
            "One slot deliberately forces conflicts; 17 slots is a measured control for this layout, not a universal optimum."
        ]
    }


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--repetitions", type=int, default=3)
    args = parser.parse_args()
    json.dump(run(args.repetitions), sys.stdout, indent=2)
    sys.stdout.write("\n")