import gradio as gr
from moviepy.audio.io.AudioFileClip import AudioFileClip
import yt_dlp
import os

def download_youtube(url, output_format, resolution, audio_bitrate):
    try:
        output_dir = "downloads"
        os.makedirs(output_dir, exist_ok=True)
        output_template = f"{output_dir}/%(title)s.%(ext)s"

        ydl_opts = {
            "format": "bestaudio/best" if output_format in ["MP3", "WAV"] else f"bestvideo[height<={resolution[:-1]}]+bestaudio/best",
            "outtmpl": output_template,
            "postprocessors": [],
        }

        if output_format == "MP3":
            ydl_opts["postprocessors"].append({
                'key': 'FFmpegExtractAudio',
                'preferredcodec': 'mp3',
                'preferredquality': audio_bitrate.replace("k", ""),
            })
        elif output_format == "WAV":
            ydl_opts["postprocessors"].append({
                'key': 'FFmpegExtractAudio',
                'preferredcodec': 'wav',
            })

        with yt_dlp.YoutubeDL(ydl_opts) as ydl:
            info_dict = ydl.extract_info(url, download=True)
            filename = ydl.prepare_filename(info_dict)
            if output_format in ["MP3", "WAV"]:
                filename = os.path.splitext(filename)[0] + f".{output_format.lower()}"
            return "Download complete", filename

    except Exception as e:
        return f"Error: {str(e)}", None


formats = ["MP3", "WAV", "MP4"]
resolutions = ["360p", "480p", "720p", "1080p"]
bitrates = ["64k", "128k", "192k", "256k", "320k"]

interface = gr.Interface(
    fn=download_youtube,
    inputs=[
        gr.Text(label="YouTube Video URL"),
        gr.Radio(formats, label="Select Format"),
        gr.Dropdown(resolutions, label="Video Resolution (for MP4)", value="720p"),
        gr.Dropdown(bitrates, label="Audio Bitrate (for MP3)", value="128k"),
    ],
    outputs=[
        gr.Text(label="Status"),
        gr.File(label="Download File"),
    ],
    title="🎵 YouTube Downloader",
    description="Download YouTube videos as MP3, WAV, or MP4 with quality selection. Built with Pytube + Gradio."
)

if __name__ == "__main__":
    os.makedirs("downloads", exist_ok=True)
    interface.launch()