"""Reconstruct the exact merged StreamRender-H3 backbone.

After snapshot_download(..., local_dir=...), run this script in that download.
Other weights are used directly without conversion. Rebuilding consumes an
additional 66.28 GB of disk; sharded weights are retained.
"""
import hashlib
import json
from pathlib import Path
import struct


def reconstruct(root):
    backbone = root / "backbone"
    info = json.loads((backbone / "sharding.json").read_text())
    output = backbone / info["merged_filename"]
    if output.exists():
        raise RuntimeError(f"Output already exists; refusing to replace: {output}")
    original = (backbone / "original_header.bin").read_bytes()
    original_header = json.loads(original[8:])
    index = json.loads((backbone / "model.safetensors.index.json").read_text())["weight_map"]
    tensors = sorted(((k, v) for k, v in original_header.items() if k != "__metadata__"),
                     key=lambda item: item[1]["data_offsets"][0])
    readers, layouts = {}, {}
    temp = output.with_suffix(".incomplete")
    if temp.exists():
        raise RuntimeError(f"Incomplete output exists; inspect before retrying: {temp}")
    digest = hashlib.sha256(original)
    try:
        with temp.open("xb") as dst:
            dst.write(original)
            for name, tensor in tensors:
                shard = index[name]
                if shard not in readers:
                    src = (backbone / shard).open("rb")
                    n = struct.unpack("<Q", src.read(8))[0]
                    layouts[shard] = (8 + n, json.loads(src.read(n)))
                    readers[shard] = src
                base, header = layouts[shard]
                start, end = header[name]["data_offsets"]
                if header[name]["dtype"] != tensor["dtype"] or header[name]["shape"] != tensor["shape"]:
                    raise RuntimeError(f"Tensor metadata mismatch: {name}")
                src = readers[shard]
                src.seek(base + start)
                remaining = end - start
                while remaining:
                    data = src.read(min(16 * 1024 * 1024, remaining))
                    if not data:
                        raise EOFError(shard)
                    dst.write(data)
                    digest.update(data)
                    remaining -= len(data)
        if digest.hexdigest() != info["source_sha256"] or temp.stat().st_size != info["source_bytes"]:
            raise RuntimeError("Reconstructed checkpoint checksum mismatch")
        temp.replace(output)
    finally:
        for src in readers.values():
            src.close()
    print("BACKBONE_VERIFIED", digest.hexdigest())


if __name__ == "__main__":
    reconstruct(Path(__file__).resolve().parent)