import gradio as gr
import yt_dlp
import os
import tempfile
from glob import glob

def download_media(urls, output_format, audio_quality, custom_bitrate, embed_metadata, embed_thumbnail, write_description):
    # Create a temporary directory to store downloads
    temp_dir = tempfile.mkdtemp()
    output_template = os.path.join(temp_dir, '%(title)s.%(ext)s')

    # Prepare the options
    ydl_opts = {
        'outtmpl': output_template,
        'format': 'bestaudio/best',
        'postprocessors': [],
        'cookiesfrombrowser': ('edge',),
    }

    # Set format options
    if output_format in ['mp3', 'opus']:
        extract_audio_postprocessor = {
            'key': 'FFmpegExtractAudio',
            'preferredcodec': output_format,
        }
        # Set postprocessor args
        if audio_quality == 'Best':
            extract_audio_postprocessor['preferredquality'] = '0'  # Best quality
        elif audio_quality == 'Custom' and custom_bitrate:
            extract_audio_postprocessor['preferredquality'] = '5'  # Default value
            extract_audio_postprocessor['postprocessor_args'] = ['-b:a', custom_bitrate]
        ydl_opts['postprocessors'].append(extract_audio_postprocessor)
    else:
        ydl_opts['format'] = 'bestvideo+bestaudio/best'

    # Additional options
    if embed_metadata:
        ydl_opts['addmetadata'] = True
        ydl_opts['postprocessors'].append({'key': 'FFmpegMetadata'})
    if embed_thumbnail:
        ydl_opts['writethumbnail'] = True
        ydl_opts['postprocessors'].append({'key': 'EmbedThumbnail'})
    if write_description:
        ydl_opts['writedescription'] = True

    # Process URLs
    url_list = urls.strip().split()
    with yt_dlp.YoutubeDL(ydl_opts) as ydl:
        ydl.download(url_list)

    # Collect output files
    downloaded_files = glob(os.path.join(temp_dir, '*'))
    # Return the list of file paths for download
    return downloaded_files

def main():
    with gr.Blocks() as demo:
        gr.Markdown("# YT-DLP Gradio Interface")
        with gr.Row():
            url_input = gr.Textbox(label="Video URL(s)", placeholder="Enter one or more URLs separated by space or newline")
        with gr.Row():
            output_format = gr.Dropdown(['mp3', 'opus', 'bestvideo+bestaudio'], label="Output Format", value='mp3')
            audio_quality = gr.Dropdown(['Best', 'Custom'], label="Audio Quality", value='Best')
            custom_bitrate = gr.Textbox(label="Custom Bitrate (e.g., 192k)", visible=False)
        with gr.Row():
            embed_metadata = gr.Checkbox(label="Embed Metadata")
            embed_thumbnail = gr.Checkbox(label="Embed Thumbnail")
            write_description = gr.Checkbox(label="Write Description")
        start_button = gr.Button("Start Download")
        output_files = gr.Files(label="Downloaded Files")

        def show_custom_bitrate(audio_quality):
            if audio_quality == 'Custom':
                return gr.update(visible=True)
            else:
                return gr.update(visible=False)

        audio_quality.change(show_custom_bitrate, inputs=audio_quality, outputs=custom_bitrate)

        start_button.click(download_media,
                           inputs=[url_input, output_format, audio_quality, custom_bitrate, embed_metadata, embed_thumbnail, write_description],
                           outputs=output_files)
    demo.launch()

if __name__ == "__main__":
    main()