#!/usr/bin/env python3
"""
Dr. Heena Libromi — AI Thumbnail Generator v2
Generates high-impact thumbnails WITH text using Gemini 3 Pro Image (Nano Banana Pro)
via Vertex AI global endpoint, then prepends as first frame into MP4.

Usage:
  python generate_thumbnails.py                        # smoke test: anxiety/day01
  python generate_thumbnails.py --all                  # all 42 videos
  python generate_thumbnails.py --topic anxiety        # 14 anxiety videos
  python generate_thumbnails.py --topic sleep --day 3  # specific video
  python generate_thumbnails.py --preview              # smoke test + save preview.jpg
"""

import argparse
import base64
import subprocess
import sys
import tempfile
import time
from pathlib import Path

# ── Config ───────────────────────────────────────────────────────────────────
BASE_DIR  = Path(__file__).parent
PROJECT   = "project-e839b049-4306-4cf9-99d"
MODEL     = "gemini-3-pro-image-preview"   # Nano Banana Pro
ENDPOINT  = f"https://aiplatform.googleapis.com/v1/projects/{PROJECT}/locations/global/publishers/google/models/{MODEL}:generateContent"

VIDEO_W, VIDEO_H = 496, 864   # target video dimensions

# ── Video titles ─────────────────────────────────────────────────────────────
TITLES = {
    "anxiety": [
        "Noticing What Anxiety Feels Like", "Your Anxiety Has a Pattern",
        "Sitting With Anxiety, Gently", "Your Breath Is Your Anchor",
        "Releasing Tension From Your Body", "Grounding Yourself in the Present",
        "Noticing Your Worry Thoughts", "Questioning the What If",
        "Letting Thoughts Pass Like Clouds", "Building Your Calm Routine",
        "Facing Small Discomforts Gently", "Creating Your Personal Toolkit",
        "Seeing How Far You Have Come", "Carrying This Forward",
    ],
    "relationship": [
        "How You Connect With Others", "The Patterns You Carry",
        "What You Bring to Relationships", "When Emotions Rise Between People",
        "Pausing Before Reacting", "Calming Yourself in Conflict",
        "The Stories We Tell Ourselves", "Seeing the Other Person's World",
        "Letting Go of Being Right", "Listening to Understand",
        "Saying What You Need", "Small Acts of Connection",
        "Accepting Imperfection", "Building What Lasts",
    ],
    "sleep": [
        "Understanding Your Sleep", "Your Personal Sleep Story",
        "Making Your Bedroom a Sleep Space", "Releasing the Day From Your Body",
        "Breathing Into Rest", "A Body Scan for Sleep",
        "When Your Mind Will Not Stop", "Letting Go of Sleep Pressure",
        "Rewriting Your Sleep Beliefs", "Building Your Wind-Down Routine",
        "Optimising Your Sleep Environment", "Consistency Over Perfection",
        "Trusting Your Body to Sleep", "Your Sleep Going Forward",
    ],
}

# ── Topic visual styles ───────────────────────────────────────────────────────
TOPIC_STYLE = {
    "anxiety": {
        "bg_mood": "dramatic stormy sky with single golden light ray breaking through dark clouds, deep teal and golden contrast",
        "accent_color": "warm golden yellow",
        "label": "ANXIETY",
        "label_color": "bright warm gold",
    },
    "relationship": {
        "bg_mood": "intimate golden hour bokeh, warm amber and deep shadow contrast, glowing light between two abstract forms",
        "accent_color": "warm rose gold",
        "label": "RELATIONSHIP",
        "label_color": "soft rose gold",
    },
    "sleep": {
        "bg_mood": "deep midnight blue starry sky, single silver moonbeam, ethereal aurora-like glow, dreamlike and mysterious",
        "accent_color": "soft silver lavender",
        "label": "SLEEP",
        "label_color": "silver white",
    },
}

# ── Prompt builder ────────────────────────────────────────────────────────────
def build_prompt(topic: str, day: int, title: str) -> str:
    style = TOPIC_STYLE[topic]
    # Split title into 2 lines max for readability
    words = title.split()
    mid = len(words) // 2
    line1 = " ".join(words[:mid]) if len(words) > 3 else title
    line2 = " ".join(words[mid:]) if len(words) > 3 else ""
    text_lines = f'"{line1}\\n{line2}"' if line2 else f'"{line1}"'

    return f"""Create a high-impact Instagram Reels thumbnail, 9:16 vertical portrait format.

BACKGROUND: {style["bg_mood"]}. Cinematic, dramatic, emotionally striking. Ultra high detail.

TEXT OVERLAY (render clearly on the image):
- Top area: Small label "{style["label"]}" in {style["label_color"]}, bold caps, 9:16 portrait
- Center-bottom: Large bold white text: {text_lines}
- Bottom corner: "DAY {day:02d}" in small {style["accent_color"]} bold text

TYPOGRAPHY STYLE: Bold modern sans-serif, clean and professional. Text should have subtle dark shadow for readability against the background.

COMPOSITION: Background fills entire frame. Text is legible and prominent. High contrast between text and background. No people. No faces. No watermarks."""


