#!/usr/bin/env python3
"""Clone a production native-sage H3 workflow from ComfyUI history and optionally
inject MiniMaxH3BlockCacheT8. Usage:
  fetch_graph.py <prompt_id> <out.json> [--inject thr] [--length N] [--steps N] [--seed S] [--prefix P]
"""
import copy, json, sys, urllib.request

HIST = "http://127.0.0.1:8188/history/"

def main():
    pid = sys.argv[1]
    out = sys.argv[2]
    args = sys.argv[3:]
    def opt(flag, cast=str):
        return cast(args[args.index(flag) + 1]) if flag in args else None

    data = json.load(urllib.request.urlopen(HIST + pid, timeout=15))
    entry = data[pid]
    graph = copy.deepcopy(entry["prompt"][2])
    meta = {
        "source_prompt_id": pid,
        "status": entry.get("status", {}).get("status_str"),
        "started": entry.get("status", {}).get("messages", [[None]])[0][1] if entry.get("status", {}).get("messages") else None,
    }

    inject_thr = opt("--inject", float)
    length = opt("--length", int)
    steps = opt("--steps", int)
    seed = opt("--seed", int)
    prefix = opt("--prefix")

    if steps:
        for n in graph.values():
            if n["class_type"] in ("KSampler", "BasicScheduler"):
                if n["class_type"] == "BasicScheduler" or "steps" in n.get("inputs", {}):
                    n["inputs"]["steps"] = steps
    if seed is not None:
        for n in graph.values():
            if n["class_type"] in ("RandomNoise", "KSampler"):
                n["inputs"]["seed"] = seed
                if n["class_type"] == "RandomNoise":
                    n["inputs"]["noise_seed"] = seed
    if length:
        for n in graph.values():
            ct = n["class_type"]
            if ct in ("EmptyMiniMaxH3LatentAV", "MiniMaxH3ImageToVideo", "MiniMaxH3ReferenceToVideo") and "length" in n.get("inputs", {}):
                n["inputs"]["length"] = length
    if prefix:
        for n in graph.values():
            if n["class_type"] == "SaveVideo":
                n["inputs"]["filename_prefix"] = prefix

    if inject_thr is not None:
        # find UNETLoader id and the node consuming its MODEL (SigmaShift normally)
        unet_id = next(nid for nid, n in graph.items() if n["class_type"] == "UNETLoader")
        consumer_id = next(nid for nid, n in graph.items()
                           if any(isinstance(v, list) and v and v[0] == unet_id for v in n.get("inputs").values()))
        bc = {"class_type": "MiniMaxH3BlockCacheT8",
              "inputs": {"model": [unet_id, 0], "residual_diff_threshold": inject_thr,
                         "start_percent": 0.08, "end_percent": 0.95, "max_consecutive_hits": 2,
                         "cache_device": "cpu", "metric_stride": 8, "verbose": True}}
        graph["BC1"] = bc
        # rewire consumer's model input to BC output (keep other inputs)
        for v in graph[consumer_id]["inputs"].values():
            if isinstance(v, list) and v and v[0] == unet_id:
                v[0] = "BC1"
        meta["injected"] = {"thr": inject_thr, "unet": unet_id, "consumer": consumer_id}

    json.dump({"graph": graph, "meta": meta}, open(out, "w"), ensure_ascii=False, indent=1)
    print(json.dumps(meta, ensure_ascii=False))
    # summary of key nodes
    for nid, n in sorted(graph.items()):
        ins = {k: (v if not isinstance(v, list) else f"[{v[0]},{v[1]}]") for k, v in n.get("inputs", {}).items() if k in ("class_type",) or True}
        keys = {k: ins[k] for k in ins if k in ("steps", "seed", "noise_seed", "length", "width", "height", "residual_diff_threshold", "filename_prefix")}
        print(f"{nid:>4} {n['class_type']:<32} {keys}")

if __name__ == "__main__":
    main()