import gradio as gr
import yt_dlp

def download_media(url, download_video):
    if download_video:
        ydl_opts = {
            'format': 'bestvideo+bestaudio/best',
            'outtmpl': 'downloads/%(title)s.%(ext)s',
        }
    else:
        ydl_opts = {
            'format': 'bestaudio/best',
            'postprocessors': [{
                'key': 'FFmpegExtractAudio',
                'preferredcodec': 'mp3',
                'preferredquality': '192',
            }],
            'outtmpl': 'downloads/%(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 get_output_component(download_video):
    if download_video:
        return gr.File(label="Downloaded Media")
    else:
        return gr.Audio(label="Downloaded Media")

# Create the Gradio interface
with gr.Blocks() as demo:
    gr.Markdown(f"# <div style='text-align: center;'> YOUTUBE Downloader</div>")
    gr.Markdown(f"## <div style='text-align: center;'> download mp3/mp4 form 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)

    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)
    gr.Markdown(f"### <div style='text-align: center;'>made with ❤ by <a href='https://huggingface.co/Hev832'>Hev832</a></div>")
    
    download_button.click(
        handle_download, 
        inputs=[url_input, download_video_checkbox], 
        outputs=[output_file, output_audio]
    )

demo.launch()