#!/usr/bin/env python3
"""Stylize a reel thumbnail with Nano Banana (Gemini image editing).

Takes the REAL frame grab (authenticity first — the skill's rule) and
restyles it as a sport-magazine cover. Two styles to start:

  modern  crisp contemporary sports-magazine editorial photo
  80s     1980s sports-magazine print aesthetic (grain, halftone, faded warmth)

Without --text you get a photo treatment only. With --text "HOOK LINE" the
model renders a full cover: masthead + bold headline typography (this is why
the default model is Nano Banana Pro — it renders text reliably).

Requires an API key in GEMINI_API_KEY (or GOOGLE_AI_API_KEY / GOOGLE_API_KEY).
Models (check /v1beta/models if a 404 comes back):
  gemini-3-pro-image        Nano Banana Pro — best text + instruction following (default)
  gemini-3.1-flash-image    Nano Banana 2 — faster/cheaper, weaker text

The subject's identity is explicitly preserved in the prompt, but ALWAYS eyeball
the result before publishing, and remind the user that AI-stylized imagery may
warrant platform disclosure.

Usage:
  stylize_thumbnail.py                                   # exports/thumbnail.jpg, both styles
  stylize_thumbnail.py --style 80s --text "NO EXCUSES" --masthead "HYROX PREP"
  stylize_thumbnail.py --video work/master.mp4 --at 3.2 --style modern
  stylize_thumbnail.py --dry-run                         # print request JSON, no network
"""
import argparse, base64, json, os, subprocess, sys, urllib.error, urllib.request
from pathlib import Path

API = "https://generativelanguage.googleapis.com/v1beta/models/{model}:generateContent"

IDENTITY = ("Keep the person's face, identity, body, pose, clothing and the "
            "overall composition unchanged — this is a real photo being "
            "restyled, not a new image. Keep 9:16 portrait framing.")

# used instead of the framing sentence when the input is landscape/square:
# the model must recompose, and it must know cropping is allowed so it
# doesn't invent prominent new scenery to fill the frame
REFRAME = ("Keep the person's face, identity, body, pose and clothing "
           "unchanged — this is a real photo being restyled, not a new "
           "image. Recompose the landscape frame to 9:16 portrait centered "
           "on the person: crop into the scene as needed and keep them "
           "prominent; do not add new prominent objects or scenery.")

STYLES = {
    "modern": (
        "Restyle this photo as a modern premium sports magazine editorial "
        "cover photo: dramatic rim lighting, deep cinematic contrast, rich "
        "but controlled color grade, subtle background cleanup, crisp "
        "detail, high-end retouching. " + IDENTITY
    ),
    "80s": (
        "CRITICAL, highest priority: the person's face must remain "
        "photographically identical to the input — same facial geometry, "
        "features, skin tone, expression and apparent age. Do not re-draw, "
        "beautify, age or stylize the face in any way. "
        "With that constraint: apply a 1980s sports magazine print finish "
        "OVER this photo (a darkroom/print treatment, not a re-imagining). "
        "Make the era unmistakable everywhere EXCEPT the face: bold visible "
        "halftone dot texture across the background and clothing, film "
        "grain, warm faded Kodachrome-style color shift, slightly crushed "
        "vintage contrast, aged paper texture and a thin off-white border "
        "at the edges. On the face keep only the gentle color grade — no "
        "halftone dots, no grain, no changes to the features. " + IDENTITY
    ),
}

COVER_TEXT = {
    "modern": (" Turn it into a full magazine cover: masthead \"{masthead}\" "
               "across the top in a clean bold modern sans-serif, and the "
               "headline \"{text}\" in large condensed type near the lower "
               "third. Render all text EXACTLY as given, correctly spelled, "
               "fully inside the frame with margins clear of the edges."),
    "80s": (" Turn it into a full magazine cover: masthead \"{masthead}\" "
            "across the top in a bold 1980s serif or slab typeface with a "
            "slight drop shadow, and the headline \"{text}\" in chunky retro "
            "cover-line type, maybe on a color bar. Render all text EXACTLY "
            "as given, correctly spelled, fully inside the frame with "
            "margins clear of the edges."),
}


def is_landscape(image):
    r = subprocess.run(["ffprobe", "-v", "quiet", "-show_entries",
                        "stream=width,height", "-of", "csv=p=0", str(image)],
                       capture_output=True, text=True)
    try:
        w, h = map(int, r.stdout.strip().split("\n")[0].split(",")[:2])
        return w >= h
    except (ValueError, IndexError):
        return False    # can't probe — assume portrait, prompt stays safe


def grab_frame(video, at, outdir):
    dst = Path(outdir) / "thumbnail.jpg"
    tmp = Path(outdir) / "thumbnail.tmp.jpg"
    r = subprocess.run(["ffmpeg", "-y", "-v", "error", "-ss", str(at),
                        "-i", str(video), "-frames:v", "1", "-q:v", "2", str(tmp)],
                       capture_output=True, text=True)
    if r.returncode != 0:
        sys.exit(f"frame grab failed:\n{r.stderr[-400:]}")
    tmp.rename(dst)
    return dst


