#!/usr/bin/env python3
"""Queue MOSS-Image prompts against the local ComfyUI (http://127.0.0.1:8188).

  python3 moss.py "a gorilla in a nightclub..." --seeds 1 2 3
  python3 moss.py -f prompts.txt --enhance --lora r256
"""
import argparse, copy, json, os, sys, time, urllib.request

HERE = os.path.dirname(os.path.abspath(__file__))
URL = "http://127.0.0.1:8188"
OUT = os.path.join(HERE, "..", "..", "output")
MODELS = {
    "base": "moss_image_base_bf16.safetensors",
    "sft": "moss_image_sft_bf16.safetensors",
}
LORAS = {
    "b21": "moss_image_sft_lora_blocks_21_to_49_rank_512_bf16.safetensors",
    "r256": "moss_image_sft_lora_rank_256_bf16.safetensors",
    "r512": "moss_image_sft_lora_rank_512_bf16.safetensors",
}


def call(path, data=None):
    req = urllib.request.Request(URL + path, data=json.dumps(data).encode() if data is not None else None,
                                 headers={"Content-Type": "application/json"})
    with urllib.request.urlopen(req) as r:
        return json.load(r)


def build(base, prompt, seed, a):
    g = copy.deepcopy(base)
    g["1"]["inputs"]["unet_name"] = MODELS.get(a.model, a.model)
    if a.swap_layers is not None or a.swap_other:
        # base model + exact donor weights for chosen layers/parts (LoRA chain is skipped)
        g["40"] = {"class_type": "UNETLoader", "inputs": {"unet_name": MODELS.get(a.donor, a.donor), "weight_dtype": "default"}}
        g["41"] = {"class_type": "MOSSBlockSwap", "inputs": {"model": ["1", 0], "donor": ["40", 0], "layers": a.swap_layers or "",
                   "parts": a.swap_parts, "include_other": a.swap_other, "ratio": a.swap_ratio}}
        g["16"]["inputs"]["model"] = ["41", 0]
        del g["30"]
    elif a.lora == "none":
        g["16"]["inputs"]["model"] = ["1", 0]
        del g["30"]
    else:
        g["30"]["inputs"]["lora_name"] = LORAS.get(a.lora, a.lora)
        g["30"]["inputs"]["strength_model"] = a.strength
    if a.lora2:
        g["31"] = {"class_type": "LoraLoaderModelOnly", "inputs": {"model": g["16"]["inputs"]["model"],
                   "lora_name": LORAS.get(a.lora2, a.lora2), "strength_model": a.strength2}}
        g["16"]["inputs"]["model"] = ["31", 0]
    if a.enhance:
        g["17"]["inputs"]["prompt"] = prompt
        g["17"]["inputs"]["sampling_mode.seed"] = seed
    else:
        g["5"]["inputs"]["text"] = prompt
        for k in ("17", "18", "19"):
            del g[k]
    if a.probe:
        g["42"] = {"class_type": "MOSSStripeProbe", "inputs": {"model": g["16"]["inputs"]["model"], "head_from": MODELS.get(a.model, a.model),
                   "tag": f"{a.probe}_s{seed}", "width": a.width, "height": a.height}}
        g["16"]["inputs"]["model"] = ["42", 0]
    if a.no_tsr:
        g["2"]["inputs"]["model"] = g["16"]["inputs"]["model"]
        del g["16"]
    if a.no_slg:
        for k in ("20", "22"):
            g[k]["inputs"]["model"] = g["2"]["inputs"]["model"]
        del g["2"]
    if a.sampler != "er_sde":
        g["21"] = {"class_type": "KSamplerSelect", "inputs": {"sampler_name": a.sampler}}
    if a.negative is not None:
        g["6"]["inputs"]["text"] = a.negative
    g["20"]["inputs"].update(noise_seed=seed, cfg=a.cfg)
    g["22"]["inputs"]["steps"] = a.steps
    g["7"]["inputs"].update(width=a.width, height=a.height)
    g["10"]["inputs"]["filename_prefix"] = "moss/" + a.prefix
    if a.window_steps:
        window(g, a)
    return g


