import argparse
import json
import subprocess
import sys
import time
from pathlib import Path

ap = argparse.ArgumentParser()
ap.add_argument("--input", action="append", default=[])
ap.add_argument("--outdir", required=True)
ap.add_argument("--params", default="{}")
a = ap.parse_args()
p = json.loads(a.params)

if not p.get("prompt"):
    print("ERROR: prompt is required")
    sys.exit(1)

model = p.get("model", "schnell")
steps = int(p.get("steps", 4))
outdir = Path(a.outdir)
out_png = outdir / f"flux_{model}_s{p.get('seed', 42)}.png"

# mflux-generate lives next to this venv's python
cli = Path(sys.executable).parent / "mflux-generate"
if not cli.exists():
    print(f"ERROR: mflux not installed ({cli} missing). Run scripts/install_mflux.sh")
    sys.exit(1)

cmd = [
    str(cli),
    "--model", model,
    "--prompt", p["prompt"],
    "--steps", str(steps),
    "--width", str(p.get("width", 1024)),
    "--height", str(p.get("height", 1024)),
    "--seed", str(p.get("seed", 42)),
    "--output", str(out_png),
]
quant = str(p.get("quantize", "8"))
if quant in ("4", "8"):
    cmd += ["--quantize", quant]
if model == "dev":
    cmd += ["--guidance", str(p.get("guidance", 3.5))]

print("+", " ".join(cmd), flush=True)
print("(first run downloads weights from HuggingFace — can take a while)", flush=True)
t0 = time.time()
res = subprocess.run(cmd, stdout=sys.stdout, stderr=subprocess.STDOUT)
if res.returncode != 0:
    print("HINT: if this failed with a gated-repo/auth error, the model needs an "
          "HF license accept + token (Settings → HuggingFace token). schnell is ungated.")
    sys.exit(res.returncode)
elapsed = round(time.time() - t0, 1)

if not out_png.exists():
    # mflux may append suffixes; grab any png it wrote
    pngs = sorted(outdir.glob("*.png"))
    if not pngs:
        print("ERROR: no image produced")
        sys.exit(1)
    out_png = pngs[-1]

(outdir / "result.json").write_text(json.dumps({"outputs": [
    {"path": out_png.name,
     "meta": {"tool": "mflux", "model": model, "steps": steps,
              "seed": p.get("seed", 42), "quantize": quant, "seconds": elapsed}}]}))
print(f"done in {elapsed}s: {out_png.name}", flush=True)