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

ROOT = Path(__file__).resolve().parents[3]
VENDOR = ROOT / "vendor" / "mflux-qwen-layered"
CLI = VENDOR / ".venv" / "bin" / "mflux-generate-qwen-layered"

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 a.input:
    print("ERROR: no input image")
    sys.exit(1)
if not CLI.exists():
    print(f"ERROR: qwen-layered not installed at {VENDOR}. "
          "Run scripts/install_qwen_layered.sh")
    sys.exit(1)

# Baked-model resolution: env override, then per-box conventional paths,
# else fall back to the HF repo with on-the-fly q8 (slow first run: 54GB).
candidates = [os.environ.get("QWEN_LAYERED_MODEL", "")]
candidates += [str(Path.home() / "qwen-layered" / "qwen-layered-q8"),
               str(Path.home() / "qwen-layered-staging" / "qwen-layered-q8")]
model_path = next((c for c in candidates if c and Path(c).is_dir()), None)

outdir = Path(a.outdir)
cmd = [
    str(CLI), "--image", str(Path(a.input[0]).resolve()),
    "--layers", str(p.get("layers", 4)),
    "--steps", str(p.get("steps", 20)),
    "--resolution", str(p.get("resolution", 640)),
    "--seed", str(p.get("seed", 0)),
    "--output-dir", str(outdir.resolve()),
]
if model_path:
    cmd += ["--model-path", model_path]
    print(f"using baked model: {model_path}", flush=True)
else:
    cmd += ["-q", "8"]
    print("no baked model found — HF fallback with on-the-fly q8 "
          "(first run downloads 54GB)", flush=True)
if p.get("prompt"):
    cmd += ["--prompt", str(p["prompt"])]

env = os.environ.copy()
env["HF_HUB_DISABLE_XET"] = "1"

print("+", " ".join(cmd), flush=True)
res = subprocess.run(cmd, cwd=str(VENDOR), stdout=sys.stdout,
                     stderr=subprocess.STDOUT, env=env)
if res.returncode != 0:
    sys.exit(res.returncode)

layers = sorted(outdir.glob("layer_*.png"))
if not layers:
    print("ERROR: qwen-layered produced no layer PNGs")
    sys.exit(1)
stem = Path(a.input[0]).stem
outputs = [{"path": str(f), "name": f"{stem}_{f.stem}.png",
            "meta": {"tool": "qwen_layered_local", "layer": i}}
           for i, f in enumerate(layers)]
(outdir / "result.json").write_text(json.dumps({"outputs": outputs}))
print(f"done: {len(layers)} layers", flush=True)