import gradio as gr
import yt_dlp
import os

def download_youtube(url, file_format):
    ydl_opts = {}
    try:
        if file_format == "audio":
            ydl_opts = {
                'format': 'bestaudio/best',
                'postprocessors': [{
                    'key': 'FFmpegExtractAudio',
                    'preferredcodec': 'mp3',
                    'preferredquality': '192',
                }],
                'outtmpl': 'downloads/%(title)s.%(ext)s',
            }
        else:
            ydl_opts = {
                'format': 'bestvideo+bestaudio/best',
                'outtmpl': 'downloads/%(title)s.%(ext)s',
            }

        with yt_dlp.YoutubeDL(ydl_opts) as ydl:
            info_dict = ydl.extract_info(url, download=True)
            file_path = ydl.prepare_filename(info_dict)
            if file_format == "audio":
                file_path = os.path.splitext(file_path)[0] + ".mp3"

        return file_path
    except Exception as e:
        return str(e)

def show_output(file_path):
    if not file_path or not os.path.exists(file_path):
        return "Error downloading file.", None, None
    if file_path.endswith(".mp3"):
        return None, gr.Audio(file_path, label="Audio Output"), None
    else:
        return None, None, gr.Video(file_path, label="Video Output")

with gr.Blocks() as demo:
    with gr.Row():
        with gr.Column():
            url_input = gr.Textbox(label="YouTube URL")
            format_input = gr.Dropdown(choices=["audio", "video"], label="Format")
            download_button = gr.Button("Download")
        with gr.Column():
            error_output = gr.Textbox(label="Error Output", visible=False)
            audio_output = gr.Audio(label="Audio Output", visible=False)
            video_output = gr.Video(label="Video Output", visible=False)

    download_button.click(download_youtube, inputs=[url_input, format_input], outputs=None).then(
        show_output, inputs=None, outputs=[error_output, audio_output, video_output]
    )

demo.launch()