"""Hero reel v4 — 6 precise shots."""
import os, sys, time, subprocess

try:
    from google import genai
    from google.genai import types
except ImportError:
    os.system(f"{sys.executable} -m pip install google-genai --quiet")
    from google import genai
    from google.genai import types

API_KEY  = os.environ.get("GEMINI_API_KEY", "AIzaSyB_1p8KvaT_rdNJGPs8HKk8bKsvUlcL6Kg")
BASE_DIR = os.path.dirname(os.path.dirname(__file__))
VID_DIR  = os.path.join(BASE_DIR, "assets", "videos", "hero", "clips")
REEL_OUT = os.path.join(BASE_DIR, "assets", "videos", "hero", "hero-reel.mp4")
os.makedirs(VID_DIR, exist_ok=True)
client   = genai.Client(api_key=API_KEY)

SHOTS = [
    {
        "name": "v4-shot-01",
        "prompt": (
            "Medium cinematic shot. A family — two adults and two children — walks through a front door "
            "into a residential home. The camera follows their feet as they step onto the beige carpet "
            "in the entryway, then slowly pans down to show their shoes leaving dirty tracks on the carpet. "
            "Warm afternoon light. Photorealistic, slow motion."
        ),
    },
    {
        "name": "v4-shot-02",
        "prompt": (
            "Slow-motion close-up cinematic shot. A wine glass tips over on a light grey fabric sofa cushion. "
            "Red wine pours out of the glass and spreads across the sofa cushion, soaking into the fabric. "
            "The dark red stain expands slowly across the grey upholstery. "
            "Warm living room light. No people visible. Photorealistic, dramatic slow motion."
        ),
    },
    {
        "name": "v4-shot-03",
        "prompt": (
            "Cinematic close-up shot slowly pushing in on a section of heavily soiled residential carpet. "
            "The beige carpet has multiple visible stains — dark spots, discoloration, pet stains, "
            "and general dirt buildup embedded in the fibers. "
            "Dramatic side lighting emphasizes the depth of the stains. No people. No equipment. Photorealistic."
        ),
    },
    {
        "name": "v4-shot-04",
        "prompt": (
            "Close-up cinematic shot. A carpet cleaning technician in a plain black shirt pushes "
            "a large upright carpet cleaning machine — like a Rug Doctor — across a dirty beige carpet. "
            "Tight shot focused on the wide flat cleaning head at the base of the machine pressing against "
            "the carpet and moving forward. The carpet behind the machine looks visibly cleaner and brighter. "
            "The machine only pulls dirt and moisture INTO itself — nothing comes out. "
            "No steam. No liquid leaving the machine. Photorealistic slow motion."
        ),
    },
    {
        "name": "v4-shot-05",
        "prompt": (
            "Wide cinematic shot slowly pushing forward through the main entrance of a modern commercial "
            "office building. Clean grey carpet covers the entire lobby floor. Glass doors, white walls, "
            "professional overhead lighting. The carpet is the visual centerpiece — clean, uniform, well-maintained. "
            "No people. No machines. Photorealistic."
        ),
    },
    {
        "name": "v4-shot-06",
        "prompt": (
            "Cinematic wide shot inside a bright clean residential living room. "
            "The camera slowly pans upward from the clean plush beige carpet to reveal the whole room. "
            "A family — two adults and a child — walks in at different moments and relaxes on the sofa. "
            "Everyone is wearing socks. The carpet is spotless and fluffy. Warm natural light. "
            "Comfortable, inviting atmosphere. Photorealistic."
        ),
    },
]

MODELS = ["veo-3.1-generate-preview", "veo-2.0-generate-001"]

def poll(op, timeout=600):
    elapsed = 0
    while not op.done:
        if elapsed >= timeout:
            return None
        print(f"  Waiting... ({elapsed}s)")
        time.sleep(15)
        elapsed += 15
        op = client.operations.get(op)
    return op

saved = []
for i, item in enumerate(SHOTS):
    out_path = os.path.join(VID_DIR, f"{item['name']}.mp4")
    print(f"\n[{i+1}/{len(SHOTS)}] {item['name']}...")
    done = False
    for model in MODELS:
        try:
            print(f"  Trying {model}...")
            op = client.models.generate_videos(
                model=model,
                prompt=item["prompt"],
                config=types.GenerateVideosConfig(
                    aspect_ratio="16:9", resolution="720p",
                    duration_seconds=6, number_of_videos=1,
                ),
            )
            op = poll(op)
            if op and op.response and op.response.generated_videos:
                vid = op.response.generated_videos[0].video
                video_bytes = client.files.download(file=vid)
                if video_bytes:
                    with open(out_path, "wb") as f:
                        f.write(video_bytes)
                    print(f"  Saved ({os.path.getsize(out_path)//1024}KB)")
                    saved.append(out_path)
                    done = True
                    break
                else:
                    try:
                        vid.save(out_path)
                        print(f"  Saved via .save()")
                        saved.append(out_path)
                        done = True
                        break
                    except Exception as e2:
                        print(f"  Download failed: {e2}")
            else:
                print(f"  No video returned from {model}")
        except Exception as e:
            print(f"  Error with {model}: {e}")
    if not done:
        print(f"  FAILED: {item['name']}")

print(f"\n{len(saved)}/{len(SHOTS)} shots saved")

if len(saved) >= 2:
    clips = [os.path.join(VID_DIR, f"{s['name']}.mp4") for s in SHOTS
             if os.path.exists(os.path.join(VID_DIR, f"{s['name']}.mp4"))]
    concat_file = os.path.join(VID_DIR, "concat-v4.txt")
    with open(concat_file, "w") as f:
        for p in clips:
            f.write(f"file '{p}'\n")
    result = subprocess.run(
        ["ffmpeg", "-y", "-f", "concat", "-safe", "0", "-i", concat_file,
         "-c:v", "libx264", "-crf", "22", "-preset", "fast", "-movflags", "+faststart", REEL_OUT],
        capture_output=True, text=True
    )
    if result.returncode == 0:
        print(f"  Reel saved ({os.path.getsize(REEL_OUT)//1024}KB)")
    else:
        print(f"  ffmpeg error: {result.stderr[-300:]}")

print("\nDone.")