#!/usr/bin/env python3
"""Generate still images as reel b-roll, coherent across cuts.

Adapted from jbaruch/speaker-toolkit's illustrations skill. The hard-won part
is not the API call — it is knowing when to edit and when to regenerate, and
what makes a SET of images look like one piece rather than a stock-photo
collage. See references/generated-stills.md.

Two things differ from the deck case it came from:

  ASPECT. A deck wants 16:9; a reel wants 9:16. Images are generated at the
  vendor's nearest PORTRAIT size (2:3 for OpenAI) rather than exactly 9:16, on
  purpose: normalize_clips.py crops to the target aspect, and a frame that is
  already exactly 9:16 leaves render_reel.py no spare pixels to pan with.

  DESTINATION. Output is a still in raw/, not a slide. normalize_clips.py turns
  it into motion via its `kb` (Ken Burns) field — a static still held for three
  seconds reads as dead air, a drifting one does not.

Coherence across a SET is the whole game. One style anchor renders on every
image; per-image prompts carry only the scene. Mixing scene content into the
anchor is what makes a set drift.

Keys: OPENAI_API_KEY for gpt-image-*, GEMINI_API_KEY for gemini-*/nano-banana-*
and imagen-*. Dispatch is by model-name prefix, so a newer model from either
vendor works without a code change.

Usage:
  gen_images.py spec.json --outdir raw/
  gen_images.py spec.json --outdir raw/ --only 2,5      # regenerate two
  gen_images.py spec.json --outdir raw/ --edit 3 "Erase the caption" \\
      --keep "Keep the rider. Keep the horizon."
"""
import argparse, base64, json, os, sys, urllib.error, urllib.request
from pathlib import Path

DEFAULT_MODEL = "gpt-image-2"
OPENAI_BASE = "https://api.openai.com/v1"
GEMINI_BASE = "https://generativelanguage.googleapis.com/v1beta/models"

# Nearest PORTRAIT size per vendor. Deliberately not exactly 9:16 — the
# pipeline crops to aspect, and spare pixels are what --pan reframes with.
PORTRAIT = {
    "openai": "1024x1536",      # 2:3
    "gemini": "2:3",            # must be sent — Gemini defaults to LANDSCAPE
    "imagen": "3:4",            # Imagen's nearest; it has no native 2:3
}

# Appended to every fresh generation. A style-only anchor is the real fix for
# content leaking between images; this is a backstop, not a substitute.
SCENE_GUARD = ("COMPOSE ONLY THE SCENE DESCRIBED ABOVE. Do not add captions, "
               "labels, logos, borders, or framing devices unless named.")

# Every edit prompt carries these two. Gemini in particular decorates
# aggressively, and flat-fills the hole when asked to erase something.
EDIT_GUARDS = ("DO NOT add any new elements.",
               "Let the background continue naturally — no flat patch.")


def model_family(model):
    """Vendor family from the model id prefix.

    Prefix dispatch rather than an allowlist: a model released after this file
    was written still renders, which is the point.
    """
    if model.startswith("gpt-image") or model.startswith("chatgpt-image"):
        return "openai"
    if model.startswith("imagen"):
        return "imagen"
    return "gemini"


def compose_prompt(anchor, scene):
    """One style anchor + one scene, plus the scene guard.

    The anchor is STYLE ONLY — medium, palette, rendering technique, light.
    Scene content in the anchor bleeds into every image in the set and is the
    usual reason a set stops looking like one piece.
    """
    parts = [p.strip() for p in (anchor, scene) if p and p.strip()]
    parts.append(SCENE_GUARD)
    return "\n\n".join(parts)


def edit_prompt(instruction, keep=None):
    """An edit instruction with the guards that make edits survivable.

    Only ever phrase an edit as a REMOVAL. Additions and modifications strip
    the style — regenerate those from the full prompt instead. The `keep` list
    is not optional in practice: without it the model removes neighbours of the
    thing you asked it to remove.
    """
    parts = [instruction.strip(), *EDIT_GUARDS]
    if keep:
        parts.append(keep.strip())
    return " ".join(parts)


def next_version(path):
    """A free filename beside `path`, never overwriting a previous attempt.

    Stills take many attempts to converge, and stomping a near-good one to try
    a variation loses something that may have been worth keeping. Disk is
    cheaper than a render you cannot get back.
    """
    path = Path(path)
    if not path.exists():
        return path
    n = 2
    while True:
        cand = path.with_name(f"{path.stem}_v{n}{path.suffix}")
        if not cand.exists():
            return cand
        n += 1


