#!/usr/bin/env python3
"""
generate-clips.py — Async video clip generation for THE TAKEOVER
Submits all shots to fal.ai queue, persists request IDs, polls for completion.

CRASH SAFE: request IDs are saved to clips/.request-ids.json immediately after
submission. Re-running after a timeout/crash will resume polling instead of
re-submitting (and re-charging) the same jobs.

Usage:
  FAL_KEY=your_key_here python3 scripts/generate-clips.py
  -- or --
  Add FAL_KEY to .env and run inside devenv shell:
  python3 scripts/generate-clips.py

Requires: pip install fal-client
"""

import asyncio
import json
import os
import sys
import time
import urllib.request
from pathlib import Path

try:
    import fal_client
except ImportError:
    print("ERROR: fal-client not installed. Run: pip install fal-client")
    sys.exit(1)

FAL_KEY = os.environ.get("FAL_KEY")
if not FAL_KEY:
    print("ERROR: FAL_KEY not set. Add it to .env or export it.")
    sys.exit(1)

os.environ["FAL_KEY"] = FAL_KEY

CLIPS_DIR = Path(__file__).parent / "storyboard-the-takeover" / "clips"
CLIPS_DIR.mkdir(parents=True, exist_ok=True)

REQUEST_IDS_FILE = CLIPS_DIR / ".request-ids.json"

MODEL = "fal-ai/kling-video/v3/pro/image-to-video"

SHOTS = [
    {
        "id": "01-establishing",
        "duration": 7,
        "image_url": "https://v3b.fal.media/files/b/0a91018c/sSMmFddy7YtVZI5BobJhN_01-establishing.jpg",
        "prompt": "Camera slowly pushes in on the party room. Warm amber light, people moving naturally. Woman near bookshelf stands slightly still amid the motion. Subtle parallax. Cinematic 35mm, shallow depth of field.",
    },
    {
        "id": "02-nora-before",
        "duration": 5,
        "image_url": "https://v3b.fal.media/files/b/0a91018c/27eRVS452NqC2OxyV7ekb_02-nora-before.jpg",
        "prompt": "Close on the woman's face. Her eyes wander slightly, unfocused. Her lips part as if about to speak but don't. A barely perceptible pause. She is present but not quite here. 35mm close-up, shallow depth of field.",
    },
    {
        "id": "03-the-shift",
        "duration": 5,
        "image_url": "https://v3b.fal.media/files/b/0a91018c/Fq39KUPN22hgF0aq1onqC_03-the-shift.jpg",
        "prompt": "Close insert on hands. The wine glass moves slowly from loose fingertips to a firm controlled palm grip. The hand settles completely still. This is the moment everything changes. 35mm close-up, warm light.",
    },
    {
        "id": "04-micha-speaks",
        "duration": 7,
        "image_url": "https://v3b.fal.media/files/b/0a91018c/errWdqNrr5sNLKC95ll97_04-micha-speaks.jpg",
        "prompt": "Woman speaks with precise measured confidence. The other woman listens, visibly impressed and slightly off-balance. The speaker is unnervingly still and calibrated. No hesitation, no trailing off. The listener smiles too quickly. 35mm medium shot.",
    },
    {
        "id": "05-the-scan",
        "duration": 7,
        "image_url": "https://v3b.fal.media/files/b/0a91018c/uqCNkvumiBI4lt0EVKm9K_05-the-scan.jpg",
        "prompt": "From behind a woman as she moves through a party. Her head turns slowly, surveying the room with deliberate precision. Others move naturally around her. This is not socialising — it is assessment. 35mm wide.",
    },
    {
        "id": "06-david",
        "duration": 7,
        "image_url": "https://v3b.fal.media/files/b/0a91018c/v9FknHz6BzXacKSYQLzQ-_06-david.jpg",
        "prompt": "Woman faces an older man at a party. He slowly checks his phone. His face changes — stunned recognition. She waits, absolutely still, already knowing what he will find. He looks up at her differently. 35mm medium.",
    },
    {
        "id": "07-the-window",
        "duration": 7,
        "image_url": "https://v3b.fal.media/files/b/0a91018c/hBp_8GSPkCZKmihICM8PI_07-the-window.jpg",
        "prompt": "Woman stands alone by the dark window. She slowly raises her hand, pressing fingers flat against the glass, and studies them. She turns them over. The party continues in soft blur behind her. 35mm wide, cool light near the glass.",
    },
    {
        "id": "08-the-reflection",
        "duration": 7,
        "image_url": "https://v3b.fal.media/files/b/0a91018c/psxLjtjqKMqX3JVbu8Axe_08-the-reflection.jpg",
        "prompt": "The woman's face and its reflection in the dark window hold each other's gaze. Both completely still. Then she blinks — and the reflection blinks with her, exactly on time. She exhales slowly. 35mm extreme close-up, high contrast, cold dark tones.",
    },
    {
        "id": "09-the-hallway",
        "duration": 7,
        "image_url": "https://v3b.fal.media/files/b/0a91018c/zj8uPA9uxCLoycpNX05qb_09-the-hallway.jpg",
        "prompt": "Woman leans against the hallway wall. She slowly takes out her phone and opens an app. Her eyes drop to the screen. Half-closed, deeply tired. Muffled party sounds through the door. Dim corridor light. 35mm medium.",
    },
    {
        "id": "10-the-app",
        "duration": 5,
        "image_url": "https://v3b.fal.media/files/b/0a91018c/gwX7PBS8cNiR4-1pGb02p_10-the-app.jpg",
        "prompt": "Close on a phone screen showing a minimal chat interface. Text appears, then a response below it. The cursor blinks once. The hand holding the phone barely trembles. Long still hold. Intimate close-up insert.",
    },
]


