#!/usr/bin/env python3
"""The libhdf5 side of the range-read measurement (docs/design/range-reads.md).

    libhdf5_reads.py FILE DATASET             # count libhdf5's reads
    libhdf5_reads.py FILE DATASET --extents   # print DATASET's stored extents

Counting: the file is opened through h5py's `fileobj` driver with a Python
file-like object that logs every `readinto`/`read` libhdf5 makes (this is how
h5py + fsspec reads remote files: each call becomes a range request unless
fsspec's own block cache absorbs it). The phases mirror range-trace: open,
list (visit every object; shape and dtype of every dataset), read (the whole
dataset). Default h5py settings (libhdf5's metadata cache, 64 KiB sieve
buffer, 1 MiB raw chunk cache) apply.

--extents prints `offset size` lines for the dataset's stored data (the
contiguous block, or every allocated chunk), the input range-trace uses to
split its read phase into metadata and raw data.
"""
import sys

import h5py


class LoggingFile:
    def __init__(self, path):
        self.f = open(path, "rb")
        self.pos = 0
        self.log = []

    def seek(self, off, whence=0):
        self.pos = self.f.seek(off, whence)
        return self.pos

    def tell(self):
        return self.pos

    def readinto(self, b):
        n = self.f.readinto(b)
        self.log.append((self.pos, n))
        self.pos += n
        return n

    def read(self, size=-1):
        data = self.f.read(size)
        self.log.append((self.pos, len(data)))
        self.pos += len(data)
        return data


def extents(path, name):
    with h5py.File(path, "r") as f:
        ds = f[name]
        if ds.chunks is None:
            off = ds.id.get_offset()
            if off is not None:
                print(off, ds.id.get_storage_size())
            return
        for i in range(ds.id.get_num_chunks()):
            info = ds.id.get_chunk_info(i)
            print(info.byte_offset, info.size)


def summarise(label, log):
    n = len(log)
    total = sum(s for _, s in log)
    distinct = len(set(log))
    print("| %s | %d | %d | %d |" % (label, n, distinct, total))


def count(path, name):
    lf = LoggingFile(path)
    f = h5py.File(lf, "r")
    opened = list(lf.log)
    lf.log.clear()

    def visit(_n, obj):
        if isinstance(obj, h5py.Dataset):
            obj.shape, obj.dtype

    f.visititems(visit)
    listed = list(lf.log)
    lf.log.clear()
    f[name][()]
    readlog = list(lf.log)
    f.close()
    print("| phase | read calls | distinct (offset, len) | bytes |")
    print("|---|---:|---:|---:|")
    summarise("open", opened)
    summarise("list", listed)
    summarise("read", readlog)
    summarise("open+list+read", opened + listed + readlog)


if __name__ == "__main__":
    if len(sys.argv) >= 4 and sys.argv[3] == "--extents":
        extents(sys.argv[1], sys.argv[2])
    else:
        count(sys.argv[1], sys.argv[2])