#!/usr/bin/env python3
"""Shoot a face's whole store gallery from the simulator, one shot per build.

The shot list is the FACE'S, at `<Face>/store-assets/shots.json`; the rest is
shared. A face joins by carrying the DEMO-CLOCK and DEMO-GAUGE anchors that
`tools/sim.py` matches. Raw 454 frames land in `store-assets/`, a copy fitted
under the store's 150 kB limit in `store-assets/store/`.

  python3 <this> --face Truss [--list] [--only 2 truss-5-violet-rhodium]
"""
import argparse
import json
import os
import sys
import time

# scripts -> store-release -> skills -> .claude -> the repo root
ROOT = os.path.abspath(os.path.join(os.path.dirname(os.path.abspath(__file__)),
                                    *([os.pardir] * 4)))
sys.path.insert(0, os.path.join(ROOT, "tools"))
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import sim                                                    # noqa: E402
from crop_dials import save_under_limit                       # noqa: E402
from PIL import Image                                         # noqa: E402

# Measured on venu3: kill, rebuild, boot, settle, save. A gallery is priced
# from this before it is launched, never guessed at.
SECONDS_PER_SHOT = 30


def load(face):
    p = os.path.join(ROOT, face, "store-assets", "shots.json")
    if not os.path.exists(p):
        raise SystemExit("no shot list at %s" % os.path.relpath(p, ROOT))
    with open(p, encoding="utf-8") as f:
        spec = json.load(f)
    for i, s in enumerate(spec["shots"]):
        for k in ("name", "time", "gauge"):
            if k not in s:
                raise SystemExit("shot %d has no %r" % (i + 1, k))
        s.setdefault("mode", "active")
        s.setdefault("set", {})
        if "battery" in s and not 0 <= float(s["battery"]) <= 100:
            raise SystemExit("shot %d: battery out of range" % (i + 1))
    names = [s["name"] for s in spec["shots"]]
    if len(set(names)) != len(names):
        raise SystemExit("two shots share a name")
    return spec


def pick(shots, only):
    if not only:
        return list(range(len(shots)))
    out = []
    for token in only:
        if token.isdigit() and 1 <= int(token) <= len(shots):
            out.append(int(token) - 1)
            continue
        hits = [i for i, s in enumerate(shots) if s["name"] == token]
        if not hits:
            raise SystemExit("no shot called %r" % token)
        out += hits
    return sorted(set(out))


def fit(raw, dst_dir, name):
    """The store copy: 454 square, under 150 kB, PNG if it will fit."""
    os.makedirs(dst_dir, exist_ok=True)
    im = Image.open(raw).convert("RGB")
    if im.size != (454, 454):
        im = im.resize((454, 454), Image.LANCZOS)
    return save_under_limit(im, os.path.join(dst_dir, name))


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--face", default=None)
    ap.add_argument("--only", nargs="*", default=None,
                    help="1-based indices or shot names")
    ap.add_argument("--list", action="store_true")
    a = ap.parse_args()

    f = sim.resolve(a.face)
    spec = load(f.name)
    shots = spec["shots"]
    device = spec.get("device", "venu3")

    if a.list:
        for i, s in enumerate(shots):
            print("%2d  %-30s %-6s %-9s %s" % (i + 1, s["name"], s["mode"],
                                               s["time"], s.get("why", "")))
        return

    todo = pick(shots, a.only)
    mins = len(todo) * SECONDS_PER_SHOT / 60.0
    print("%s: %d shot(s) on %s, one build each, about %.0f min"
          % (f.name, len(todo), device, mins))
    if os.path.isdir(f.saved) and os.listdir(f.saved):
        raise SystemExit("stale stash in %s - inspect it, put back what is "
                         "yours, then delete it" % f.saved)

    out = os.path.join(f.dir, "store-assets")
    os.makedirs(out, exist_ok=True)
    started, failed = time.time(), []
    try:
        for n, i in enumerate(todo, 1):
            s = shots[i]
            print("[%d/%d] %s  %s  %s" % (n, len(todo), s["name"], s["mode"],
                                          s["time"]), flush=True)
            sim.set_clock(f, tuple(int(x) for x in s["time"].split(":")))
            sim.set_gauges(f, *(int(x) for x in s["gauge"].split(",")))
            if s["set"]:
                sim.set_props(f, {k: int(v) for k, v in s["set"].items()})
            dest = os.path.join(out, s["name"] + ".png")
            if not sim.deploy(f, device):
                failed.append(s["name"] + " (build)")
                continue
            if "battery" in s:
                sim.set_battery(s["battery"])
            if not sim.capture(f, dest, s["mode"]):
                failed.append(s["name"] + " (capture)")
                continue
            path, size = fit(dest, os.path.join(out, "store"), s["name"])
            left = (len(todo) - n) * (time.time() - started) / n / 60.0
            print("      %s  %d kB   %.0f min left"
                  % (os.path.relpath(path, ROOT), size / 1024, left),
                  flush=True)
    finally:
        # The source goes back and the binary goes away whatever happened: a
        # .prg left here carries a frozen clock, and one reached a wrist.
        sim.revert(f)
        try:
            os.remove(os.path.join(f.dir, "bin", f.app + ".prg"))
        except OSError:
            pass

    print("done in %.0f min" % ((time.time() - started) / 60.0))
    if failed:
        raise SystemExit("FAILED: " + ", ".join(failed))


if __name__ == "__main__":
    main()