def load_request_ids() -> dict:
    if REQUEST_IDS_FILE.exists():
        return json.loads(REQUEST_IDS_FILE.read_text())
    return {}


def save_request_id(shot_id: str, request_id: str):
    ids = load_request_ids()
    ids[shot_id] = request_id
    REQUEST_IDS_FILE.write_text(json.dumps(ids, indent=2))


async def submit_shot(shot: dict, saved_ids: dict) -> tuple[str, str | None]:
    """Submit a job and return (shot_id, request_id). Skips if already submitted."""
    shot_id = shot["id"]

    if shot_id in saved_ids:
        print(f"  ↩  {shot_id} — already submitted (request_id: {saved_ids[shot_id][:16]}…)")
        return shot_id, saved_ids[shot_id]

    out_path = CLIPS_DIR / f"{shot_id}.mp4"
    if out_path.exists():
        print(f"  ⏭  {shot_id} — clip already exists, skipping")
        return shot_id, "DONE"

    print(f"  📤 {shot_id} — submitting…")
    try:
        handler = await fal_client.submit_async(
            MODEL,
            arguments={
                "image_url": shot["image_url"],
                "prompt": shot["prompt"],
                "duration": str(shot["duration"]),
                "aspect_ratio": "16:9",
                "negative_prompt": "blur, distortion, text, watermarks, fast motion",
            },
        )
        request_id = handler.request_id
        save_request_id(shot_id, request_id)
        print(f"  ✓  {shot_id} — queued (request_id: {request_id[:16]}…)")
        return shot_id, request_id
    except Exception as e:
        print(f"  ✗  {shot_id} — submit failed: {e}")
        return shot_id, None


async def poll_and_download(shot_id: str, request_id: str) -> bool:
    """Poll for completion and download. Returns True on success."""
    if request_id == "DONE":
        return True

    out_path = CLIPS_DIR / f"{shot_id}.mp4"
    if out_path.exists():
        print(f"  ⏭  {shot_id} — already downloaded")
        return True

    print(f"  ⏳ {shot_id} — waiting for result…")
    try:
        result = await fal_client.result_async(MODEL, request_id)
        video_url = (result.get("video") or {}).get("url") or result.get("url")
        if not video_url:
            print(f"  ✗  {shot_id} — no video URL in result: {result}")
            return False

        urllib.request.urlretrieve(video_url, out_path)
        size_kb = out_path.stat().st_size // 1024
        print(f"  ✓  {shot_id} — {size_kb}KB saved to {out_path.name}")

        # Clean up request ID record
        ids = load_request_ids()
        ids.pop(shot_id, None)
        REQUEST_IDS_FILE.write_text(json.dumps(ids, indent=2))
        return True
    except Exception as e:
        print(f"  ✗  {shot_id} — poll failed: {e}")
        return False


async def main():
    print(f"\n🎬 THE TAKEOVER — clip generation")
    print(f"   Model: {MODEL}")
    print(f"   Output: {CLIPS_DIR}\n")

    saved_ids = load_request_ids()

    # Phase 1: submit all jobs (idempotent — skips already-submitted)
    print("── Phase 1: Submit ──")
    submit_tasks = [submit_shot(shot, saved_ids) for shot in SHOTS]
    submitted = await asyncio.gather(*submit_tasks)

    pending = [(sid, rid) for sid, rid in submitted if rid and rid != "DONE"]
    skipped = [(sid, rid) for sid, rid in submitted if rid == "DONE"]
    failed_submit = [(sid, rid) for sid, rid in submitted if not rid]

    print(f"\n   {len(pending)} queued, {len(skipped)} already done, {len(failed_submit)} failed to submit\n")

    if failed_submit:
        print("Failed submissions:")
        for sid, _ in failed_submit:
            print(f"  ✗ {sid}")

    if not pending:
        print("Nothing to poll.")
    else:
        # Phase 2: poll for results (sequential to avoid hammering API)
        print("── Phase 2: Download ──")
        successes = 0
        for shot_id, request_id in pending:
            ok = await poll_and_download(shot_id, request_id)
            if ok:
                successes += 1
            else:
                print(f"     Re-run this script to retry failed clips.")

    # Final ffmpeg command
    clips = sorted(CLIPS_DIR.glob("*.mp4"))
    if clips:
        print(f"\n── {len(clips)}/10 clips ready ──")
        print(f"\nAssemble with ffmpeg:")
        print(f"  cd {CLIPS_DIR}")
        print(f"  ls -1 *.mp4 | sort | sed 's/^/file /' > concat.txt")
        print(f"  ffmpeg -f concat -safe 0 -i concat.txt -c copy ../the-takeover-draft.mp4")


if __name__ == "__main__":
    asyncio.run(main())
