"""
Regenerate one specific shot and re-concatenate the hero reel.
Usage: python3 tools/regen-shot.py shot-02-staircase
"""
import os, sys, time, subprocess
from google import genai
from google.genai import types
import urllib.request

API_KEY  = os.environ.get("GEMINI_API_KEY", "")
BASE_DIR = os.path.dirname(os.path.dirname(__file__))
CLIPS_DIR = os.path.join(BASE_DIR, "assets", "videos", "hero", "clips")
REEL_OUT  = os.path.join(BASE_DIR, "assets", "videos", "hero", "hero-reel.mp4")

if not API_KEY:
    print("Set GEMINI_API_KEY env var"); sys.exit(1)

# Import SHOTS from gen-video.py (dash in name requires importlib)
import importlib.util
spec = importlib.util.spec_from_file_location(
    "gen_video", os.path.join(os.path.dirname(__file__), "gen-video.py")
)
_mod = importlib.util.module_from_spec(spec)
spec.loader.exec_module(_mod)
SHOTS = _mod.SHOTS

shot_name = sys.argv[1] if len(sys.argv) > 1 else None
target = next((s for s in SHOTS if s["name"] == shot_name), None)
if not target:
    print(f"Shot '{shot_name}' not found. Available: {[s['name'] for s in SHOTS]}")
    sys.exit(1)

client = genai.Client(api_key=API_KEY)
out_path = os.path.join(CLIPS_DIR, f"{target['name']}.mp4")

print(f"Regenerating {target['name']}...")
op = client.models.generate_videos(
    model="veo-3.0-generate-001",
    prompt=target["prompt"],
    config=types.GenerateVideosConfig(
        aspect_ratio="16:9", resolution="720p",
        duration_seconds=6, number_of_videos=1,
    ),
)
elapsed = 0
while not op.done:
    print(f"  Waiting... ({elapsed}s)")
    time.sleep(15); elapsed += 15
    op = client.operations.get(op)

if not (op.response and op.response.generated_videos):
    print("No video returned"); sys.exit(1)

vid = op.response.generated_videos[0].video
video_bytes = None
try:
    video_bytes = client.files.download(file=vid)
except Exception:
    pass
if video_bytes:
    with open(out_path, "wb") as f: f.write(video_bytes)
elif hasattr(vid, "uri") and vid.uri:
    uri = vid.uri + ("&" if "?" in vid.uri else "?") + f"key={API_KEY}"
    urllib.request.urlretrieve(uri, out_path)
else:
    print("Download failed"); sys.exit(1)

print(f"Saved {out_path} ({os.path.getsize(out_path)//1024}KB)")

# Re-concat in shot order
clips = [os.path.join(CLIPS_DIR, f"{s['name']}.mp4") for s in SHOTS
         if os.path.exists(os.path.join(CLIPS_DIR, f"{s['name']}.mp4"))]
list_file = os.path.join(CLIPS_DIR, "concat.txt")
with open(list_file, "w") as f:
    for c in clips: f.write(f"file '{c}'\n")

print(f"Concatenating {len(clips)} clips...")
r = subprocess.run(
    ["ffmpeg", "-y", "-f", "concat", "-safe", "0", "-i", list_file,
     "-c:v", "libx264", "-crf", "22", "-preset", "fast",
     "-movflags", "+faststart", REEL_OUT],
    capture_output=True, text=True
)
if r.returncode == 0:
    print(f"Reel updated: {REEL_OUT} ({os.path.getsize(REEL_OUT)//1024}KB)")
else:
    print(f"ffmpeg error: {r.stderr[-200:]}")