# ── Vertex AI auth ────────────────────────────────────────────────────────────
def get_token() -> str:
    import google.auth
    import google.auth.transport.requests
    creds, _ = google.auth.default(scopes=["https://www.googleapis.com/auth/cloud-platform"])
    creds.refresh(google.auth.transport.requests.Request())
    return creds.token


# ── Generate image via Nano Banana Pro ───────────────────────────────────────
def generate_thumbnail(prompt: str, output_path: Path, token: str) -> bool:
    import requests as _req

    for attempt in range(4):  # up to 4 attempts
        resp = _req.post(
            ENDPOINT,
            headers={"Authorization": f"Bearer {token}", "Content-Type": "application/json"},
            json={
                "contents": [{"role": "user", "parts": [{"text": prompt}]}],
                "generationConfig": {"response_modalities": ["IMAGE", "TEXT"]},
            },
            timeout=180,
        )

        if resp.status_code == 429:
            wait = 45 * (attempt + 1)  # 45s, 90s, 135s
            print(f"  ⏳ Rate limit, waiting {wait}s (attempt {attempt+1}/4)...")
            time.sleep(wait)
            continue

        if not resp.ok:
            print(f"  ✗ API {resp.status_code}: {resp.text[:300]}")
            return False

        parts = resp.json().get("candidates", [{}])[0].get("content", {}).get("parts", [])
        img_parts = [p for p in parts if p.get("inlineData", {}).get("mimeType", "").startswith("image/")]

        if not img_parts:
            text_parts = [p.get("text", "") for p in parts if "text" in p]
            print(f"  ✗ No image returned. Text: {text_parts[:1]}")
            return False

        output_path.write_bytes(base64.b64decode(img_parts[0]["inlineData"]["data"]))
        return True

    print(f"  ✗ Failed after 4 attempts (rate limit)")
    return False


# ── Prepend thumbnail as first frame ─────────────────────────────────────────
def prepend_as_first_frame(video_path: Path, thumb_path: Path) -> bool:
    """
    Prepend thumbnail as a ~0.5s still at the start of the video using 2-step concat.
    Step 1: JPEG → tiny h264 clip (ultrafast, 0.5s, no audio)
    Step 2: concat clip + original via stream copy (fast, no re-encode)
    """
    tmp_clip = video_path.parent / f"_thumbclip_{video_path.stem}.mp4"
    tmp_out  = video_path.parent / f"_tmp_{video_path.name}"
    concat_list = video_path.parent / f"_concat_{video_path.stem}.txt"

    try:
        # Step 1: thumbnail JPEG → h264 clip WITH silent audio (preserves audio in concat)
        r1 = subprocess.run([
            "ffmpeg", "-y", "-loglevel", "error",
            "-loop", "1", "-i", str(thumb_path),
            "-f", "lavfi", "-i", "anullsrc=r=48000:cl=stereo",
            "-vf", f"scale={VIDEO_W}:{VIDEO_H}:force_original_aspect_ratio=decrease,"
                   f"pad={VIDEO_W}:{VIDEO_H}:(ow-iw)/2:(oh-ih)/2,format=yuv420p,setsar=1",
            "-c:v", "libx264", "-preset", "ultrafast", "-crf", "23",
            "-c:a", "aac", "-b:a", "128k",
            "-t", "0.5",
            str(tmp_clip),
        ], capture_output=True)

        if r1.returncode != 0:
            print(f"  ✗ ffmpeg step1: {r1.stderr.decode()[:300]}")
            return False

        # Step 2: concat via stream copy (near-instant)
        concat_list.write_text(
            f"file '{tmp_clip.resolve()}'\nfile '{video_path.resolve()}'\n"
        )
        r2 = subprocess.run([
            "ffmpeg", "-y", "-loglevel", "error",
            "-f", "concat", "-safe", "0", "-i", str(concat_list),
            "-c", "copy",
            "-movflags", "+faststart",
            str(tmp_out),
        ], capture_output=True)

        if r2.returncode != 0:
            print(f"  ✗ ffmpeg step2: {r2.stderr.decode()[:300]}")
            return False

        tmp_out.replace(video_path)
        return True

    finally:
        tmp_clip.unlink(missing_ok=True)
        concat_list.unlink(missing_ok=True)
        if tmp_out.exists():
            tmp_out.unlink(missing_ok=True)


