#!/usr/bin/env python3
"""
probe_i2v_dual_frame.py — I2V with start + end frame + audio.

Routes through StepRunner.execute_video() with start_frame, end_frame,
and generate_audio. Works with any model that supports all three
(kling-v3, seeddance-2.0, kling-o3, etc.).

Usage:
    python3 recoil/pipeline/tools/probe_i2v_dual_frame.py \\
      --project driver-beware \\
      --model kling-v3 \\
      --start-frame /path/to/start.png \\
      --end-frame /path/to/end.png \\
      --prompt "..." \\
      --duration 6 --audio \\
      [--shot-id NAME] [--aspect 16:9] [--negative-prompt "..."]
"""

import argparse
import os
import sys
import time
from pathlib import Path

_HERE = os.path.dirname(os.path.abspath(__file__))
_RECOIL_ROOT = os.path.dirname(os.path.dirname(_HERE))
_PROJECTS_ROOT_PARENT = os.path.dirname(_RECOIL_ROOT)
if _RECOIL_ROOT not in sys.path:
    sys.path.insert(0, _RECOIL_ROOT)
if _PROJECTS_ROOT_PARENT not in sys.path:
    sys.path.insert(0, _PROJECTS_ROOT_PARENT)

from recoil.core.paths import ensure_pipeline_importable  # noqa: E402

ensure_pipeline_importable()


def main() -> int:
    parser = argparse.ArgumentParser(
        description="Probe I2V with start + end frame + audio (Kling V3 / Seedance / Kling O3)."
    )
    parser.add_argument("--project", required=True)
    parser.add_argument("--model", required=True, help="kling-v3 / seeddance-2.0 / kling-o3")
    parser.add_argument("--start-frame", required=True)
    parser.add_argument("--end-frame", default=None, help="Optional. Omit for single-frame I2V.")
    parser.add_argument("--prompt", required=True)
    parser.add_argument("--duration", type=int, default=6)
    parser.add_argument("--aspect", default="16:9")
    parser.add_argument("--audio", action="store_true", help="Enable generate_audio")
    parser.add_argument("--negative-prompt", default=None)
    parser.add_argument("--shot-id", default=None)
    parser.add_argument("--dry-run", action="store_true")
    args = parser.parse_args()

    start_path = Path(args.start_frame).expanduser().resolve()
    if not start_path.exists():
        print(f"ERROR: --start-frame not found: {start_path}")
        return 1
    end_path = None
    if args.end_frame:
        end_path = Path(args.end_frame).expanduser().resolve()
        if not end_path.exists():
            print(f"ERROR: --end-frame not found: {end_path}")
            return 1

    model_tag = args.model.upper().replace("-", "_").replace(".", "")
    shot_id = args.shot_id or f"I2V_DUAL_{model_tag}_{int(time.time())}"

    print()
    print("Mode:        I2V dual-frame (probe)")
    print(f"Project:     {args.project}")
    print(f"Shot ID:     {shot_id}")
    print(f"Model:       {args.model}")
    print(f"Aspect:      {args.aspect}")
    print(f"Duration:    {args.duration}s")
    print(f"Start:       {start_path.name}")
    print(f"End:         {end_path.name if end_path else '(none — single-frame I2V)'}")
    print(f"Audio:       {'ON' if args.audio else 'OFF'}")
    if args.negative_prompt:
        print(f"Negative:    {args.negative_prompt}")
    print(f"Prompt ({len(args.prompt.split())} words):")
    print("  " + args.prompt.replace("\n", "\n  "))

    if args.dry_run:
        print()
        print("=== DRY RUN — not submitting ===")
        return 0

    from recoil.execution.execution_store import ExecutionStore  # noqa: E402
    from recoil.execution.step_runner import StepRunner  # noqa: E402
    from recoil.execution.step_types import ProjectPaths  # noqa: E402

    store = ExecutionStore(args.project)
    paths = ProjectPaths.for_episode(args.project, 1)
    paths.video_dir.mkdir(parents=True, exist_ok=True)

    runner = StepRunner(store=store, paths=paths, validate_frames=False)
    runner._dispatch_path = "probe_i2v_dual_frame.py"

    t0 = time.time()
    result = runner.execute_video(
        shot_id=shot_id,
        prompt=args.prompt,
        model=args.model,
        start_frame=start_path,
        end_frame=end_path,
        duration=args.duration,
        aspect_ratio=args.aspect,
        generate_audio=args.audio,
        negative_prompt=args.negative_prompt,
    )
    elapsed = time.time() - t0

    status = "OK" if result.success else "FAIL"
    print()
    print(f"[{status}] {shot_id} -> {result.output_path} (${result.cost_usd or 0:.2f})")
    if result.error:
        print(f"Error: {result.error}")
    print(f"Done in {elapsed:.0f}s")
    return 0 if result.success else 1


if __name__ == "__main__":
    sys.exit(main())