#!/usr/bin/env python3
"""
nano-to-grok-video: Animate a still image (or prompt) into a video via Grok Imagine API.
Usage:
  uv run generate_video.py --prompt "..." --image path/to/image.png
  uv run generate_video.py --prompt "..." --image https://example.com/img.png
  uv run generate_video.py --prompt "..." --text-to-video
"""

import argparse
import base64
import mimetypes
import os
import sys
import time
import requests
from datetime import datetime
from pathlib import Path

API_BASE = "https://api.x.ai/v1"
MODEL = "grok-imagine-video"
POLL_INTERVAL = 8   # seconds between status checks
POLL_TIMEOUT = 600  # 10 minutes max


def get_api_key() -> str:
    key = os.getenv("XAI_API_KEY")
    if not key:
        sys.exit("Error: XAI_API_KEY environment variable not set.")
    return key


def encode_image(path: str) -> str:
    """Encode a local image as a base64 data URI."""
    p = Path(path)
    mime, _ = mimetypes.guess_type(str(p))
    mime = mime or "image/png"
    with open(p, "rb") as f:
        data = base64.b64encode(f.read()).decode()
    return f"data:{mime};base64,{data}"


def prepare_image(image_arg: str) -> str:
    """Return URL or base64 data URI for image input."""
    if image_arg.startswith("http://") or image_arg.startswith("https://"):
        return image_arg
    return encode_image(image_arg)


def start_generation(api_key: str, payload: dict) -> str:
    """Submit video generation request; return request_id."""
    resp = requests.post(
        f"{API_BASE}/videos/generations",
        headers={"Authorization": f"Bearer {api_key}", "Content-Type": "application/json"},
        json=payload,
        timeout=60,
    )
    if not resp.ok:
        sys.exit(f"API error {resp.status_code}: {resp.text}")
    return resp.json()["request_id"]


def poll_until_done(api_key: str, request_id: str) -> str:
    """Poll until video is ready; return video URL."""
    headers = {"Authorization": f"Bearer {api_key}"}
    deadline = time.time() + POLL_TIMEOUT
    print(f"Polling (request_id={request_id})...", flush=True)
    while time.time() < deadline:
        resp = requests.get(f"{API_BASE}/videos/{request_id}", headers=headers, timeout=30)
        data = resp.json()
        status = data.get("status", "unknown")
        if status == "done":
            return data["video"]["url"]
        elif status in ("failed", "expired"):
            sys.exit(f"Video generation {status}: {data}")
        print(f"  status={status}, waiting {POLL_INTERVAL}s...", flush=True)
        time.sleep(POLL_INTERVAL)
    sys.exit("Timed out waiting for video generation.")


def download_video(url: str, output_path: str) -> str:
    """Download video to output_path."""
    resp = requests.get(url, timeout=120, stream=True)
    resp.raise_for_status()
    with open(output_path, "wb") as f:
        for chunk in resp.iter_content(chunk_size=8192):
            f.write(chunk)
    return output_path


def main():
    parser = argparse.ArgumentParser(description="Animate images to video via Grok Imagine")
    parser.add_argument("--prompt", required=True, help="Motion/animation description prompt")
    parser.add_argument("--image", help="Path or URL to source image (for image-to-video)")
    parser.add_argument("--duration", type=int, default=8, choices=range(3, 11),
                        metavar="[3-10]", help="Video duration in seconds (default: 8)")
    parser.add_argument("--aspect-ratio", default="16:9",
                        choices=["16:9", "9:16", "1:1"], help="Aspect ratio (default: 16:9)")
    parser.add_argument("--resolution", default="720p", choices=["480p", "720p"],
                        help="Output resolution (default: 720p)")
    parser.add_argument("--output", help="Output .mp4 path (auto-named if omitted)")
    args = parser.parse_args()

    api_key = get_api_key()

    # Build payload
    timestamp = datetime.now().strftime("%Y-%m-%d-%H-%M-%S")
    output_path = args.output or f"{timestamp}-grok-video.mp4"

    payload = {
        "model": MODEL,
        "prompt": args.prompt,
        "duration": args.duration,
        "aspect_ratio": args.aspect_ratio,
        "resolution": args.resolution,
    }

    if args.image:
        print(f"Preparing image: {args.image}", flush=True)
        payload["image"] = prepare_image(args.image)
        print("Submitting image-to-video request...", flush=True)
    else:
        print("Submitting text-to-video request...", flush=True)

    request_id = start_generation(api_key, payload)
    video_url = poll_until_done(api_key, request_id)

    print(f"Downloading video from {video_url}...", flush=True)
    download_video(video_url, output_path)

    abs_path = str(Path(output_path).resolve())
    print(f"Saved: {abs_path}")
    print(f"MEDIA:{abs_path}")  # OpenClaw auto-attach signal


if __name__ == "__main__":
    main()