#!/usr/bin/env python3
"""Parse a DCPM metadata.bin index. Stage 1: read only, write nothing.

Record layout, decoded from the bytes:
    int64 name_len | name | int64 ndim | int64 shape[ndim]
    int32 dtype | int32 ? | int64 ? | int32 ? | int64 file_len | file
    int64 offset  | ... trailing fields
We locate records by scanning for the length-prefixed names rather than
assuming a fixed stride, so a variable-length shape list cannot desync us.
"""
import struct, sys, re, json

def parse(path):
    b = open(path, "rb").read()
    assert b[:4] == b"DCPM", f"not DCPM: {b[:4]!r}"
    recs = []
    for m in re.finditer(rb"param_(\d{6})", b):
        i = m.start()
        nlen = struct.unpack_from("<q", b, i - 8)[0]
        if nlen != m.end() - i:            # name must be length-prefixed
            continue
        j = m.end()
        ndim = struct.unpack_from("<q", b, j)[0]; j += 8
        if not (0 <= ndim <= 8): continue
        shape = list(struct.unpack_from("<" + "q" * ndim, b, j)); j += 8 * ndim
        dtype = struct.unpack_from("<i", b, j)[0]; j += 4
        f1    = struct.unpack_from("<i", b, j)[0]; j += 4
        f2    = struct.unpack_from("<q", b, j)[0]; j += 8
        f3    = struct.unpack_from("<i", b, j)[0]; j += 4
        flen  = struct.unpack_from("<q", b, j)[0]; j += 8
        fname = b[j:j + flen].decode(errors="replace"); j += flen
        off   = struct.unpack_from("<q", b, j)[0]; j += 8
        recs.append(dict(idx=int(m.group(1)), name=m.group(0).decode(), ndim=ndim,
                         shape=shape, dtype=dtype, f1=f1, f2=f2, f3=f3,
                         file=fname, offset=off))
    recs.sort(key=lambda r: r["idx"])
    return recs

if __name__ == "__main__":
    recs = parse(sys.argv[1])
    print(f"  parsed {len(recs)} param records")
    dts = {}
    for r in recs: dts.setdefault(r["dtype"], 0); dts[r["dtype"]] += 1
    print(f"  dtype codes seen: {dts}")
    files = {}
    for r in recs: files.setdefault(r["file"], 0); files[r["file"]] += 1
    print(f"  source files: {files}")
    print(f"\n  {'idx':>4} {'shape':<20} {'dtype':>5} {'file':<16} {'offset':>14}")
    for r in recs[:14]:
        print(f"  {r['idx']:>4} {str(r['shape']):<20} {r['dtype']:>5} {r['file']:<16} {r['offset']:>14,}")
    print("   ...")
    for r in recs[-6:]:
        print(f"  {r['idx']:>4} {str(r['shape']):<20} {r['dtype']:>5} {r['file']:<16} {r['offset']:>14,}")
    tot = sum(int(__import__("math").prod(r["shape"])) for r in recs)
    print(f"\n  total elements across all params: {tot:,}")
    json.dump(recs, open("/tmp/dcpm_index.json", "w"))
    print("  wrote /tmp/dcpm_index.json")