#!/usr/bin/env python3
"""Copy finished FLUX assets out of MODELBEAST into assets/generated/<name>.png.

Reads queue.lock.json (job_id/name/seed written by enqueue.py), finds each
generated image in the MB asset store by its seed-based filename, and copies it
in under the friendly asset name. Run after the MB jobs finish.

    /opt/homebrew/bin/uv run python assets/collect.py
"""
import json, sys, shutil
from pathlib import Path

MB = Path.home() / "MODELBEAST"
sys.path.insert(0, str(MB))
import server.db as db  # noqa: E402

HERE = Path(__file__).resolve().parent
lock = json.loads((HERE / "queue.lock.json").read_text())
outdir = HERE / "generated"
outdir.mkdir(exist_ok=True)

con = db.connect()
done = missing = 0
for it in lock:
    fname = f"flux_flux2-klein-4b_s{it['seed']}.png"
    row = con.execute("SELECT path FROM assets WHERE name=? ORDER BY created_at DESC LIMIT 1", (fname,)).fetchone()
    if row and Path(row["path"]).exists():
        shutil.copy(row["path"], outdir / f"{it['name']}.png")
        done += 1
    else:
        print(f"  not ready: {it['name']}")
        missing += 1
print(f"collected {done}/{len(lock)} into {outdir}" + (f" ({missing} not ready)" if missing else ""))