#!/usr/bin/env python3
"""Pass 13 comparison renders — 5 configurations.

Pass 13 (EP001_PASS_013_A_WREN): Wren + Torch, 3 segments, 12s total.
  Seg 0: MS — Wren holding Torch by throat (4s)
  Seg 1: CU — Wren scanning Torch's face (3s)
  Seg 2: MCU — Wren stiff robotic processing (5s)

Renders:
  1. Kling V3 multi-shot (3 segments as sequence)
  2. Kling O3 multi-shot with elements (WREN + TORCH refs)
  3. Kling V3 continuous single-take (combined prompt, 10s)
  4. Veo 3.1 multi-shot (3 segments sequential)
  5. Veo 3.1 continuous single-take (combined prompt + reference_images)
"""
import sys
import time
from pathlib import Path

STARSEND_ROOT = Path(__file__).parent.parent
sys.path.insert(0, str(STARSEND_ROOT))

from lib.execution_store import ExecutionStore
from lib.elements import ElementManager, extract_batch_location
from lib.prompt_engine import build_multi_prompt_sequence
from orchestrator.step_runner import StepRunner
from orchestrator.step_types import ProjectPaths

PROJECT = "tartarus"
EPISODE = 1

# Pass 13 segments
SEGMENTS = [
    {
        "shot_id": "EP001_SH29",
        "shot_type": "MS",
        "duration_s": 4,
        "prompt": "MS: Wren standing, right hand gripping Torch's throat, left arm at side, head tilted toward Torch; Torch standing, both hands gripping Wren's right forearm, head facing Wren android head tilting slightly, woman gripping his arm, suspended in air [uncanny valley, helpless suspension, terrifying stillness]",
    },
    {
        "shot_id": "EP001_SH27",
        "shot_type": "CU",
        "duration_s": 3,
        "prompt": "CU: Wren standing, head tilted down, gaze tracking across Torch's face; Torch standing, head tilted up, gaze fixed on Wren's eyes eyes darting in micro-movements, scanning, analyzing [cold calculation, machine logic, predatory focus]",
    },
    {
        "shot_id": "EP001_SH28",
        "shot_type": "MCU",
        "duration_s": 5,
        "prompt": "MCU: Wren standing rigidly, right hand resting out of frame, left hand resting out of frame, head facing screen-left with vacant gaze stiff vocalization, processing pause, unblinking stare [robotic coldness, disorientation, computational logic]",
    },
]

COMBINED_PROMPT = (
    "Continuous single-take, no cuts. Dim futuristic maintenance shaft. "
    "MS: Wren stands holding Torch by the throat, android head tilting slightly, "
    "Torch gripping his forearm, suspended in air — uncanny valley stillness. "
    "Camera moves closer. Wren's head tilts down, gaze tracking across Torch's face, "
    "eyes darting in micro-movements, scanning, analyzing — cold calculation, "
    "machine logic, predatory focus. "
    "Wren's expression goes vacant, head facing screen-left, stiff body language, "
    "processing pause, unblinking stare — robotic coldness, computational logic."
)

ELEMENT_CHAR_IDS = ["WREN", "TORCH"]
LOCATION_ID = "int_lower_decks_maintenance_shaft"

# Ref images for Veo (max 3)
TORCH_HERO = "/Users/joeturnerlin/Dropbox/CLAUDE_PROJECTS/projects/tartarus/output/refs/characters/torch/hero_beauty.png"
WREN_FRONT = "/Users/joeturnerlin/Dropbox/CLAUDE_PROJECTS/projects/tartarus/output/refs/characters/wren/wren_front.png"


def build_plan_shots():
    """Build plan_shots dicts matching what the sequence endpoint expects."""
    shots = []
    for seg in SEGMENTS:
        shots.append({
            "shot_id": seg["shot_id"],
            "shot_type": seg["shot_type"],
            "duration": seg["duration_s"],
            "_api_duration": max(seg["duration_s"], 3),
            "routing_data": {"target_editorial_duration_s": seg["duration_s"]},
            "location_id": LOCATION_ID,
        })
    return shots


def build_prompt_overrides():
    """Build prompt_overrides dict so build_multi_prompt_sequence uses
    the actual Pass 13 coverage prompts, not generic plan prompts."""
    return {seg["shot_id"]: seg["prompt"] for seg in SEGMENTS}


def reset_shot_status(store, shot_id):
    """Reset shot to a state allowing video_submitted transition."""
    try:
        shot = store.get_shot(shot_id)
        if shot:
            current = shot.get("status", "")
            if current not in ("video_pending", "previs_approved"):
                store.update_shot(shot_id, status="video_pending")
    except Exception:
        pass


