main.py

#!/usr/bin/env python3
"""Example CLI for submitting and managing simulation jobs using YAML config.

Simulations generate video asynchronously without requiring a live connection.
Define your script in a YAML file, submit it, and download the result.

Setup:
    # Create .env.local with your API key (or export ODYSSEY_API_KEY)
    echo 'ODYSSEY_API_KEY=ody_your_key_here' > .env.local

Usage:
    # Submit a simulation from YAML config
    uv run python examples/simulate/main.py submit config.yaml

    # Check job status
    uv run python examples/simulate/main.py status <job_id>

    # List your simulation jobs
    uv run python examples/simulate/main.py list

    # Wait for completion and download
    uv run python examples/simulate/main.py wait <job_id> --output video.mp4

    # Cancel a pending job
    uv run python examples/simulate/main.py cancel <job_id>

Example YAML config (config.yaml):
    prompt: "A cat sleeping on a couch"
    portrait: true
    duration: 10s

    interactions:
      - time: 3s
        prompt: "The cat wakes up"
      - time: 6s
        prompt: "The cat stretches"

    # Optional: image for i2v
    # image: /path/to/image.jpg
"""

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

import aiohttp
import yaml
from dotenv import load_dotenv

# Load .env.local from example directory or current working directory
load_dotenv(Path(__file__).parent / ".env.local")
load_dotenv(Path.cwd() / ".env.local")

from odyssey import Odyssey, SimulationJobDetail, SimulationJobStatus  # noqa: E402


