#!/usr/bin/env python3
"""
build_upload_bundle.py — Package prompts + refs for manual frontier model upload.

Plan-first: loads shot data from Starsend plans when available,
falls back to Recoil storyboards.

For models without API access (Kling web UI, Veo playground), packages everything
needed for manual upload: prompt text, reference images, and step-by-step instructions.

Supports:
- Model-to-prompt routing (compiled_prompts from plan JSON)
- I2V mode (keyframe as start frame + identity refs)
- Character + Prop + Location ref routing per model type
- A/B variant bundles with configurable ref ordering
- Kling O3 multi-reference T2V character ref packaging

Usage:
    python -m tools.build_upload_bundle --episode 1 --shots 1-5 --model kling-v3-direct --mode i2v --project starsend-test
    python -m tools.build_upload_bundle --episode 1 --shots 7 --model kling-o3-direct --variant A --ref-order hero-last --project starsend-test
    python -m tools.build_upload_bundle --episode 1 --shots 1,3,5 --model veo-3.1 --project starsend-test
    python -m tools.build_upload_bundle --episode 1 --shots 1-5 --model kling-v3-direct --dry-run --project starsend-test
"""

import argparse
import json
import logging
import shutil
import sys
from pathlib import Path
from typing import Optional

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

from recoil.pipeline._lib.recoil_bridge import (
    load_storyboard, load_breakdown, load_project_config,
    get_character_refs, resolve_character_for_episode, get_shot_by_id,
)
from recoil.execution.asset_manager import AssetManager, ReferenceImage
from recoil.pipeline._lib.prompt_engine import (
    build_cinematic_prompt, build_two_character_prompt,
    build_prompt_from_plan, _is_plan_shot,
)
from recoil.core.model_profiles import get_profile, get_max_refs, get_aspect_ratios
from recoil.core.paths import PIPELINE_ROOT, projects_root, STATE_NAMESPACE, ProjectPaths
from orchestrator.scene_planner import classify_shot_tier

# ── Constants ────────────────────────────────────────────────────────

BUNDLES_ROOT = None  # Set per-project in build_bundle()

# Model → compiled_prompts key mapping
_PROMPT_KEY_MAP = {
    "kling-v3-direct": "kling_i2v",
    "kling-o3-direct": "kling_i2v",       # fallback; prefer kling_t2v if exists
    "veo-3.1": "veo_t2v",
    "hunyuan-video-1.5": "kling_i2v",  # same motion-description format
}

# Models that support I2V (keyframe as start frame)
_I2V_MODELS = {"kling-v3-direct", "hunyuan-video-1.5"}

# Kling O3 character ref order (recency bias: hero last = strongest anchor)
_KLING_O3_REF_ORDER_HERO_LAST = [
    "kit_angle_grid.png",
    "kit_profile.png",
    "kit_front.png",
    "kit_hero.png",
]

_KLING_O3_REF_ORDER_HERO_FIRST = [
    "kit_hero.png",
    "kit_angle_grid.png",
    "kit_profile.png",
    "kit_front.png",
]

logger = logging.getLogger("starsend.bundle")


def parse_shot_range(shots_str: str) -> list[int]:
    """Parse shot range string: '1-5' or '1,3,5' or '2'."""
    result = []
    for part in shots_str.split(","):
        part = part.strip()
        if "-" in part:
            start, end = part.split("-", 1)
            result.extend(range(int(start), int(end) + 1))
        else:
            result.append(int(part))
    return sorted(set(result))


def _get_plan_shot(plan: dict, shot_id: int) -> Optional[dict]:
    """Find a shot in the plan JSON by shot number."""
    target = f"EP{plan['episode_id'].replace('EP', '')}_{f'SH{shot_id:02d}'}"
    # Also try raw format
    for shot in plan.get("shots", []):
        sid = shot.get("shot_id", "")
        if sid == target or sid.endswith(f"SH{shot_id:02d}"):
            return shot
    return None


