import os
import base64
from typing import List
import replicate
from dotenv import load_dotenv


load_dotenv()
replicate_client = replicate.Client(api_token= os.getenv("REPLICATE_API_KEY"))

def video_generation(
    image_bytes,
    prompt: str,
    aspect_ratio: str,
    resolution: str
):
    encoded_image = base64.b64encode(image_bytes).decode("utf-8")
    output = replicate_client.run(
        "google/veo-3",
        input={
            "image": f"data:image/png;base64,{encoded_image}",
            "prompt": prompt,
            "resolution": resolution,
            "aspect_ratio": aspect_ratio
        }
    )
    urls: List[str] = []
    if isinstance(output, list) and output:
        first = output[0]
        url = getattr(first, "url", str(first))
        urls = [url]
    elif isinstance(output, str):
        urls = [output]
    elif hasattr(output, "url"):
        urls = [getattr(output, "url")]
    if urls:
        return urls[0]