def format_duration(seconds: float | None) -> str:
    """Format duration in human readable form."""
    if seconds is None:
        return "N/A"
    if seconds < 60:
        return f"{seconds:.1f}s"
    minutes = int(seconds // 60)
    secs = seconds % 60
    return f"{minutes}m {secs:.0f}s"


def format_timestamp(ts: str | None) -> str:
    """Format ISO timestamp to readable form."""
    if not ts:
        return "N/A"
    return ts[:19].replace("T", " ")


def parse_time(time_str: str | int | float) -> int:
    """Parse time string like '3s', '1.5m', or raw ms int to milliseconds."""
    if isinstance(time_str, int):
        return time_str
    if isinstance(time_str, float):
        return int(time_str)

    time_str = str(time_str).strip().lower()

    if time_str.endswith("ms"):
        return int(float(time_str[:-2]))
    elif time_str.endswith("s"):
        return int(float(time_str[:-1]) * 1000)
    elif time_str.endswith("m"):
        return int(float(time_str[:-1]) * 60 * 1000)
    else:
        # Assume seconds if no unit
        return int(float(time_str) * 1000)


def status_color(status: SimulationJobStatus) -> str:
    """Return ANSI color code for status."""
    colors = {
        SimulationJobStatus.PENDING: "\033[33m",  # Yellow
        SimulationJobStatus.DISPATCHED: "\033[36m",  # Cyan
        SimulationJobStatus.PROCESSING: "\033[34m",  # Blue
        SimulationJobStatus.COMPLETED: "\033[32m",  # Green
        SimulationJobStatus.FAILED: "\033[31m",  # Red
        SimulationJobStatus.CANCELLED: "\033[90m",  # Gray
    }
    return colors.get(status, "")


def reset_color() -> str:
    """Return ANSI reset code."""
    return "\033[0m"


def load_image_as_format(image_path: str, fmt: str) -> tuple:
    """Load image in the specified format for testing different input types.

    Args:
        image_path: Path to the image file.
        fmt: One of 'path', 'pil', 'bytes', 'numpy'.

    Returns:
        Tuple of (image_data, format_description).
    """
    import cv2

    if fmt == "path":
        return image_path, "file path"

    elif fmt == "pil":
        from PIL import Image

        img = Image.open(image_path)
        return img, "PIL Image"

    elif fmt == "bytes":
        with open(image_path, "rb") as f:
            data = f.read()
        return data, f"raw bytes ({len(data):,} bytes)"

    elif fmt == "numpy":
        # Load with OpenCV and convert BGR -> RGB
        img = cv2.imread(image_path)
        if img is None:
            raise ValueError(f"Failed to load image: {image_path}")
        img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        return img_rgb, f"numpy array {img_rgb.shape}"

    else:
        raise ValueError(f"Unknown format: {fmt}")


def load_config(config_path: str) -> dict:
    """Load and validate YAML config file."""
    path = Path(config_path)
    if not path.exists():
        print(f"Error: Config file not found: {config_path}")
        sys.exit(1)

    with open(path) as f:
        config = yaml.safe_load(f)

    if not config:
        print("Error: Empty config file")
        sys.exit(1)

    # Validate required fields
    if "prompt" not in config:
        print("Error: Config must have 'prompt' field")
        sys.exit(1)

    return config


def build_script_from_config(config: dict, image_format: str = "path") -> tuple[list[dict], str | None]:
    """Build simulation script from YAML config.

    Returns:
        Tuple of (script, image_format_description).
    """
    script: list[dict] = []
    format_desc = None

    # Get duration (default 10s)
    duration_str = config.get("duration", "10s")
    duration_ms = parse_time(duration_str)

    # Start entry
    start_entry: dict = {
        "timestamp_ms": 0,
        "start": {"prompt": config["prompt"]},
    }

    # Handle image - load in specified format
    if "image" in config:
        image_path = config["image"]
        # Resolve relative paths from config file location
        if not os.path.isabs(image_path):
            config_dir = config.get("_config_dir", ".")
            image_path = os.path.join(config_dir, image_path)
        image_data, format_desc = load_image_as_format(image_path, image_format)
        start_entry["start"]["image"] = image_data

    script.append(start_entry)

    # Interaction entries
    interactions = config.get("interactions", [])
    for interaction in interactions:
        if isinstance(interaction, str):
            # Simple string - auto-space it
            continue  # Handle below

        time_ms = parse_time(interaction.get("time", 0))
        prompt = interaction.get("prompt", "")

        if prompt:
            script.append(
                {
                    "timestamp_ms": time_ms,
                    "interact": {"prompt": prompt},
                }
            )

    # Handle simple interaction list (auto-space them)
    if interactions and all(isinstance(i, str) for i in interactions):
        interval = duration_ms // (len(interactions) + 1)
        for idx, prompt in enumerate(interactions):
            script.append(
                {
                    "timestamp_ms": interval * (idx + 1),
                    "interact": {"prompt": prompt},
                }
            )

    # End entry
    script.append(
        {
            "timestamp_ms": duration_ms,
            "end": {},
        }
    )

    # Sort by timestamp
    script.sort(key=lambda x: x["timestamp_ms"])

    return script, format_desc


async def submit_job(client: Odyssey, config_path: str, image_format: str = "path") -> None:
    """Submit a simulation job from YAML config."""
    config = load_config(config_path)

    # Store config directory for relative path resolution
    config["_config_dir"] = str(Path(config_path).parent)

    # Build script
    script, format_desc = build_script_from_config(config, image_format)
    portrait = config.get("portrait", True)

    # Display what we're submitting
    print("Submitting simulation job...")
    print(f"  Config: {config_path}")
    print(f"  Prompt: {config['prompt']}")
    if "image" in config:
        print(f"  Image: {config['image']}")
        print(f"  Image format: {format_desc}")
    print(f"  Orientation: {'portrait' if portrait else 'landscape'}")
    print(f"  Duration: {parse_time(config.get('duration', '10s')) / 1000}s")

    interactions = [e for e in script if "interact" in e]
    if interactions:
        print(f"  Interactions: {len(interactions)}")
        for entry in interactions:
            time_s = entry["timestamp_ms"] / 1000
            print(f"    [{time_s:.1f}s] {entry['interact']['prompt']}")
    print()

    try:
        job = await client.simulate(script=script, portrait=portrait)
        print("Job submitted successfully!")
        print(f"  Job ID: {job.job_id}")
        print(f"  Status: {status_color(job.status)}{job.status.value}{reset_color()}")
        if job.estimated_wait_minutes:
            print(f"  Estimated wait: ~{job.estimated_wait_minutes:.0f} minutes")
        print()
        print("Check status with:")
        print(f"  uv run python examples/simulate/main.py status {job.job_id}")
        print()
        print("Wait for completion and download:")
        print(f"  uv run python examples/simulate/main.py wait {job.job_id} --output video.mp4")
    except Exception as e:
        print(f"Error submitting job: {e}")
        sys.exit(1)


async def get_status(client: Odyssey, job_id: str) -> SimulationJobDetail:
    """Get and display job status."""
    print(f"Fetching status for job {job_id}...")
    job = await client.get_simulate_status(job_id)

    print(f"\nJob: {job.job_id}")
    print("-" * 60)
    print(f"Status: {status_color(job.status)}{job.status.value}{reset_color()}")
    print(f"Priority: {job.priority}")
    print(f"Created: {format_timestamp(job.created_at)}")

    if job.dispatched_at:
        print(f"Dispatched: {format_timestamp(job.dispatched_at)}")
    if job.started_at:
        print(f"Started: {format_timestamp(job.started_at)}")
    if job.completed_at:
        print(f"Completed: {format_timestamp(job.completed_at)}")
    if job.assigned_region:
        print(f"Region: {job.assigned_region}")
    if job.retry_count > 0:
        print(f"Retries: {job.retry_count}")
    if job.error_message:
        print(f"Error: {job.error_message}")

    if job.streams:
        print(f"\nOutput Streams ({len(job.streams)}):")
        for stream in job.streams:
            print(f"  Stream {stream.script_index}: {stream.stream_id}")
            if stream.duration_seconds:
                print(f"    Duration: {format_duration(stream.duration_seconds)}")
            if stream.frame_count:
                print(f"    Frames: {stream.frame_count}")
            if stream.video_url:
                print(f"    Video: {stream.video_url[:70]}...")
            if stream.thumbnail_url:
                print(f"    Thumbnail: {stream.thumbnail_url[:70]}...")

    return job


async def list_jobs(
    client: Odyssey,
    status: SimulationJobStatus | None,
    active: bool,
    limit: int | None,
) -> None:
    """List simulation jobs."""
    print("Fetching simulation jobs...")
    result = await client.list_simulations(
        status=status,
        active=active,
        limit=limit,
    )

    if not result.jobs:
        print("No simulation jobs found.")
        return

    print(f"\nFound {result.total} jobs (showing {len(result.jobs)}):\n")
    print(f"{'Job ID':<38} {'Status':<12} {'Created':<20} {'Error'}")
    print("-" * 90)

    for job in result.jobs:
        status_str = f"{status_color(job.status)}{job.status.value:<12}{reset_color()}"
        created = format_timestamp(job.created_at)
        error = (job.error_message or "")[:30]
        print(f"{job.job_id:<38} {status_str} {created:<20} {error}")


async def wait_for_job(
    client: Odyssey,
    job_id: str,
    output: str | None,
    poll_interval: int,
    timeout: int,
) -> None:
    """Wait for job to complete and optionally download."""
    print(f"Waiting for job {job_id} to complete...")
    print(f"  Poll interval: {poll_interval}s, Timeout: {timeout}s")
    print()

    start_time = time.time()
    last_status = None

    while True:
        elapsed = time.time() - start_time
        if elapsed > timeout:
            print(f"\nTimeout after {timeout}s. Job may still be processing.")
            print(f"Check status with: uv run python examples/simulate/main.py status {job_id}")
            sys.exit(1)

        job = await client.get_simulate_status(job_id)

        # Print status updates
        if job.status != last_status:
            print(f"[{elapsed:.0f}s] Status: {status_color(job.status)}{job.status.value}{reset_color()}")
            last_status = job.status

        # Check terminal states
        if job.status == SimulationJobStatus.COMPLETED:
            print(f"\nJob completed in {elapsed:.1f}s!")

            if job.streams:
                stream = job.streams[0]
                print(f"  Duration: {format_duration(stream.duration_seconds)}")
                print(f"  Frames: {stream.frame_count}")

                if output and stream.video_url:
                    await download_video(stream.video_url, output)
                elif stream.video_url:
                    print("\nVideo URL (valid for ~1 hour):")
                    print(f"  {stream.video_url}")
            return

        if job.status == SimulationJobStatus.FAILED:
            print(f"\nJob failed: {job.error_message}")
            sys.exit(1)

        if job.status == SimulationJobStatus.CANCELLED:
            print("\nJob was cancelled.")
            sys.exit(1)

        # Wait before next poll
        await asyncio.sleep(poll_interval)


async def download_video(url: str, output: str) -> None:
    """Download video from URL."""
    output_path = Path(output)
    print(f"\nDownloading to {output_path}...")

    async with (
        aiohttp.ClientSession() as session,
        session.get(url) as response,
    ):
        if not response.ok:
            print(f"Error downloading: {response.status} {response.reason}")
            sys.exit(1)

        total_size = int(response.headers.get("content-length", 0))
        downloaded = 0

        with open(output_path, "wb") as f:
            async for chunk in response.content.iter_chunked(8192):
                f.write(chunk)
                downloaded += len(chunk)
                if total_size:
                    pct = (downloaded / total_size) * 100
                    print(f"\rDownloading: {pct:.1f}% ({downloaded:,} / {total_size:,} bytes)", end="")

    print(f"\nSaved to {output_path}")


async def cancel_job(client: Odyssey, job_id: str) -> None:
    """Cancel a simulation job."""
    print(f"Cancelling job {job_id}...")
    try:
        result = await client.cancel_simulation(job_id)
        print("Job cancelled successfully.")
        print(f"  Status: {status_color(result.status)}{result.status.value}{reset_color()}")
    except ValueError as e:
        print(f"Error: {e}")
        sys.exit(1)


async def main() -> None:
    """Main entry point."""
    parser = argparse.ArgumentParser(
        description="Odyssey Simulation CLI - YAML-based job submission",
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    parser.add_argument(
        "--api-key",
        default=os.environ.get("ODYSSEY_API_KEY", ""),
        help="Odyssey API key (or set ODYSSEY_API_KEY env var)",
    )

    subparsers = parser.add_subparsers(dest="command", help="Command to run")

    # Submit command
    submit_parser = subparsers.add_parser("submit", help="Submit a simulation job from YAML config")
    submit_parser.add_argument("config", help="Path to YAML config file")
    submit_parser.add_argument(
        "--image-format",
        choices=["path", "pil", "bytes", "numpy"],
        default="path",
        help="How to load the image before passing to API (default: path)",
    )

    # Status command
    status_parser = subparsers.add_parser("status", help="Get job status")
    status_parser.add_argument("job_id", help="Job ID to check")

    # List command
    list_parser = subparsers.add_parser("list", help="List simulation jobs")
    list_parser.add_argument("--active", action="store_true", help="Show only active jobs")
    list_parser.add_argument("--limit", type=int, default=20, help="Maximum jobs to return")

    # Wait command
    wait_parser = subparsers.add_parser("wait", help="Wait for job completion")
    wait_parser.add_argument("job_id", help="Job ID to wait for")
    wait_parser.add_argument("--output", "-o", help="Download video to this file when complete")
    wait_parser.add_argument("--poll", type=int, default=5, help="Poll interval in seconds (default: 5)")
    wait_parser.add_argument("--timeout", type=int, default=600, help="Timeout in seconds (default: 600)")

    # Cancel command
    cancel_parser = subparsers.add_parser("cancel", help="Cancel a job")
    cancel_parser.add_argument("job_id", help="Job ID to cancel")

    args = parser.parse_args()

    if not args.command:
        parser.print_help()
        print("\n" + "=" * 60)
        print("Example YAML config file:")
        print("=" * 60)
        print("""
# simulation.yaml - Text-to-video example
prompt: "A cat sleeping on a couch"
portrait: true
duration: 10s

interactions:
  - time: 3s
    prompt: "The cat wakes up and looks around"
  - time: 6s
    prompt: "The cat stretches and yawns"

# --- OR with image (i2v) ---

# prompt: "Robot starts dancing"
# image: ./robot.jpg
# portrait: false
# duration: 8s
""")
        sys.exit(1)

    if not args.api_key:
        print("Error: API key required. Set ODYSSEY_API_KEY or use --api-key")
        sys.exit(1)

    client = Odyssey(api_key=args.api_key)

    try:
        if args.command == "submit":
            await submit_job(client, args.config, args.image_format)
        elif args.command == "status":
            await get_status(client, args.job_id)
        elif args.command == "list":
            await list_jobs(client, None, args.active, args.limit)
        elif args.command == "wait":
            await wait_for_job(client, args.job_id, args.output, args.poll, args.timeout)
        elif args.command == "cancel":
            await cancel_job(client, args.job_id)
    finally:
        await client.disconnect()


if __name__ == "__main__":
    try:
        asyncio.run(main())
    except KeyboardInterrupt:
        print("\nCancelled.")