#!/usr/bin/env python3
"""Convert a checkpoint's expert tensors into the TITAN expert-contiguous layout: one 4 KiB-aligned blob per (layer, expert)
holding its tensors back to back, so that the SSD tier reads an expert with a single pread. Dense tensors stay in the
original safetensors. Output: <out_dir>/experts.titan (blob) + <out_dir>/experts.titan.json (index). Never modifies the source.
"""
import argparse, hashlib, json, os, sys, time
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from titan.registry import load_spec

ALIGN = 4096

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("model_dir")
    ap.add_argument("--out", default=None, help="default: ~/.cache/titan/<sha1 of realpath>/")
    a = ap.parse_args()
    spec = load_spec(a.model_dir)
    out = a.out or os.path.expanduser(f"~/.cache/titan/{hashlib.sha1(os.path.realpath(a.model_dir).encode()).hexdigest()[:12]}")
    os.makedirs(out, exist_ok=True)
    blob = os.path.join(out, "experts.titan")
    index = {"source": os.path.realpath(a.model_dir), "model_type": spec.model_type, "align": ALIGN, "experts": {}}
    fds = {}
    t0 = time.perf_counter(); written = 0
    with open(blob, "wb") as w:
        pos = 0
        for l in range(spec.num_layers):
            for e in range(spec.num_experts):
                es = spec.experts[(l, e)]
                parts = {}; off = 0
                chunks = []
                for name, loc in es.tensors.items():  # fixed order: as in the registry (gate, up, down × weight, scales, biases)
                    fd = fds.get(loc.file) or fds.setdefault(loc.file, os.open(loc.file, os.O_RDONLY))
                    b = os.pread(fd, loc.nbytes, loc.offset)
                    assert len(b) == loc.nbytes
                    parts[name] = {"off": off, "n": loc.nbytes, "dtype": loc.dtype, "shape": list(loc.shape)}
                    chunks.append(b); off += loc.nbytes
                data = b"".join(chunks)
                pad = (-len(data)) % ALIGN
                w.write(data); w.write(b"\0" * pad)
                index["experts"][f"{l},{e}"] = {"offset": pos, "nbytes": len(data), "tensors": parts}
                pos += len(data) + pad; written += len(data)
            print(f"layer {l} done, {pos/1e9:.2f} GB", file=sys.stderr)
    index["blob_bytes"] = pos
    json.dump(index, open(blob + ".json", "w"))
    print(json.dumps({"out": out, "experts": len(index["experts"]), "bytes": written, "blob_bytes": pos, "seconds": time.perf_counter() - t0}))

if __name__ == "__main__":
    main()