import gradio as gr
import yt_dlp

def download_media(url, media_type):
    ydl_opts = {
        'format': 'bestaudio/best' if media_type == 'Audio' else 'bestvideo+bestaudio/best',
        'outtmpl': '%(title)s.%(ext)s',
    }
    with yt_dlp.YoutubeDL(ydl_opts) as ydl:
        result = ydl.extract_info(url, download=True)
        filename = ydl.prepare_filename(result)
    return filename

def download_audio(url):
    return download_media(url, 'Audio')

def download_video(url):
    return download_media(url, 'Video')

with gr.Blocks() as demo:
    gr.Markdown("# YouTube Audio/Video Downloader")

    url_input = gr.Textbox(label="YouTube URL")
    media_type = gr.Radio(label="Select Media Type", choices=["Audio", "Video"], value="Audio")
    output = gr.Audio(label="Downloaded File")

    download_button = gr.Button("Download")
    download_button.click(
        fn=lambda url, media: download_audio(url) if media == 'Audio' else download_video(url),
        inputs=[url_input, media_type],
        outputs=output
    )

demo.launch()