def run_kling_multishot(runner, store, plan_shots, model, elements_payload=None, label="",
                        prompt_overrides=None):
    """Run a Kling multi-shot sequence render."""
    shot_ids = [s["shot_id"] for s in plan_shots]
    for sid in shot_ids:
        reset_shot_status(store, sid)

    sequence = build_multi_prompt_sequence(
        plan_shots,
        batch_char_ids=sorted(ELEMENT_CHAR_IDS) if elements_payload else None,
        has_location_element=False,
        total_elements=len(elements_payload.get("elements", [])) if elements_payload else 0,
        prompt_overrides=prompt_overrides,
    )

    print(f"\n{'='*60}")
    print(f"  {label}")
    print(f"  Model: {model}")
    print(f"  Shots: {len(plan_shots)}")
    print(f"  Elements: {'yes' if elements_payload else 'no'}")
    print(f"{'='*60}")

    results = runner.execute_multi_shot(
        batch=plan_shots,
        multi_prompt_sequence=sequence,
        model=model,
        aspect_ratio="9:16",
        elements_payload=elements_payload,
    )

    ok = sum(1 for r in results if r.success)
    cost = sum(r.cost_usd for r in results)
    print(f"  Result: {ok}/{len(results)} ok, ${cost:.3f}")
    for r in results:
        if r.success:
            print(f"    {r.shot_id}: {r.output_path}")
        else:
            print(f"    {r.shot_id}: FAIL — {r.error}")
    return results


def run_single_take(runner, store, shot_id, model, duration, prompt, label="",
                    reference_images=None):
    """Run a single continuous take."""
    reset_shot_status(store, shot_id)

    print(f"\n{'='*60}")
    print(f"  {label}")
    print(f"  Model: {model}")
    print(f"  Duration: {duration}s")
    print(f"  Refs: {len(reference_images) if reference_images else 0}")
    print(f"{'='*60}")

    result = runner.execute_video(
        shot_id=shot_id,
        prompt=prompt,
        model=model,
        duration=duration,
        aspect_ratio="9:16",
        reference_images=reference_images,
    )

    if result.success:
        print(f"  OK: {result.output_path} (${result.cost_usd:.3f})")
    else:
        print(f"  FAIL: {result.error}")
    return result


def run_veo_multishot(runner, store, plan_shots, label=""):
    """Run Veo multi-shot (sequential fallback — one shot at a time with refs)."""
    shot_ids = [s["shot_id"] for s in plan_shots]
    for sid in shot_ids:
        reset_shot_status(store, sid)

    sequence = build_multi_prompt_sequence(plan_shots)

    print(f"\n{'='*60}")
    print(f"  {label}")
    print(f"  Model: veo-3.1")
    print(f"  Shots: {len(plan_shots)} (sequential with refs)")
    print(f"{'='*60}")

    # Veo goes through _execute_sequential_shots, but we need reference_images
    # injected into each shot's payload. Currently the sequential fallback
    # doesn't support reference_images. We'll submit each shot individually
    # via execute_video with reference_images.
    ref_images = [TORCH_HERO, WREN_FRONT]
    results = []
    for i, (shot, seq_entry) in enumerate(zip(plan_shots, sequence)):
        sid = shot["shot_id"]
        reset_shot_status(store, sid)
        duration = shot.get("_api_duration", 5)
        prompt = seq_entry.get("prompt", "")

        print(f"  Submitting shot {i+1}/{len(plan_shots)}: {sid} ({duration}s)")
        result = runner.execute_video(
            shot_id=sid,
            prompt=prompt,
            model="veo-3.1",
            duration=duration,
            aspect_ratio="9:16",
            reference_images=ref_images,
        )
        results.append(result)
        if result.success:
            print(f"    OK: {result.output_path} (${result.cost_usd:.3f})")
        else:
            print(f"    FAIL: {result.error}")

    ok = sum(1 for r in results if r.success)
    cost = sum(r.cost_usd for r in results)
    print(f"  Total: {ok}/{len(results)} ok, ${cost:.3f}")
    return results


def main():
    paths = ProjectPaths.for_episode(PROJECT, EPISODE)
    store = ExecutionStore(project=PROJECT)
    runner = StepRunner(store=store, paths=paths)
    plan_shots = build_plan_shots()

    # Build elements for O3
    elements_payload, has_loc, total_el = ElementManager.build_elements_with_info(
        ELEMENT_CHAR_IDS, PROJECT, location_id=LOCATION_ID,
    )
    print(f"Elements built: {total_el} total, location={has_loc}")

    prompt_overrides = build_prompt_overrides()

    t0 = time.time()

    # 1. Kling V3 multi-shot (no elements)
    run_kling_multishot(
        runner, store, plan_shots,
        model="kling-o3",
        elements_payload=None,
        label="RENDER 1: Kling V3 multi-shot (no elements)",
        prompt_overrides=prompt_overrides,
    )

    # 2. Kling O3 multi-shot with elements
    run_kling_multishot(
        runner, store, plan_shots,
        model="kling-o3",
        elements_payload=elements_payload,
        label="RENDER 2: Kling O3 multi-shot (WREN + TORCH elements)",
        prompt_overrides=prompt_overrides,
    )

    # 3. Kling V3 continuous single-take (uses COMBINED_PROMPT directly, not prompt_engine)
    run_single_take(
        runner, store,
        shot_id="EP001_SH29",
        model="kling-o3",
        duration=10,
        prompt=COMBINED_PROMPT,
        label="RENDER 3: Kling V3 continuous single-take (10s)",
    )

    # 4-5: Veo SKIPPED — reference_images returns INVALID_ARGUMENT (API tier limitation)
    # and rate limiting. Will investigate Veo ref support separately.
    print("\n  [SKIP] Veo renders — reference_images not supported on current API tier")

    elapsed = time.time() - t0
    print(f"\n{'='*60}")
    print(f"  ALL 5 RENDERS COMPLETE — {elapsed/60:.1f} min total")
    print(f"{'='*60}")


if __name__ == "__main__":
    main()