"""
patch_int8_add_format.py
Adds 'format': 'int8_tensorwise' to each per-layer comfy_quant config in an
INT8-Fast-saved safetensors so ComfyUI's stock Load Diffusion Model (PR #14636)
can dispatch it correctly. Writes a sibling file with '_fixed' appended to the
basename, preserving the original.
"""

import json
import os
import sys
import torch
from safetensors import safe_open
import safetensors.torch as sft


def patch_one(path_in: str) -> int:
    if not os.path.isfile(path_in):
        print(f"  skip (not a file): {path_in}")
        return 1
    if not path_in.lower().endswith(".safetensors"):
        print(f"  skip (not .safetensors): {path_in}")
        return 1

    base, ext = os.path.splitext(path_in)
    path_out = base + "_fixed" + ext

    if os.path.exists(path_out):
        print(f"  output already exists, refusing to overwrite: {path_out}")
        return 1

    print(f"  reading: {path_in}")
    with safe_open(path_in, framework="pt") as f:
        src_meta = dict(f.metadata() or {})
        sd = {k: f.get_tensor(k) for k in f.keys()}

    # Drop any malformed global header from previous patch attempts; per-layer
    # comfy_quant byte tensors are the canonical source of truth.
    if "_quantization_metadata" in src_meta:
        print("  removing stale _quantization_metadata header")
        src_meta.pop("_quantization_metadata", None)

    patched = 0
    already_ok = 0
    for k in list(sd.keys()):
        if not k.endswith(".comfy_quant"):
            continue
        try:
            layer_conf = json.loads(sd[k].numpy().tobytes())
        except Exception as e:
            print(f"  WARNING: could not decode {k}: {e}")
            continue
        if "format" in layer_conf:
            already_ok += 1
            continue
        layer_conf["format"] = "int8_tensorwise"
        sd[k] = torch.tensor(
            list(json.dumps(layer_conf).encode("utf-8")),
            dtype=torch.uint8,
        )
        patched += 1

    if patched == 0 and already_ok == 0:
        print("  ERROR: no per-layer comfy_quant byte tensors found.")
        print("  This file doesn't look like it was saved by INT8-Fast's INT8ModelSave.")
        return 1

    if patched == 0:
        print(f"  nothing to do: all {already_ok} layers already have 'format'")
        return 0

    print(f"  writing: {path_out}")
    sft.save_file(sd, path_out, metadata=src_meta)
    print(f"  patched {patched} layer(s); {already_ok} already had 'format'")
    return 0


def main(argv):
    if len(argv) < 2:
        print("Usage: patch_int8_add_format.py <file.safetensors> [more files...]")
        return 1
    rc = 0
    for p in argv[1:]:
        print(f"\n--- {os.path.basename(p)} ---")
        rc |= patch_one(p)
    return rc


if __name__ == "__main__":
    sys.exit(main(sys.argv))