# More than one name is in use for the same Google key, so read them all
# rather than making the caller rename an environment variable.
KEY_ENV = {
    "openai": ("OPENAI_API_KEY",),
    "gemini": ("GEMINI_API_KEY", "GOOGLE_AI_API_KEY", "GOOGLE_AI_API",
               "GOOGLE_API_KEY"),
}
KEY_ENV["imagen"] = KEY_ENV["gemini"]


def api_key(family, env=None):
    """First key set for this vendor, or exit naming every name checked."""
    env = os.environ if env is None else env
    names = KEY_ENV[family]
    for name in names:
        if env.get(name):
            return env[name]
    sys.exit(f"No key for {family}. Set one of: {', '.join(names)}. Generated "
             f"stills are optional; the rest of the pipeline is unaffected.")


def post(url, payload, headers, timeout=300):
    req = urllib.request.Request(
        url, data=json.dumps(payload).encode(),
        headers={"Content-Type": "application/json", **headers})
    try:
        with urllib.request.urlopen(req, timeout=timeout) as r:
            return json.loads(r.read())
    except urllib.error.HTTPError as e:
        body = e.read().decode()[:400]
        sys.exit(f"{url.split('/')[-1]} failed ({e.code}): {body}")
    except urllib.error.URLError as e:
        sys.exit(f"Could not reach the image API: {e.reason}")


def generate(model, prompt, out_path):
    """Render one image and write it. Returns the path written."""
    family = model_family(model)
    key = api_key(family)
    if family == "openai":
        data = post(f"{OPENAI_BASE}/images/generations",
                    {"model": model, "prompt": prompt,
                     "size": PORTRAIT["openai"], "n": 1},
                    {"Authorization": f"Bearer {key}"})
        b64 = data["data"][0]["b64_json"]
    elif family == "imagen":
        data = post(f"{GEMINI_BASE}/{model}:predict?key={key}",
                    {"instances": [{"prompt": prompt}],
                     "parameters": {"sampleCount": 1,
                                    "aspectRatio": PORTRAIT["imagen"]}}, {})
        b64 = data["predictions"][0]["bytesBase64Encoded"]
    else:
        # aspectRatio is not optional here. Without it Gemini returns
        # LANDSCAPE (measured 1408x768), which a 9:16 crop would gut.
        data = post(f"{GEMINI_BASE}/{model}:generateContent?key={key}",
                    {"contents": [{"parts": [{"text": prompt}]}],
                     "generationConfig": {
                         "imageConfig": {"aspectRatio": PORTRAIT["gemini"]}}},
                    {})
        b64 = None
        for part in data["candidates"][0]["content"]["parts"]:
            if "inlineData" in part:
                b64 = part["inlineData"]["data"]
                break
        if not b64:
            sys.exit(f"{model} returned no image — it may not be an image model.")
    out_path = Path(out_path)
    out_path.parent.mkdir(parents=True, exist_ok=True)
    out_path.write_bytes(base64.b64decode(b64))
    return out_path


def edit_image(model, src, prompt, out_path):
    """Edit an existing image. OpenAI only — Imagen has no edit endpoint.

    Multipart rather than JSON: the images/edits endpoint takes the source
    image as a file part, so this is the one call that cannot go through post().
    """
    family = model_family(model)
    if family == "imagen":
        sys.exit(f"{model} has no edit endpoint. Regenerate with --only "
                 f"instead, or switch to an editable model.")
    if family != "openai":
        sys.exit(f"Editing is wired for gpt-image-* only; {model} is not. "
                 f"Regenerate with --only instead.")
    import mimetypes, uuid
    key = api_key(family)
    boundary = uuid.uuid4().hex
    src = Path(src)
    ctype = mimetypes.guess_type(str(src))[0] or "image/png"
    body = b""
    for name, value in (("model", model), ("prompt", prompt),
                        ("size", PORTRAIT["openai"]), ("n", "1")):
        body += (f"--{boundary}\r\nContent-Disposition: form-data; "
                 f'name="{name}"\r\n\r\n{value}\r\n').encode()
    body += (f"--{boundary}\r\nContent-Disposition: form-data; "
             f'name="image"; filename="{src.name}"\r\n'
             f"Content-Type: {ctype}\r\n\r\n").encode()
    body += src.read_bytes() + b"\r\n"
    body += f"--{boundary}--\r\n".encode()

    req = urllib.request.Request(
        f"{OPENAI_BASE}/images/edits", data=body,
        headers={"Authorization": f"Bearer {key}",
                 "Content-Type": f"multipart/form-data; boundary={boundary}"})
    try:
        with urllib.request.urlopen(req, timeout=300) as r:
            data = json.loads(r.read())
    except urllib.error.HTTPError as e:
        sys.exit(f"images/edits failed ({e.code}): {e.read().decode()[:400]}")
    out_path = Path(out_path)
    out_path.write_bytes(base64.b64decode(data["data"][0]["b64_json"]))
    return out_path


