#!/usr/bin/env python3
"""CLI — Generate video from text prompt (T2V) or edit existing video (V2V).

Usage:
    python -m cli.generate "A dragon breathing fire"
    python -m cli.generate "A dragon breathing fire" --aspect landscape -o dragon.mp4
    python -m cli.generate "Make it anime" --edit MEDIA_ID
"""

import argparse
import asyncio
import os
import sys

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from omniflash import (
    ExtensionBridge, generate_video, edit_video,
    poll_status, download_video, ASPECTS, DEFAULT_PROJECT,
)
from omniflash.generators.i2v import upload_image, generate_video_i2v, generate_video_fl, generate_video_r2v


async def run(args):
    aspect = ASPECTS.get(args.aspect, "VIDEO_ASPECT_RATIO_PORTRAIT")

    bridge = ExtensionBridge()
    await bridge.start()

    if not await bridge.wait_for_extension(timeout=30):
        return

    # Auto-upload local image files
    async def resolve_image(path_or_id):
        if os.path.exists(path_or_id):
            mid = await upload_image(bridge, path_or_id)
            if mid:
                print(f"📤 Uploaded: {path_or_id} → {mid[:12]}...")
            return mid
        return path_or_id

    if args.start and args.end:
        # First+Last frame mode
        start_id = await resolve_image(args.start)
        end_id = await resolve_image(args.end)
        if not start_id or not end_id:
            await bridge.close()
            return
        media_ids = await generate_video_fl(bridge, args.prompt, aspect, args.project_id,
                                             start_image_id=start_id, end_image_id=end_id,
                                             duration=args.duration)
    elif args.start:
        # I2V mode (start image only)
        start_id = await resolve_image(args.start)
        if not start_id:
            await bridge.close()
            return
        media_ids = await generate_video_i2v(bridge, args.prompt, aspect, args.project_id,
                                              image_media_id=start_id, duration=args.duration)
    elif args.ref:
        # Reference images mode
        ref_ids = []
        for r in args.ref:
            mid = await resolve_image(r)
            if mid:
                ref_ids.append(mid)
        if not ref_ids:
            await bridge.close()
            return
        media_ids = await generate_video_r2v(bridge, args.prompt, aspect, args.project_id,
                                              ref_media_ids=ref_ids, duration=args.duration)
    elif args.edit:
        media_ids = await edit_video(bridge, args.prompt, aspect, args.project_id,
                                     video_media_id=args.edit, duration=args.duration)
    else:
        media_ids = await generate_video(bridge, args.prompt, aspect, args.project_id,
                                         duration=args.duration, count=args.count)

    if not media_ids:
        await bridge.close()
        return

    for i, media_id in enumerate(media_ids):
        label = f"[{i+1}/{len(media_ids)}] " if len(media_ids) > 1 else ""
        print(f"{label}Polling {media_id[:12]}...")
        if not await poll_status(bridge, media_id, args.project_id):
            continue

        if len(media_ids) == 1:
            out_path = args.output
        else:
            base, ext = os.path.splitext(args.output)
            out_path = f"{base}_{i+1}{ext}"

        # Setup temp dir for download
        out_dir = os.path.dirname(out_path) or "."
        temp_dir = os.path.join(out_dir, ".temp")
        os.makedirs(temp_dir, exist_ok=True)
        temp_path = os.path.join(temp_dir, os.path.basename(out_path))

        if await download_video(bridge, media_id, temp_path):
            # Watermark removal disabled — save the raw downloaded video as-is.
            os.replace(temp_path, out_path)

            # Cleanup empty .temp dir
            try:
                os.rmdir(temp_dir)
            except OSError:
                pass

            print(f"🎉 Done! {out_path}")

    await bridge.close()


def main():
    parser = argparse.ArgumentParser(description="Omni Flash — Video Generator")
    parser.add_argument("prompt", help="Text prompt for video")
    parser.add_argument("--output", "-o", default="omni_output.mp4", help="Output file")
    parser.add_argument("--aspect", "-a", choices=["portrait", "landscape"], default="portrait")
    parser.add_argument("--duration", "-d", type=int, choices=[4, 6, 8, 10], default=10)
    parser.add_argument("--count", "-c", type=int, choices=[1, 2, 3, 4], default=1)
    parser.add_argument("--edit", "-e", metavar="MEDIA_ID",
                        help="Edit existing video (V2V mode)")
    parser.add_argument("--start", "-s", metavar="IMAGE",
                        help="Start frame image (file path or media_id)")
    parser.add_argument("--end", metavar="IMAGE",
                        help="End frame image (use with --start for FL mode)")
    parser.add_argument("--ref", "-r", nargs="+", metavar="IMAGE",
                        help="Reference images for R2V mode")
    parser.add_argument("--project-id", "-p", default=DEFAULT_PROJECT)
    parser.add_argument("--no-clean", action="store_true",
                        help="Skip automatic watermark removal")
    args = parser.parse_args()
    asyncio.run(run(args))


if __name__ == "__main__":
    main()