import argparse
import os
import requests
import subprocess
from pathlib import Path
from PIL import Image

def download_images(urls, out_dir="frames"):
    os.makedirs(out_dir, exist_ok=True)
    image_paths = []
    for idx, url in enumerate(urls):
        ext = url.split(".")[-1].split("?")[0]
        ext = "jpg" if ext not in ["png", "jpeg", "jpg"] else ext
        out_path = f"{out_dir}/img{idx:03d}.{ext}"
        r = requests.get(url)
        if r.ok:
            with open(out_path, "wb") as f:
                f.write(r.content)
            image_paths.append(out_path)
    return image_paths

def convert_to_video(images, audio_file, output="output.mp4"):
    temp_dir = Path("temp")
    temp_dir.mkdir(exist_ok=True)

    # Rename images to sequential format
    for idx, img_path in enumerate(images):
        ext = img_path.split(".")[-1]
        new_name = temp_dir / f"frame{idx:03d}.{ext}"
        Image.open(img_path).save(new_name)

    # Get duration of audio
    cmd = [
        "ffprobe", "-v", "error", "-show_entries", "format=duration",
        "-of", "default=noprint_wrappers=1:nokey=1", audio_file
    ]
    result = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    duration = float(result.stdout.decode().strip())
    img_duration = duration / len(images)

    # ffmpeg command to create video
    ffmpeg_cmd = [
        "ffmpeg", "-y", "-r", f"1/{img_duration}", "-i", f"{temp_dir}/frame%03d.jpg",
        "-i", audio_file, "-c:v", "libx264", "-t", f"{duration}", "-pix_fmt", "yuv420p",
        "-c:a", "aac", "-shortest", output
    ]

    subprocess.run(ffmpeg_cmd)
    print(f" Video created: {output}")

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("audio", help="Path to TTS mp3 file")
    parser.add_argument("--urls", nargs="+", required=True, help="List of image URLs")
    args = parser.parse_args()

    print(" Downloading images...")
    images = download_images(args.urls)

    print(" Creating video...")
    convert_to_video(images, args.audio)

if __name__ == "__main__":
    main()