def main():
    ap = argparse.ArgumentParser(
        description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("spec", help="JSON: {style_anchor, model?, images:[{name,prompt}]}")
    ap.add_argument("--outdir", default="raw",
                    help="Where stills land (default raw/, like any footage)")
    ap.add_argument("--only", help="Comma list of 1-based indices to render")
    ap.add_argument("--edit", nargs=2, metavar=("N", "INSTRUCTION"),
                    help="Edit image N instead of generating. REMOVAL ONLY — "
                         "additions and modifications strip the style, so "
                         "regenerate those with --only N")
    ap.add_argument("--keep", help="What the edit must preserve, e.g. "
                                   "'Keep the rider. Keep the horizon.'")
    ap.add_argument("--model", help=f"Override the spec's model (default {DEFAULT_MODEL})")
    ap.add_argument("--dry-run", action="store_true",
                    help="Print the composed prompts and stop — no API calls, "
                         "no cost")
    a = ap.parse_args()

    spec = json.loads(Path(a.spec).read_text())
    anchor = spec.get("style_anchor", "")
    # plan_broll.py emits slots with empty prompts for the sentences it could
    # not judge. An unfilled slot is not an error — most sentences do not want
    # a picture — so they are dropped here rather than rendered blank.
    images = [i for i in spec.get("images", []) if (i.get("prompt") or "").strip()]
    skipped = len(spec.get("images", [])) - len(images)
    if not images:
        sys.exit(f"{a.spec} has no image with a prompt. If this came from "
                 f"plan_broll.py, fill the prompt on the slots worth "
                 f"illustrating first.")
    if skipped:
        print(f"Skipping {skipped} slot(s) with no prompt.", file=sys.stderr)
    if not (spec.get("style_anchor") or "").strip():
        print("No style_anchor set. Images will not share a look, which is the "
              "usual reason a set reads as stock photography.", file=sys.stderr)
    model = a.model or spec.get("model", DEFAULT_MODEL)
    outdir = Path(a.outdir)

    if a.edit:
        n, instruction = int(a.edit[0]), a.edit[1]
        if not 1 <= n <= len(images):
            sys.exit(f"--edit {n} is outside 1..{len(images)}")
        if not a.keep:
            print("No --keep given. An edit routinely takes neighbours of its "
                  "target with it — name what must survive.", file=sys.stderr)
        src = outdir / f"{images[n - 1]['name']}.png"
        if not src.exists():
            sys.exit(f"{src} does not exist yet — generate it first.")
        prompt = edit_prompt(instruction, a.keep)
        if a.dry_run:
            print(f"--- edit {src.name} ---\n{prompt}")
            return
        dest = next_version(src)
        edit_image(model, src, prompt, dest)
        print(json.dumps({"model": model, "edited": str(src),
                          "written": str(dest)}, indent=2))
        return

    wanted = range(1, len(images) + 1)
    if a.only:
        wanted = [int(n) for n in a.only.split(",")]

    results = []
    for i in wanted:
        item = images[i - 1]
        prompt = compose_prompt(anchor, item["prompt"])
        dest = next_version(outdir / f"{item['name']}.png")
        if a.dry_run:
            print(f"--- {dest.name} ---\n{prompt}\n")
            continue
        print(f"[{i}/{len(images)}] {dest.name}...", flush=True)
        generate(model, prompt, dest)
        results.append(str(dest))

    if a.dry_run:
        return
    print(json.dumps({"model": model, "written": results}, indent=2))
    print("\nThese are stills. normalize_clips.py gives them Ken Burns motion "
          "via the `kb` field — a static still held for three seconds reads as "
          "dead air.", file=sys.stderr)


if __name__ == "__main__":
    main()