import argparse
import json
import subprocess
import sys
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 a.input:
    print("ERROR: no input image")
    sys.exit(1)
cli = Path(sys.executable).parent / "mflux-upscale-seedvr2"
if not cli.exists():
    print(f"ERROR: {cli} missing. Run scripts/install_mflux.sh")
    sys.exit(1)

src = Path(a.input[0])
outdir = Path(a.outdir)
out = outdir / f"{src.stem}_upscaled.png"
cmd = [str(cli), "-i", str(src.resolve()),
       "--resolution", str(p.get("resolution", "2x")),
       "--softness", str(p.get("softness", 0.0)),
       "--seed", str(p.get("seed", 42)),
       "--output", str(out.resolve())]
quant = str(p.get("quantize", "none"))
if quant in ("4", "8"):
    cmd += ["-q", quant]

print("+", " ".join(cmd), flush=True)
print("(first run downloads SeedVR2 weights)", flush=True)
res = subprocess.run(cmd, cwd=str(outdir), stdout=sys.stdout, stderr=subprocess.STDOUT)
if res.returncode != 0:
    sys.exit(res.returncode)

pngs = sorted(outdir.rglob("*.png"))
if not pngs:
    print("ERROR: no upscaled image produced")
    sys.exit(1)
img = out if out.exists() else pngs[-1]
(outdir / "result.json").write_text(json.dumps(
    {"outputs": [{"path": img.name, "meta": {"tool": "seedvr2", "resolution": p.get("resolution", "2x")}}]}))
print(f"done: {img.name}", flush=True)