import torch
from .utils import (
    VIDEO_ASPECT_RATIOS,
    call_api_with_fallback,
    resolve_api_key,
    resolve_generation_url,
    upload_image,
)


def _get_video_url(data: dict) -> str:
    url = data.get("generation_url", "")
    if not url:
        raise RuntimeError(f"Pruna API returned no generation_url. Response: {data}")
    return resolve_generation_url(url)


class PrunaVideo:
    """Generate a video from a text prompt and/or an image using the Pruna p-video model.

    When an image is connected the model animates it; without an image it generates
    from the text prompt alone. Returns the video URL as a STRING.

    Pricing (p-video):
      720p draft OFF  — $0.02/sec  |  720p draft ON  — $0.005/sec
      1080p draft OFF — $0.04/sec  |  1080p draft ON — $0.01/sec
    """

    CATEGORY = "pruna ai"
    RETURN_TYPES = ("STRING",)
    RETURN_NAMES = ("video_url",)
    FUNCTION = "generate"

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "prompt": ("STRING", {"multiline": True, "default": ""}),
                "duration": ("INT", {"default": 5, "min": 1, "max": 20, "step": 1}),
                "resolution": (["720p", "1080p"], {"default": "720p"}),
                "aspect_ratio": (VIDEO_ASPECT_RATIOS, {"default": "16:9"}),
                "draft": ("BOOLEAN", {"default": False}),
                "api_key": ("STRING", {"default": ""}),
            },
            "optional": {
                "image": ("IMAGE",),
            },
        }

    def generate(
        self,
        prompt: str,
        duration: int,
        resolution: str,
        aspect_ratio: str,
        draft: bool,
        api_key: str,
        image: torch.Tensor = None,
    ):
        key = resolve_api_key(api_key)
        payload: dict = {
            "input": {
                "prompt": prompt,
                "duration": duration,
                "resolution": resolution,
                "aspect_ratio": aspect_ratio,
                "draft": draft,
            }
        }

        if image is not None:
            payload["input"]["image"] = upload_image(image, key)

        data = call_api_with_fallback("p-video", payload, key)
        return (_get_video_url(data),)