#!/usr/bin/env python3
"""Compile production/flow/*_SHOTS.md into Flow Auto-Pilot prompt banks.

The flowext extension types prompt TEXT only — attached characters and the
model/output-count are composer state set by hand once per batch. So shots are
grouped into one bank per (tier, attach-set); John opens the RSG Flow project,
attaches that bank's characters, sets the model (Fast/Quality) and 1x, then the
bank is enqueued and drains unattended. 2 takes per shot per the spend doctrine.

Usage:
  make_flow_bank.py               # write banks to ~/Documents/flowext/banks/rsg/
  make_flow_bank.py --enqueue N   # POST bank number N to localhost:8017/enqueue
"""
import json
import re
import sys
import urllib.request
from pathlib import Path

ROOT = Path(__file__).resolve().parent.parent
SHEETS = ["TRACK1", "TRACK2", "TRACK3", "TRACK4", "TRACK5", "TRACK6", "TRACK7",
          "SHIFTS", "EXTRAS"]
BANKS = Path.home() / "Documents/flowext/banks/rsg"
SKIP = {"Logging", "t1_whitelabel", "attract_loop", "intro_qual"}  # done or cut
TAKES = 2
SERVER = "http://localhost:8017"


def parse():
    shots = []  # (id, chars tuple, quality bool, prompt)
    for name in SHEETS:
        text = (ROOT / f"production/flow/{name}_SHOTS.md").read_text()
        for m in re.finditer(r"^## (.+?)\n(.*?)(?=^## |\Z)", text, re.M | re.S):
            header, body = m.group(1), m.group(2)
            sid = header.split()[0]
            if sid in SKIP:
                continue
            quality = "QUALITY" in header
            am = re.search(r"\[attach: ([^\]]+)\]", header) or \
                 re.search(r"\[([^\]]+)\]", header)
            chars = () if not am or am.group(1) == "no characters" else \
                tuple(c.strip() for c in am.group(1).split(","))
            lines = [l.strip() for l in body.strip().splitlines() if l.strip()]
            if lines and lines[0].startswith("(") and lines[0].endswith(")"):
                lines = lines[1:]  # production note, not paste text
            prompt = re.sub(r"\s+", " ", " ".join(lines)).strip()
            if len(prompt) < 80:  # header wrap fragments / non-prompt sections
                print(f"  !! skipping {sid}: body too short ({len(prompt)} ch)")
                continue
            shots.append((sid, chars, quality, prompt))
    return shots


def slug(chars):
    return "+".join(re.sub(r"^the ", "", c.lower()).replace(" ", "") for c in chars) or "nochars"


def main():
    shots = parse()
    groups = {}  # (quality, chars) -> [shots] — keep sheet order within a group
    for s in shots:
        groups.setdefault((s[2], s[1]), []).append(s)
    # fast banks first (bulk of the game), quality after; big groups first
    order = sorted(groups, key=lambda k: (k[0], -len(groups[k])))
    BANKS.mkdir(parents=True, exist_ok=True)
    for old in BANKS.glob("*.jsonl"):
        old.unlink()
    for n, key in enumerate(order, 1):
        quality, chars = key
        model = "Veo 3.1 - Quality" if quality else "Veo 3.1 - Fast"
        path = BANKS / f"{n:02d}__{'quality' if quality else 'fast'}__{slug(chars)}.jsonl"
        with open(path, "w") as f:
            for sid, _, _, prompt in groups[key]:
                for take in range(1, TAKES + 1):
                    f.write(json.dumps({
                        "id": f"rsg_{sid}__t{take}", "kind": "video", "model": model,
                        "title": f"{sid} take {take}", "prompt": prompt,
                        "category": "rsg"}) + "\n")
        gens = len(groups[key]) * TAKES
        cr = gens * (100 if quality else 10)
        print(f"{path.name}: {len(groups[key])} shots x{TAKES} = {gens} gens ~{cr} cr"
              f"  [attach: {', '.join(chars) or 'none'}]")
    if len(sys.argv) > 2 and sys.argv[1] == "--enqueue":
        bank = sorted(BANKS.glob(f"{int(sys.argv[2]):02d}__*.jsonl"))[0]
        for line in bank.read_text().splitlines():
            req = urllib.request.Request(SERVER + "/enqueue", line.encode(),
                                         {"Content-Type": "application/json"})
            urllib.request.urlopen(req, timeout=5).read()
        print(f"enqueued {bank.name}")


if __name__ == "__main__":
    main()