def window(g, a):
    """Apply --window-lora only for sampling steps a..b (1-indexed) via ComfyUI hook keyframes on both conds.
    --window-steps "1-N" applies it for the whole run; strength 0 everywhere is a no-op control."""
    lo, hi = (int(v) for v in a.window_steps.split("-"))
    n = a.steps
    g["70"] = {"class_type": "CreateHookLoraModelOnly", "inputs": {"lora_name": a.window_lora, "strength_model": a.window_strength}}
    kf = ([(0.0, 0.0)] if lo > 1 else []) + [((lo - 1) / n, 1.0)] + ([(hi / n, 0.0)] if hi < n else [])
    prev = None
    for i, (pct, mult) in enumerate(kf):
        node = {"class_type": "CreateHookKeyframe", "inputs": {"strength_mult": mult, "start_percent": pct}}
        if prev:
            node["inputs"]["prev_hook_kf"] = [prev, 0]
        prev = str(71 + i)
        g[prev] = node
    g["75"] = {"class_type": "SetHookKeyframes", "inputs": {"hooks": ["70", 0], "hook_kf": [prev, 0]}}
    for cid, src in (("76", "5"), ("77", "6")):
        g[cid] = {"class_type": "ConditioningSetProperties", "inputs": {"cond_NEW": [src, 0], "strength": 1.0, "set_cond_area": "default", "hooks": ["75", 0]}}
    g["20"]["inputs"]["positive"], g["20"]["inputs"]["negative"] = ["76", 0], ["77", 0]


def main():
    p = argparse.ArgumentParser()
    p.add_argument("prompt", nargs="*")
    p.add_argument("-f", "--file", help="one prompt per paragraph (blank-line separated)")
    p.add_argument("--seeds", type=int, nargs="+", default=[16])
    p.add_argument("--enhance", action="store_true", help="rewrite the prompt with Qwen3.5 first")
    p.add_argument("--model", default="sft", help="base | sft | <filename>")
    p.add_argument("--lora", default="b21", help="b21 | r256 | none | <filename>")
    p.add_argument("--strength", type=float, default=-0.5)
    p.add_argument("--donor", default="sft", help="model whose exact weights are swapped in (with --swap-layers)")
    p.add_argument("--swap-layers", help='layers to take from --donor, e.g. "21-41" or "3,4"; enables block swap')
    p.add_argument("--swap-parts", default="all", help="all | comma list of qkv,o,ffn_up_gate,ffn_down,norm")
    p.add_argument("--swap-other", action="store_true", help="also swap non-layer weights (embedders, modulation, head)")
    p.add_argument("--swap-ratio", type=float, default=1.0)
    p.add_argument("--probe", help="record per-layer stripe spectra to output/moss/probe/<probe>_s<seed>.npz")
    p.add_argument("--window-steps", help='e.g. "5-12": apply --window-lora only for these sampling steps (1-indexed)')
    p.add_argument("--window-lora", default="moss_image_sft_banding_fix_L37-41_bf16.safetensors")
    p.add_argument("--window-strength", type=float, default=1.0)
    p.add_argument("--lora2", help="second LoRA chained after the first, e.g. --lora r512 --strength 1 --lora2 b21 --strength2 -0.5")
    p.add_argument("--strength2", type=float, default=1.0)
    p.add_argument("--no-tsr", action="store_true")
    p.add_argument("--no-slg", action="store_true")
    p.add_argument("--sampler", default="er_sde", help="er_sde (workflow's SamplerER_SDE) or any KSampler sampler name")
    p.add_argument("--steps", type=int, default=25)
    p.add_argument("--cfg", type=float, default=3.0)
    p.add_argument("--width", type=int, default=1024)
    p.add_argument("--height", type=int, default=1024)
    p.add_argument("--negative")
    p.add_argument("--prefix", default="moss")
    a = p.parse_args()

    prompts = [" ".join(a.prompt)] if a.prompt else []
    if a.file:
        prompts += [x.strip() for x in open(a.file).read().split("\n\n") if x.strip()]
    if not prompts:
        p.error("give a prompt or -f file")

    base = json.load(open(os.path.join(HERE, "moss_image_t2i_api.json")))
    jobs = []
    for prompt in prompts:
        for seed in a.seeds:
            r = call("/prompt", {"prompt": build(base, prompt, seed, a)})
            if r.get("node_errors"):
                sys.exit(json.dumps(r["node_errors"], indent=1))
            jobs.append((r["prompt_id"], prompt, seed))
    print(f"queued {len(jobs)}", flush=True)

    for pid, prompt, seed in jobs:
        while True:
            h = call(f"/history/{pid}").get(pid)
            if h and h.get("status", {}).get("completed") is not None:
                break
            time.sleep(1)
        if h["status"].get("status_str") != "success":
            print(f"FAILED seed {seed}: {json.dumps(h['status'].get('messages', [])[-1:])[:800]}")
            continue
        outs = h["outputs"]
        for img in outs.get("10", {}).get("images", []):
            print(os.path.normpath(os.path.join(OUT, img["subfolder"], img["filename"])), f"seed={seed}")
        if "19" in outs:
            print("  enhanced:", outs["19"]["text"][0][:1500])


if __name__ == "__main__":
    main()