#!/usr/bin/env python3
"""
LaoZhang Veo 3.1 Async API client.

Based on: https://docs.laozhang.ai/api-capabilities/veo/veo-31-async-api

Key advantage: 失败不扣费 (no charge on failure)
"""

import json
import time
import requests
from pathlib import Path
from typing import Optional


BASE_URL = "https://api.laozhang.ai/v1"


def call_veo3_laozhang_async(
    prompt: str,
    api_key: str,
    first_frame_path: Optional[str] = None,
    last_frame_path: Optional[str] = None,
    model: str = "veo-3.1-fl",
    logger=None
) -> dict:
    """
    Create video generation task using LaoZhang async API.

    Args:
        prompt: Video generation prompt
        api_key: LaoZhang API key
        first_frame_path: Path to first frame image (uploaded as file)
        last_frame_path: Path to last frame image for interpolation (uploaded as file)
        model: Model to use (must have -fl suffix for image-to-video)
        logger: Optional logger instance

    Returns:
        Dict with task_id for polling
    """
    msg = f"  Model: {model}"
    print(msg)
    if logger:
        logger.info(msg)

    # Collect image files if provided
    image_paths = []
    if first_frame_path and Path(first_frame_path).exists():
        image_paths.append(first_frame_path)
        msg = f"  First frame: {Path(first_frame_path).name}"
        print(msg)
        if logger:
            logger.info(msg)

    if last_frame_path and Path(last_frame_path).exists():
        image_paths.append(last_frame_path)
        msg = f"  Last frame: {Path(last_frame_path).name}"
        print(msg)
        if logger:
            logger.info(msg)

    msg = "  Sending async request to LaoZhang..."
    print(msg)
    if logger:
        logger.info(msg)

    headers = {"Authorization": f"Bearer {api_key}"}

    if image_paths:
        # Image-to-video: Use multipart/form-data with file uploads
        files = []
        for path in image_paths:
            files.append(("input_reference", (Path(path).name, open(path, 'rb'), "image/jpeg")))

        data = {"model": model, "prompt": prompt}
        response = requests.post(
            f"{BASE_URL}/videos",
            headers=headers,
            files=files,
            data=data,
            timeout=60
        )

        # Close file handles
        for _, (_, f, _) in files:
            f.close()
    else:
        # Text-to-video: Use JSON format
        headers["Content-Type"] = "application/json"
        data = {"model": model, "prompt": prompt}
        response = requests.post(
            f"{BASE_URL}/videos",
            headers=headers,
            json=data,
            timeout=60
        )

    if response.status_code != 200:
        error_msg = f"LaoZhang async API error: {response.status_code} - {response.text[:500]}"
        if logger:
            logger.error(error_msg)
        raise RuntimeError(error_msg)

    result = response.json()

    # Response format: {"id": "video_xxx", ...}
    task_id = result.get("id") or result.get("task_id")
    if not task_id:
        error_msg = f"No task_id in response: {result}"
        if logger:
            logger.error(error_msg)
        raise RuntimeError(error_msg)

    msg = f"  Task created: {task_id}"
    print(msg)
    if logger:
        logger.info(msg)
    return {"task_id": task_id, "raw": result}


