#!/usr/bin/env python3
"""Apply (or undo) the pqhaz heretic patch IN PLACE in the official DeepSeek-V4.1-Flash-1017 shards, for engines that
read attention weights straight from the official checkpoint (Morrowmake). Only the 38 `layers.N.attn.wo_b.weight`
tensors change (same dtype F8_E4M3, same shape, same offsets); FP8 scales are untouched, as the patch was built.

  heretic_inplace.py apply   OFFICIAL_DIR PATCH.safetensors BACKUP.safetensors
  heretic_inplace.py restore OFFICIAL_DIR PATCH.safetensors BACKUP.safetensors
  heretic_inplace.py status  OFFICIAL_DIR PATCH.safetensors BACKUP.safetensors

apply: writes an exact byte backup of every original tensor first (a safetensors file with the same names), then
overwrites, fsyncs and reads back; refuses a shard with more than one hard link. restore: writes the backup back.
Idempotent: a tensor already in the wanted state is skipped. Stdlib only."""
import hashlib, json, os, struct, sys


def header(path):
    with open(path, "rb") as f:
        n = struct.unpack("<Q", f.read(8))[0]
        h = json.loads(f.read(n))
    return h, 8 + n


def region(path, name):
    h, base = header(path)
    t = h[name]
    a, b = t["data_offsets"]
    return t["dtype"], t["shape"], base + a, b - a


def read_at(path, off, size):
    with open(path, "rb") as f:
        f.seek(off)
        return f.read(size)


def sha(b):
    return hashlib.sha256(b).hexdigest()


def write_safetensors(path, tensors):  # tensors: name -> (dtype, shape, bytes)
    hdr, off = {}, 0
    for k in sorted(tensors):
        dt, shp, data = tensors[k]
        hdr[k] = {"dtype": dt, "shape": shp, "data_offsets": [off, off + len(data)]}
        off += len(data)
    hb = json.dumps(hdr, separators=(",", ":")).encode()
    hb += b" " * ((8 - len(hb) % 8) % 8)
    tmp = path + ".tmp"
    with open(tmp, "wb") as f:
        f.write(struct.pack("<Q", len(hb)) + hb)
        for k in sorted(tensors):
            f.write(tensors[k][2])
        f.flush(); os.fsync(f.fileno())
    os.replace(tmp, path)


def main():
    mode, odir, patch, backup = sys.argv[1:5]
    wm = json.load(open(os.path.join(odir, "model.safetensors.index.json")))["weight_map"]
    ph, _ = header(patch)
    names = sorted(k for k in ph if k != "__metadata__")
    plan = []
    for k in names:
        shard = os.path.join(odir, wm[k])
        dt, shp, off, size = region(shard, k)
        pdt, pshp, poff, psize = region(patch, k)
        assert (dt, shp, size) == (pdt, pshp, psize), (k, dt, shp, size, pdt, pshp, psize)
        plan.append((k, shard, off, size, poff))
    print(f"{len(plan)} tensors in {len({p[1] for p in plan})} shards")

    if mode == "status":
        bh = header(backup)[0] if os.path.exists(backup) else {}
        n_patched = n_stock = n_other = 0
        for k, shard, off, size, poff in plan:
            cur = sha(read_at(shard, off, size))
            if cur == sha(read_at(patch, poff, size)): n_patched += 1
            elif bh and cur == sha(read_at(backup, region(backup, k)[2], size)): n_stock += 1
            else: n_other += 1
        print(f"heretic {n_patched}, stock (matches backup) {n_stock}, other {n_other}")
        return

    if mode == "apply":
        for k, shard, *_ in plan:
            assert os.stat(shard).st_nlink == 1, f"{shard} has other hard links; refusing to write in place"
        if not os.path.exists(backup):
            orig = {}
            for k, shard, off, size, poff in plan:
                data = read_at(shard, off, size)
                if data == read_at(patch, poff, size):
                    sys.exit(f"{k} is already patched and there is no backup: refusing (stock bytes unknown)")
                orig[k] = (region(shard, k)[0], region(shard, k)[1], data)
            write_safetensors(backup, orig)
            print(f"backup written: {backup} ({os.path.getsize(backup) / 2**30:.2f} GiB)")
        src = lambda k, poff, size: read_at(patch, poff, size)
    elif mode == "restore":
        assert os.path.exists(backup), "no backup"
        src = lambda k, poff, size: read_at(backup, region(backup, k)[2], size)
    else:
        sys.exit("mode: apply | restore | status")

    done = skipped = 0
    for k, shard, off, size, poff in plan:
        want = src(k, poff, size)
        if read_at(shard, off, size) == want:
            skipped += 1; continue
        fd = os.open(shard, os.O_RDWR)
        try:
            assert os.pwrite(fd, want, off) == size
            os.fsync(fd)
        finally:
            os.close(fd)
        assert sha(read_at(shard, off, size)) == sha(want), f"read-back mismatch on {k}"
        done += 1
    print(f"{mode}: {done} written, {skipped} already in place, all read back OK")


if __name__ == "__main__":
    main()