# ── Per-video processing ──────────────────────────────────────────────────────
def process_video(topic: str, day: int, token: str, save_preview: Path = None) -> bool:
    video_path = BASE_DIR / topic / f"day{day:02d}.mp4"
    if not video_path.exists():
        print(f"  ✗ Not found: {video_path}")
        return False

    title = TITLES[topic][day - 1]
    print(f"\n[{topic}] day{day:02d} — {title}")

    prompt = build_prompt(topic, day, title)

    with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as tmp:
        thumb_path = Path(tmp.name)

    try:
        t0 = time.time()
        if not generate_thumbnail(prompt, thumb_path, token):
            return False
        size_kb = thumb_path.stat().st_size // 1024
        print(f"  ✓ Thumbnail generated ({size_kb}KB, {time.time()-t0:.1f}s)")

        # Save a copy for preview if requested
        if save_preview:
            import shutil
            shutil.copy2(thumb_path, save_preview)
            print(f"  ✓ Preview saved → {save_preview}")

        if not prepend_as_first_frame(video_path, thumb_path):
            return False
        print(f"  ✓ Embedded as first frame → {video_path.name}")
        return True

    finally:
        thumb_path.unlink(missing_ok=True)


# ── Extract first frame for verification ─────────────────────────────────────
def extract_first_frame(video_path: Path, out_path: Path) -> bool:
    result = subprocess.run([
        "ffmpeg", "-y", "-loglevel", "error",
        "-i", str(video_path),
        "-frames:v", "1",
        str(out_path),
    ], capture_output=True)
    return result.returncode == 0


# ── Main ──────────────────────────────────────────────────────────────────────
def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--all",         action="store_true", help="Process all 42 videos")
    parser.add_argument("--topic",       choices=["anxiety", "relationship", "sleep"])
    parser.add_argument("--day",         type=int)
    parser.add_argument("--preview",     action="store_true", help="Save thumbnail as preview.jpg and first frame")
    parser.add_argument("--start-delay", type=int, default=0, help="Initial wait in seconds before starting")
    args = parser.parse_args()

    if args.start_delay:
        print(f"Waiting {args.start_delay}s before starting (staggered launch)...")
        time.sleep(args.start_delay)

    print(f"Model   : {MODEL}  (Nano Banana Pro)")
    print(f"Project : {PROJECT}")
    print(f"Endpoint: global\n")

    print("Authenticating via ADC...")
    try:
        token = get_token()
        print("✓ Auth OK\n")
    except Exception as e:
        print(f"✗ Auth failed: {e}")
        sys.exit(1)

    # Determine scope
    if args.all:
        tasks = [(t, d) for t in ["anxiety", "relationship", "sleep"] for d in range(1, 15)]
        print(f"Processing all 42 videos...\n")
    elif args.topic and args.day:
        tasks = [(args.topic, args.day)]
    elif args.topic:
        tasks = [(args.topic, d) for d in range(1, 15)]
    else:
        tasks = [("anxiety", 1)]
        print("── SMOKE TEST: anxiety/day01 ──")

    preview_path = BASE_DIR / "thumb_preview.jpg" if (args.preview or not args.all) else None

    # Load progress log (skip already-done videos)
    progress_file = BASE_DIR / ".thumbnail_progress.json"
    done = set()
    if progress_file.exists():
        import json as _json
        done = set(_json.loads(progress_file.read_text()))

    success, fail = 0, 0
    for i, (topic, day) in enumerate(tasks):
        key = f"{topic}/day{day:02d}"
        if key in done:
            print(f"[{topic}] day{day:02d} — already done, skipping")
            success += 1
            continue

        if i > 0 and i % 25 == 0:
            token = get_token()  # refresh token

        save_prev = preview_path if (i == 0 and not done) else None
        ok = process_video(topic, day, token, save_preview=save_prev)
        if ok:
            success += 1
            done.add(key)
            import json as _json
            progress_file.write_text(_json.dumps(sorted(done)))
        else:
            fail += 1
        if len(tasks) > 1:
            time.sleep(5)  # 5s between requests to stay within quota

    print(f"\n{'─'*45}")
    print(f"✓ {success} succeeded   ✗ {fail} failed")

    if not args.all and fail == 0:
        # Extract first frame for visual check
        first_frame = BASE_DIR / "first_frame_check.jpg"
        video = BASE_DIR / "anxiety" / "day01.mp4"
        if extract_first_frame(video, first_frame):
            print(f"\nFirst frame extracted → {first_frame.name}")
            print(f"Check at: https://videos.dramankhanna.cloud/first_frame_check.jpg")
        print(f"\nThumbnail preview → https://videos.dramankhanna.cloud/thumb_preview.jpg")
        print(f"\nIf happy, run: python generate_thumbnails.py --all")

    return 0 if fail == 0 else 1


if __name__ == "__main__":
    sys.exit(main())