def _get_compiled_prompt(plan_shot: dict, model: str) -> Optional[str]:
    """Get the model-appropriate compiled prompt from plan data.

    Tries model-specific key first, then falls back through alternatives.
    """
    prompts = plan_shot.get("compiled_prompts", {})
    if not prompts:
        return None

    # Try model-specific key
    primary_key = _PROMPT_KEY_MAP.get(model)
    if primary_key and primary_key in prompts:
        return prompts[primary_key]

    # Kling O3: prefer kling_t2v, fall back to kling_i2v
    if model == "kling-o3-direct":
        for key in ("kling_t2v", "kling_i2v"):
            if key in prompts:
                return prompts[key]

    # Veo: try veo_t2v
    if model == "veo-3.1" and "veo_t2v" in prompts:
        return prompts["veo_t2v"]

    # Fallback to keyframe_nbp (still-image prompt, better than nothing)
    if "keyframe_nbp" in prompts:
        return prompts["keyframe_nbp"]

    return None


def _find_keyframe(project: str, episode: int, shot_id: int) -> Optional[Path]:
    """Find the approved keyframe for a shot in sequences/ep_NNN/."""
    frames_dir = ProjectPaths.for_project(project).episode_prep_dir(episode)
    if not frames_dir.exists():
        return None

    # Look for keyframe files matching this shot
    patterns = [
        f"shot_{shot_id:03d}_keyframe_take*.png",
        f"shot_{shot_id:03d}_keyframe.png",
    ]
    for pattern in patterns:
        matches = sorted(frames_dir.glob(pattern))
        if matches:
            return matches[-1]  # Latest take
    return None


def _get_character_ref_paths(project: str, char_id: str) -> list[Path]:
    """Get character reference image paths from project refs."""
    subject_root = ProjectPaths.for_project(project).asset_subject_dir("char", char_id.lower())
    if not subject_root.exists():
        return []

    image_exts = {".png", ".jpg", ".jpeg", ".webp"}
    return sorted(
        p for p in subject_root.iterdir()
        if p.is_file() and p.suffix.lower() in image_exts and "candidates" not in str(p)
    )


def _get_kling_o3_character_refs(project: str, char_id: str, ref_order: str = "hero-last") -> list[ReferenceImage]:
    """Build ordered character refs for Kling O3 T2V (4 refs, specific order).

    Args:
        project: Project name.
        char_id: Character ID (e.g. "KIT").
        ref_order: "hero-last" (Variant A, default) or "hero-first" (Variant B).

    Returns:
        List of ReferenceImage in the specified order.
    """
    subject_root = ProjectPaths.for_project(project).asset_subject_dir("char", char_id.lower())
    if not subject_root.exists():
        return []

    order = _KLING_O3_REF_ORDER_HERO_LAST if ref_order == "hero-last" else _KLING_O3_REF_ORDER_HERO_FIRST
    refs = []
    for i, filename in enumerate(order):
        path = subject_root / filename
        if path.exists():
            refs.append(ReferenceImage(
                path=path,
                label=f"[CHARACTER REF {i + 1}: {char_id} - {path.stem}]",
                weight=i + 1,
                ref_type="identity",
            ))
        else:
            logger.warning(f"Missing Kling O3 ref: {path}")

    return refs


def _get_location_refs(location_id: str, project: str = None) -> list[ReferenceImage]:
    """Get location moodboard refs for ENV shots (Veo 3.1 etc.)."""
    subject_root = ProjectPaths.for_project(project).asset_subject_dir("loc", location_id.lower())
    if not subject_root.exists():
        return []

    image_exts = {".png", ".jpg", ".jpeg", ".webp"}
    refs = []
    for i, p in enumerate(sorted(subject_root.iterdir())):
        if p.is_file() and p.suffix.lower() in image_exts:
            refs.append(ReferenceImage(
                path=p,
                label=f"[LOCATION REF: {location_id} - {p.stem}]",
                weight=i + 1,
                ref_type="scene",
            ))
    return refs


