import os
import gradio as gr
import yt_dlp

def download_media(url, download_video):
    download_dir = 'downloads'
    os.makedirs(download_dir, exist_ok=True)
    
    if download_video:
        ydl_opts = {
            'format': 'bestvideo+bestaudio/best',
            'outtmpl': f'{download_dir}/%(title)s.%(ext)s',
        }
    else:
        ydl_opts = {
            'format': 'bestaudio/best',
            'postprocessors': [{
                'key': 'FFmpegExtractAudio',
                'preferredcodec': 'mp3',
                'preferredquality': '192',
            }],
            'outtmpl': f'{download_dir}/%(title)s.%(ext)s',
        }

    with yt_dlp.YoutubeDL(ydl_opts) as ydl:
        info_dict = ydl.extract_info(url, download=True)
        file_title = ydl.prepare_filename(info_dict)

        if download_video:
            output_file = file_title
        else:
            output_file = file_title.rsplit('.', 1)[0] + '.mp3'

    return output_file

def handle_download(url, download_video):
    output_file = download_media(url, download_video)
    if download_video:
        return gr.update(value=output_file, visible=True), gr.update(visible=False)
    else:
        return gr.update(visible=False), gr.update(value=output_file, visible=True)

with gr.Blocks(theme=gr.themes.Soft(primary_hue="orange", secondary_hue="orange")) as demo:
    gr.Markdown(f"# <div style='text-align: center;'>YOUTUBE Downloader</div>")
    gr.Markdown(f"## <div style='text-align: center;'>Download MP3/MP4 from YouTube URL</div>")
    
    with gr.Row():
        url_input = gr.Textbox(label="YouTube URL")
    with gr.Row():
        download_video_checkbox = gr.Checkbox(label="Download Video", value=False)
    with gr.Row():    
        download_button = gr.Button("Download")
    with gr.Row():
        output_audio = gr.Audio(label="Downloaded Media", visible=False)
        output_file = gr.File(label="Downloaded Media", visible=False)

    download_button.click(
        handle_download, 
        inputs=[url_input, download_video_checkbox], 
        outputs=[output_file, output_audio]
    )

    gr.Markdown(f"### <div style='text-align: center;'>Made with ❤ by <a href='https://huggingface.co/Gradio-Blocks'>Gradio-Blocks-Party-Member</a></div>")

demo.launch()