"""
Generate per-step latent snapshots using StepwiseKSampler.
One prompt × 8 quant levels × 1 seed = 8 runs (for step-by-step visualization).
"""

import argparse
import os
import time
from config import GGUF_MODELS, MODEL_ORDER, PROMPTS
from utils import check_comfyui, queue_prompt, wait_for_completion, \
                  build_stepwise_workflow


def main():
    parser = argparse.ArgumentParser(description="Generate stepwise latents")
    parser.add_argument("--output-dir", required=True)
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--prompt-id", default="A01",
                        help="Single prompt ID for stepwise analysis")
    args = parser.parse_args()

    os.makedirs(args.output_dir, exist_ok=True)
    check_comfyui()

    prompt_text = PROMPTS[args.prompt_id]
    total = len(MODEL_ORDER)

    print(f"Stepwise generation: {args.prompt_id}, seed={args.seed}")
    print(f"Models: {total}\n")

    for i, tag in enumerate(MODEL_ORDER):
        model_file = GGUF_MODELS[tag]
        prefix = f"stepwise_{tag.lower()}_s{args.seed}"
        print(f"[{i+1}/{total}] {tag} -> {prefix} ...", end=" ", flush=True)

        wf = build_stepwise_workflow(model_file, prompt_text, args.seed, prefix)
        t0 = time.time()
        try:
            prompt_id = queue_prompt(wf)
            ok = wait_for_completion(prompt_id)
            print(f"{'ok' if ok else 'FAIL'} ({time.time()-t0:.1f}s)")
        except Exception as e:
            print(f"ERROR: {e}")


if __name__ == "__main__":
    main()