def _get_prop_refs(plan_shot: dict, project: str = None) -> list[ReferenceImage]:
    """Get prop reference images for a shot based on plan asset_data."""
    props = plan_shot.get("asset_data", {}).get("props", [])
    if not props:
        return []

    assets = AssetManager()
    all_refs = []
    for prop_entry in props:
        prop_id = prop_entry.get("prop_id", "")
        if prop_id:
            refs = assets.get_prop_ref(prop_id, project=project)
            all_refs.extend(refs)
    return all_refs


def build_bundle(
    episode: int,
    shot_ids: list[int],
    model: str,
    project: str = None,
    dry_run: bool = False,
    mode: str = "auto",
    variant: Optional[str] = None,
    ref_order: str = "hero-last",
) -> Optional[Path]:
    """Build an upload bundle for a set of shots.

    Creates a directory structure with prompt text files, reference images,
    and a README with step-by-step upload instructions.

    Args:
        episode: Episode number.
        shot_ids: Shot IDs to include.
        model: Target model ID (e.g. "kling-v3-direct", "kling-o3-direct", "veo-3.1").
        project: Project name.
        dry_run: If True, show what would be created without writing.
        mode: "i2v" (keyframe as start frame), "t2v" (text-to-video), or "auto" (detect from model).
        variant: Optional variant label (e.g. "A", "B") for A/B testing.
        ref_order: Ref ordering for Kling O3: "hero-last" or "hero-first".

    Returns:
        Path to the bundle directory, or None on dry_run.
    """
    if project is None:
        from recoil.core.paths import DEFAULT_PROJECT
        project = DEFAULT_PROJECT

    model_profile = get_profile(model)
    max_refs = model_profile.get("max_reference_images", 4)
    aspect_ratios = model_profile.get("supported_aspect_ratios", ["9:16"])
    modality = model_profile.get("modality", "video")

    # Auto-detect mode from model profile
    if mode == "auto":
        if model in _I2V_MODELS:
            mode = "i2v"
        else:
            mode = "t2v"

    # Load plan JSON for compiled prompts
    plan_path = ProjectPaths.for_project(project).plans_dir / f"ep_{episode:03d}_plan.json"
    plan = {}
    if plan_path.exists():
        plan = json.loads(plan_path.read_text(encoding="utf-8"))

    # Load storyboard as fallback for shot metadata (may not exist for plan-only projects)
    sb = None
    try:
        sb = load_storyboard(episode, project)
    except FileNotFoundError:
        if plan:
            # Use plan as storyboard-compatible dict
            sb = plan
            sb.setdefault("_source", "plan")
        else:
            raise

    bd = None
    try:
        bd = load_breakdown(project)
    except FileNotFoundError:
        bd = {}

    config = None
    try:
        config = load_project_config(project)
    except FileNotFoundError:
        config = {}

    assets = AssetManager()

    # Bundle directory (project-local)
    bundles_root = ProjectPaths.for_project(project).bundles_dir
    variant_suffix = f"_variant_{variant}" if variant else ""
    mode_suffix = f"_{mode}" if modality == "video" else ""
    bundle_name = f"ep_{episode:03d}_{model.replace('-', '_')}{mode_suffix}{variant_suffix}"
    bundle_dir = bundles_root / bundle_name

    if dry_run:
        print(f"=== Upload Bundle (DRY RUN) ===")
        print(f"Bundle: {bundle_dir}")
        print(f"Model: {model} ({model_profile.get('display_name', model)})")
        print(f"Mode: {mode}")
        print(f"Shots: {len(shot_ids)}")
        print(f"Max refs per shot: {max_refs}")
        if variant:
            print(f"Variant: {variant} (ref order: {ref_order})")
        print()

    if not dry_run:
        bundle_dir.mkdir(parents=True, exist_ok=True)

    bundle = {
        "episode": episode,
        "project": project,
        "model": model,
        "model_display_name": model_profile.get("display_name", model),
        "mode": mode,
        "variant": variant,
        "ref_order": ref_order if model == "kling-o3-direct" else None,
        "shots": [],
    }

    readme_lines = [
        f"# Upload Bundle — EP{episode:03d} for {model_profile.get('display_name', model)}",
        "",
        f"Model: {model}",
        f"Provider: {model_profile.get('provider', 'unknown')}",
        f"Mode: {mode.upper()}",
        f"Max reference images: {max_refs}",
    ]
    if variant:
        readme_lines.append(f"Variant: {variant} (ref order: {ref_order})")
    readme_lines.extend([
        "",
        "## Instructions",
        "",
    ])

    for shot_id in shot_ids:
        shot = get_shot_by_id(sb, shot_id)
        plan_shot = _get_plan_shot(plan, shot_id) if plan else None

        if shot is None and plan_shot is None:
            logger.warning(f"Shot {shot_id} not found, skipping")
            continue

        # Prefer plan_shot for metadata
        working_shot = plan_shot or shot

        # Extract shot metadata
        if _is_plan_shot(working_shot):
            shot_name = working_shot.get("shot_id", f"shot_{shot_id}")
            routing = working_shot.get("routing_data", {})
            asset_data = working_shot.get("asset_data", {})
            is_env = routing.get("is_env_only", False)
            chars = [c.get("char_id", "") for c in asset_data.get("characters", [])]
            location_id = asset_data.get("location_id", "")
            props = asset_data.get("props", [])
        else:
            shot_name = working_shot.get("name", f"shot_{shot_id}")
            chars = working_shot.get("characters_in_shot", [])
            is_env = len(chars) == 0
            location_id = ""
            props = []

        tier = classify_shot_tier(working_shot)

        # ── Build prompt ─────────────────────────────────────
        prompt = None
        prompt_source = "unknown"

        # For video models: JIT compile from live bible + skeleton
        if modality == "video" and plan_shot:
            try:
                from recoil.pipeline._lib.jit_prompt import hydrate_skeleton
                from recoil.pipeline._lib.prompt_engine import build_kling_i2v_prompt, build_kling_t2v_prompt, build_video_prompt_from_plan
                skel = plan_shot.get("prompt_data", {}).get("prompt_skeleton", {})
                if skel and bd:
                    hydrated = hydrate_skeleton(skel, bd, episode=episode, asset_data=plan_shot.get("asset_data"))
                    if model in ("kling-v3-direct", "hunyuan-video-1.5"):
                        prompt = build_kling_i2v_prompt(plan_shot)
                    elif model == "kling-o3-direct":
                        prompt = build_kling_t2v_prompt(plan_shot)
                    elif model == "veo-3.1":
                        prompt = build_video_prompt_from_plan(plan_shot, bd, {}, episode=episode)
                    else:
                        prompt = build_kling_t2v_prompt(plan_shot)
                    if prompt:
                        prompt_source = f"jit_prompt ({model})"
            except ImportError:
                pass
            # Fallback to compiled_prompts if JIT failed
            if prompt is None:
                prompt = _get_compiled_prompt(plan_shot, model)
                if prompt:
                    prompt_source = f"compiled_prompts ({_PROMPT_KEY_MAP.get(model, 'fallback')})"

        # Fallback: build from plan data (image-style prompt)
        if prompt is None:
            if _is_plan_shot(working_shot):
                prompt = build_prompt_from_plan(working_shot, bd, config, episode=episode)
                prompt_source = "build_prompt_from_plan"
            else:
                character_data = {}
                for char_key in chars:
                    try:
                        resolved = resolve_character_for_episode(char_key, episode, project)
                        character_data[char_key.lower()] = {
                            "name": resolved["display_name"],
                            "visual": resolved["visual_description"],
                            "wardrobe": resolved["wardrobe_desc"],
                            "identity_type": "non_human" if "android" in resolved.get(
                                "visual_description", "").lower() else "human",
                        }
                    except (FileNotFoundError, KeyError):
                        pass

                if len(chars) >= 2 and len(character_data) >= 2:
                    char_a_data = character_data.get(chars[0].lower(), {})
                    char_b_data = character_data.get(chars[1].lower(), {})
                    prompt = build_two_character_prompt(
                        working_shot, sb, char_a_data, char_b_data, config)
                else:
                    prompt = build_cinematic_prompt(
                        working_shot, sb, character_data, config, is_env=is_env)
                prompt_source = "legacy_prompt_engine"

        # ── Build reference list ─────────────────────────────
        ref_files: list[ReferenceImage] = []
        notes_lines: list[str] = []

        # --- Kling O3 T2V: 4 ordered character refs ---
        if model == "kling-o3-direct" and chars:
            for char_key in chars[:1]:  # Primary character only for O3
                char_refs = _get_kling_o3_character_refs(project, char_key, ref_order)
                ref_files.extend(char_refs)
            notes_lines.append("Upload as CHARACTER REFERENCE IMAGES in Kling O3")
            notes_lines.append(f"Ref order: {ref_order} (recency bias: last image gets most attention)")

        # --- I2V mode: keyframe as start frame + identity refs ---
        elif mode == "i2v":
            keyframe = _find_keyframe(project, episode, shot_id)
            if keyframe:
                ref_files.append(ReferenceImage(
                    path=keyframe,
                    label=f"[START FRAME / I2V INPUT: keyframe for SH{shot_id:02d}]",
                    weight=0,  # First position (uploaded as start frame, not ref)
                    ref_type="keyframe",
                ))
                notes_lines.append("Upload ref_01_keyframe as START FRAME / I2V input")
            else:
                notes_lines.append(f"WARNING: No keyframe found for SH{shot_id:02d} — generate via Console first")

            # ADR-M02: Kling I2V still needs identity Elements even with start_frame
            if chars and model in ("kling-v3-direct",):
                for char_key in chars[:1]:
                    char_paths = _get_character_ref_paths(project, char_key)
                    identity_refs = assets.get_identity_refs(char_key, char_paths, max_refs=2)
                    ref_files.extend(identity_refs)
                notes_lines.append("Create CHARACTER ELEMENT in Kling web UI using identity refs")

            # Include prop refs for shots with props
            if props:
                prop_refs = _get_prop_refs(plan_shot or working_shot, project=project)
                if prop_refs:
                    ref_files.extend(prop_refs[:2])  # Max 2 prop refs
                    notes_lines.append("Create PROP ELEMENT in Kling web UI using prop refs")

        # --- Veo 3.1 ENV shots: location moodboards ---
        elif model == "veo-3.1" and is_env and location_id:
            loc_refs = _get_location_refs(location_id, project=project)
            ref_files.extend(loc_refs[:max_refs])
            notes_lines.append("Upload location moodboards as reference images in Veo")

        # --- Veo 3.1 non-ENV: location + limited character refs ---
        elif model == "veo-3.1" and not is_env:
            # Veo has 3-ref limit — use 2 location + 1 character or skip location
            if location_id:
                loc_refs = _get_location_refs(location_id, project=project)
                ref_files.extend(loc_refs[:1])  # 1 location ref
            for char_key in chars[:1]:
                char_paths = _get_character_ref_paths(project, char_key)
                identity_refs = assets.get_identity_refs(char_key, char_paths, max_refs=2)
                ref_files.extend(identity_refs)
            notes_lines.append("Veo 3-ref limit: prioritized character identity over location")

        # --- Default: character identity refs ---
        elif not is_env:
            for char_key in chars[:2]:
                char_paths = _get_character_ref_paths(project, char_key)
                per_char = max(1, max_refs // max(len(chars), 1))
                identity_refs = assets.get_identity_refs(char_key, char_paths, max_refs=per_char)
                ref_files.extend(identity_refs)

        # Trim to model's max refs (but don't trim keyframe which is a start frame, not a ref)
        actual_refs = [r for r in ref_files if r.ref_type != "keyframe"]
        keyframe_refs = [r for r in ref_files if r.ref_type == "keyframe"]
        actual_refs = actual_refs[:max_refs]
        ref_files = keyframe_refs + actual_refs

        shot_entry = {
            "shot_id": shot_id,
            "shot_name": shot_name,
            "tier": tier,
            "is_env": is_env,
            "characters": chars,
            "props": [p.get("prop_id", "") for p in props],
            "prompt_length": len(prompt) if prompt else 0,
            "prompt_source": prompt_source,
            "ref_count": len(ref_files),
            "mode": mode,
        }
        bundle["shots"].append(shot_entry)

        if dry_run:
            print(f"  Shot {shot_id}: {shot_name} ({tier})")
            print(f"    Characters: {chars or 'ENV'}")
            print(f"    Props: {[p.get('prop_id', '') for p in props] or 'none'}")
            print(f"    Mode: {mode}")
            print(f"    Prompt source: {prompt_source}")
            print(f"    Refs: {len(ref_files)}")
            for r in ref_files:
                exists = r.path.exists()
                print(f"      {'[OK]' if exists else '[MISSING]'} {r.label} ({r.path.name})")
            if prompt:
                print(f"    Prompt: {prompt[:120]}...")
            print()
            continue

        # Create shot directory
        shot_dir = bundle_dir / f"shot_{shot_id:02d}_{shot_name}"
        shot_dir.mkdir(exist_ok=True)
        refs_dir = shot_dir / "refs"
        refs_dir.mkdir(exist_ok=True)

        # Write prompt
        if prompt:
            (shot_dir / "prompt.txt").write_text(prompt, encoding="utf-8")

        # Copy reference images
        for i, ref in enumerate(ref_files):
            if ref.path.exists():
                ext = ref.path.suffix
                # Use descriptive naming for refs
                if ref.ref_type == "keyframe":
                    dest = refs_dir / f"ref_{i + 1:02d}_keyframe{ext}"
                elif ref.ref_type == "identity":
                    dest = refs_dir / f"ref_{i + 1:02d}_identity_{ref.path.stem}{ext}"
                elif ref.ref_type == "prop":
                    dest = refs_dir / f"ref_{i + 1:02d}_prop_{ref.path.stem}{ext}"
                elif ref.ref_type == "scene":
                    dest = refs_dir / f"ref_{i + 1:02d}_location_{ref.path.stem}{ext}"
                else:
                    dest = refs_dir / f"ref_{i + 1:02d}_{ref.path.stem}{ext}"
                shutil.copy2(ref.path, dest)
                shot_entry[f"ref_{i + 1}"] = ref.label

        # Write shot notes
        shot_type = ""
        camera_info = ""
        action_info = ""
        emotion_info = ""
        if plan_shot:
            pd = plan_shot.get("prompt_data", {})
            shot_type = pd.get("shot_type", "")
            camera_info = f"{pd.get('camera_movement', 'static')}, {pd.get('focal_length', '')}"
            skeleton = pd.get("prompt_skeleton", {})
            action_info = skeleton.get("action_line", "")
            emotion_info = skeleton.get("emotion_line", "")
        elif shot:
            shot_type = shot.get("shot_type", "MS")
            camera_info = f"{shot.get('camera_angle', 'eye')} angle, {shot.get('focal_length', '50mm')}"
            action_info = shot.get("action", "")
            emotion_info = shot.get("emotion", "")

        notes = [
            f"Shot {shot_id}: {shot_name}",
            f"Type: {shot_type}",
            f"Camera: {camera_info}",
            f"Tier: {tier}",
            f"Mode: {mode.upper()}",
            "",
            f"Action: {action_info}",
            f"Emotion: {emotion_info}",
            "",
        ]

        if notes_lines:
            notes.append("--- Model-specific instructions ---")
            notes.extend(notes_lines)
            notes.append("")

        notes.append("Reference upload order:")
        for i, ref in enumerate(ref_files):
            notes.append(f"  {i + 1}. [{ref.ref_type.upper()}] {ref.label} — {ref.path.name}")

        (shot_dir / "notes.txt").write_text("\n".join(notes), encoding="utf-8")

        # README per-shot instructions
        readme_lines.append(f"### Shot {shot_id}: {shot_name}")
        readme_lines.append("")
        readme_lines.append(f"1. Open `shot_{shot_id:02d}_{shot_name}/prompt.txt`")

        if mode == "i2v":
            readme_lines.append(f"2. Upload `refs/ref_01_keyframe.png` as **Start Frame / I2V input**")
            readme_lines.append(f"3. Upload remaining refs as reference images / create Elements")
        elif model == "kling-o3-direct":
            readme_lines.append(f"2. Upload character refs from `refs/` in order listed in `notes.txt`")
            readme_lines.append(f"3. Ref order: {ref_order} (last image = strongest identity anchor)")
        else:
            readme_lines.append(f"2. Upload refs from `refs/` in the order listed in `notes.txt`")

        readme_lines.append(f"4. Set aspect ratio to `9:16` (vertical)")

        if model_profile.get("supports_negative_prompt"):
            readme_lines.append(f"5. Negative prompt: `deformed, blurry, cartoon, anime, illustration`")

        readme_lines.append("")

    if not dry_run:
        # Write bundle manifest
        (bundle_dir / "bundle.json").write_text(
            json.dumps(bundle, indent=2), encoding="utf-8"
        )

        # Write README
        (bundle_dir / "README.md").write_text("\n".join(readme_lines), encoding="utf-8")

        logger.info(f"Bundle created: {bundle_dir}")
        return bundle_dir

    return None


def main():
    parser = argparse.ArgumentParser(description="Build upload bundle for manual frontier model use")
    parser.add_argument("--episode", type=int, required=True, help="Episode number")
    parser.add_argument("--shots", required=True, help="Shot range (e.g. '1-5' or '1,3,5')")
    parser.add_argument("--model", default="kling-v3-direct", help="Target model (default: kling-v3-direct)")
    parser.add_argument("--mode", default="auto", choices=["auto", "i2v", "t2v"],
                        help="I2V (keyframe as start frame) or T2V (text-to-video). Default: auto-detect from model.")
    parser.add_argument("--variant", default=None, help="Variant label for A/B testing (e.g. 'A', 'B')")
    parser.add_argument("--ref-order", default="hero-last", choices=["hero-last", "hero-first"],
                        help="Ref ordering for Kling O3 (default: hero-last)")
    parser.add_argument("--dry-run", action="store_true", help="Show what would be created")
    parser.add_argument("--project", default=None)
    args = parser.parse_args()

    logging.basicConfig(level=logging.INFO, format="%(asctime)s %(message)s",
                        datefmt="%H:%M:%S")

    shot_ids = parse_shot_range(args.shots)
    print(f"Building upload bundle for EP{args.episode:03d}, shots {shot_ids}, model {args.model}, mode {args.mode}")

    result = build_bundle(
        episode=args.episode,
        shot_ids=shot_ids,
        model=args.model,
        project=args.project,
        dry_run=args.dry_run,
        mode=args.mode,
        variant=args.variant,
        ref_order=args.ref_order,
    )

    if result:
        print(f"\nBundle ready at: {result}")
        print(f"Contents:")
        for f in sorted(result.rglob("*")):
            if f.is_file():
                size = f.stat().st_size
                rel = f.relative_to(result)
                print(f"  {rel} ({size:,} bytes)")


if __name__ == "__main__":
    main()