def request_body(image_path, prompt):
    data = base64.b64encode(Path(image_path).read_bytes()).decode()
    return {
        "contents": [{"parts": [
            {"inline_data": {"mime_type": "image/jpeg", "data": data}},
            {"text": prompt},
        ]}],
        "generationConfig": {"imageConfig": {"aspectRatio": "9:16"}},
    }


def call_api(model, body, key):
    req = urllib.request.Request(
        API.format(model=model),
        data=json.dumps(body).encode(),
        headers={"x-goog-api-key": key, "Content-Type": "application/json"})
    try:
        with urllib.request.urlopen(req, timeout=180) as r:
            return json.loads(r.read())
    except urllib.error.HTTPError as e:
        detail = e.read().decode(errors="replace")[:600]
        hint = ("\nModel names change — list current ones: curl -H 'x-goog-api-key: "
                "$GEMINI_API_KEY' https://generativelanguage.googleapis.com/v1beta/models"
                if e.code == 404 else "")
        sys.exit(f"Gemini API error {e.code}:\n{detail}{hint}")
    except urllib.error.URLError as e:
        sys.exit(f"Gemini API unreachable: {e.reason}")


def extract_image(resp):
    for cand in resp.get("candidates", []):
        for part in cand.get("content", {}).get("parts", []):
            blob = part.get("inlineData") or part.get("inline_data")
            if blob and blob.get("data"):
                return base64.b64decode(blob["data"])
    block = resp.get("promptFeedback", {}).get("blockReason")
    text = next((p.get("text") for c in resp.get("candidates", [])
                 for p in c.get("content", {}).get("parts", []) if p.get("text")), None)
    sys.exit("No image in response"
             + (f" (blocked: {block})" if block else "")
             + (f" — model said: {text[:300]}" if text else ""))


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("image", nargs="?", default="exports/thumbnail.jpg",
                    help="Frame to stylize (default exports/thumbnail.jpg)")
    ap.add_argument("--video", help="Grab the frame from this video instead")
    ap.add_argument("--at", type=float, default=1.5,
                    help="Timestamp for --video frame grab (default 1.5)")
    ap.add_argument("--style", default="all", choices=[*STYLES, "all"])
    ap.add_argument("--text", help="Cover headline — enables full magazine-cover "
                                   "treatment (masthead + typography)")
    ap.add_argument("--masthead", help="Masthead name for --text mode "
                                       "(default: project folder name)")
    ap.add_argument("--model", default="gemini-3-pro-image",
                    help="gemini-3-pro-image (default) | gemini-3.1-flash-image")
    ap.add_argument("--outdir", default="exports")
    ap.add_argument("--dry-run", action="store_true",
                    help="Print request JSON (image bytes elided), no network")
    args = ap.parse_args()

    outdir = Path(args.outdir)
    outdir.mkdir(parents=True, exist_ok=True)
    src = grab_frame(args.video, args.at, outdir) if args.video else Path(args.image)
    if not src.exists():
        sys.exit(f"{src} not found — run export_variants.py first, or pass "
                 "--video work/master.mp4")

    key = next((os.environ[v] for v in
                ("GEMINI_API_KEY", "GOOGLE_AI_API_KEY", "GOOGLE_API_KEY")
                if os.environ.get(v)), None)
    if not key and not args.dry_run:
        sys.exit("No API key found (checked GEMINI_API_KEY, GOOGLE_AI_API_KEY, "
                 "GOOGLE_API_KEY) — stylized thumbnails need a Gemini API key "
                 "(https://aistudio.google.com/apikey). The plain frame grab "
                 "still works without it.")

    masthead = (args.masthead or Path.cwd().name.replace("-", " ").upper())
    styles = list(STYLES) if args.style == "all" else [args.style]
    reframe = is_landscape(src)
    if reframe:
        print(f"{src.name} is landscape — the model will recompose to 9:16 "
              "around the subject (some background gets synthesized; check it)")
    made = []
    for style in styles:
        prompt = STYLES[style]
        if reframe:
            prompt = prompt.replace(IDENTITY, REFRAME)
        if args.text:
            prompt += COVER_TEXT[style].format(masthead=masthead, text=args.text)
        body = request_body(src, prompt)
        if args.dry_run:
            body["contents"][0]["parts"][0]["inline_data"]["data"] = \
                f"<{src} base64 elided>"
            print(f"--- {style} -> {API.format(model=args.model)}")
            print(json.dumps(body, indent=2))
            continue
        print(f"Stylizing ({style}, {args.model})...", flush=True)
        img = extract_image(call_api(args.model, body, key))
        dst = outdir / f"thumbnail-{style}.jpg"
        tmp = outdir / f"thumbnail-{style}.tmp.jpg"
        tmp.write_bytes(img)
        tmp.rename(dst)
        made.append(dst)
        print(f"  {dst} ({len(img)/1e3:.0f} KB)")

    if made:
        print("\nPresent ALL options to the user — plain frame grab + styled "
              "variants — and let them pick the cover. Check faces for "
              "fidelity before publishing; AI-stylized imagery may warrant "
              "platform disclosure.")


if __name__ == "__main__":
    main()