def poll_veo3_laozhang_async(
    task_id: str,
    api_key: str,
    poll_interval: int = 10,
    timeout: int = 600,
    logger=None
) -> dict:
    """
    Poll LaoZhang async API for task completion.

    Args:
        task_id: Task ID from call_veo3_laozhang_async()
        api_key: LaoZhang API key
        poll_interval: Seconds between polls
        timeout: Maximum wait time in seconds
        logger: Optional logger instance

    Returns:
        Dict with video URL
    """
    start_time = time.time()
    consecutive_errors = 0
    max_errors = 5

    while time.time() - start_time < timeout:
        try:
            # Query task status
            response = requests.get(
                f"{BASE_URL}/videos/{task_id}",
                headers={
                    "Authorization": f"Bearer {api_key}"
                },
                timeout=30
            )

            if response.status_code != 200:
                raise RuntimeError(f"LaoZhang poll error: {response.status_code} - {response.text[:500]}")

            result = response.json()
            consecutive_errors = 0  # Reset on success

        except Exception as e:
            consecutive_errors += 1
            elapsed = int(time.time() - start_time)
            print(f"  Network error ({consecutive_errors}/{max_errors}): {e}")
            if consecutive_errors >= max_errors:
                raise RuntimeError(f"Too many consecutive network errors: {e}")
            time.sleep(poll_interval)
            continue

        # Check status OUTSIDE try block so exceptions aren't caught
        # Response format: {"id": "...", "status": "queued|processing|completed|failed", ...}
        status = result.get("status", "")

        if status == "completed":
            # Video URL is in the status response itself
            video_url = (
                result.get("url") or
                result.get("video_url") or
                result.get("result_url")
            )

            if not video_url:
                error_msg = f"No video URL in completed response: {result}"
                if logger:
                    logger.error(error_msg)
                raise RuntimeError(error_msg)

            msg = f"  Video generation completed: {video_url}"
            if logger:
                logger.info(msg)
            return {"uri": video_url, "state": "COMPLETE", "data": result}

        if status == "failed":
            error_msg = result.get("error", {}).get("message") if isinstance(result.get("error"), dict) else "Unknown error"
            full_error = f"LaoZhang generation failed: {error_msg}"
            if logger:
                logger.error(full_error)
            raise RuntimeError(full_error)

        elapsed = int(time.time() - start_time)
        msg = f"  LaoZhang Veo3 generating... ({elapsed}s elapsed, status: {status})"
        print(msg)
        if logger:
            logger.info(msg)

        time.sleep(poll_interval)

    raise TimeoutError(f"LaoZhang Veo3 generation timed out after {timeout}s")


def download_video_laozhang(video_url: str, output_path: str) -> str:
    """Download video from URL."""
    response = requests.get(video_url, stream=True, timeout=300)

    if response.status_code != 200:
        raise RuntimeError(f"Download error: {response.status_code}")

    with open(output_path, 'wb') as f:
        for chunk in response.iter_content(chunk_size=8192):
            f.write(chunk)

    return output_path


if __name__ == "__main__":
    import sys
    import os
    from dotenv import load_dotenv

    load_dotenv(Path(__file__).parent.parent / ".env")

    api_key = os.environ.get("LAOZHANG_KEY")
    if not api_key:
        print("Error: LAOZHANG_KEY not set in .env")
        sys.exit(1)

    if len(sys.argv) < 2:
        print("Usage:")
        print("  python laozhang_veo3_async.py generate <prompt> [first_frame_path] [last_frame_path]")
        sys.exit(1)

    cmd = sys.argv[1]

    if cmd == "generate":
        if len(sys.argv) < 3:
            print("Usage: python laozhang_veo3_async.py generate <prompt> [first_frame_path] [last_frame_path]")
            sys.exit(1)

        prompt = sys.argv[2]
        first_frame_path = sys.argv[3] if len(sys.argv) > 3 else None
        last_frame_path = sys.argv[4] if len(sys.argv) > 4 else None

        print(f"Prompt: {prompt}")
        if first_frame_path:
            print(f"First frame: {first_frame_path}")
        if last_frame_path:
            print(f"Last frame: {last_frame_path}")

        # Auto-select model based on whether frames are provided
        has_frames = first_frame_path or last_frame_path
        model = "veo-3.1-fl" if has_frames else "veo-3.1"

        result = call_veo3_laozhang_async(
            prompt=prompt,
            api_key=api_key,
            first_frame_path=first_frame_path,
            last_frame_path=last_frame_path,
            model=model
        )

        print(f"\nTask created: {result['task_id']}")
        print("Polling for completion...")

        video_result = poll_veo3_laozhang_async(
            task_id=result["task_id"],
            api_key=api_key,
            poll_interval=5,
            timeout=600
        )

        print(f"\nVideo URL: {video_result['uri']}")

    else:
        print(f"Unknown command: {cmd}")
