import gradio as gr
import yt_dlp
import os
import time
import threading

def download_video(url, resolution):
    # Define resolution options
    resolutions = {
        "2160p": "bestvideo[height<=2160]+bestaudio/best",
        "1440p": "bestvideo[height<=1440]+bestaudio/best",
        "1080p": "bestvideo[height<=1080]+bestaudio/best",
        "720p": "bestvideo[height<=720]+bestaudio/best",
        "480p": "bestvideo[height<=480]+bestaudio/best",
    }

    try:
        # Set yt-dlp options
        ydl_opts = {
            'format': resolutions[resolution],
            'outtmpl': '%(title)s.%(ext)s',
            'noplaylist': True,
            'merge_output_format': 'mp4',
            'postprocessors': [{
                'key': 'FFmpegVideoConvertor',
                'preferedformat': 'mp4',  # Ensure output is in MP4 format
            }],
        }

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

            # Check if the file already exists
            if os.path.exists(filename):
                os.remove(filename)  # Remove the existing file

            ydl.download([url])

        yield filename  # Yield the filename to display the video
        time.sleep(10)  # Wait for 20 seconds before deleting the file
        if os.path.exists(filename):
            os.remove(filename)  # Remove the file after 20 seconds

    except Exception as e:
        yield None
        print("An error occurred:", e)

def download_audio(url):
    try:
        # Set yt-dlp options for audio download
        ydl_opts = {
            'format': 'bestaudio/best',
            'outtmpl': '%(title)s.%(ext)s',
            'noplaylist': True,
            'postprocessors': [{
                'key': 'FFmpegExtractAudio',
                'preferredcodec': 'mp3',
            }],
        }

        with yt_dlp.YoutubeDL(ydl_opts) as ydl:
            info_dict = ydl.extract_info(url, download=True)
            filename = ydl.prepare_filename(info_dict).replace('.webm', '.mp3').replace('.m4a', '.mp3')

            # Check if the file already exists
            if os.path.exists(filename):
                os.remove(filename)  # Remove the existing file

            ydl.download([url])

        yield filename  # Yield the filename to display the audio
        time.sleep(20)  # Wait for 20 seconds before deleting the file
        if os.path.exists(filename):
            os.remove(filename)  # Remove the file after 20 seconds

    except Exception as e:
        yield None
        print("An error occurred:", e)

# Create Gradio interface
with gr.Blocks() as iface:
    with gr.Tab("MP4 Downloader"):
        gr.Markdown("### Download YouTube videos as MP4")
        url_input_mp4 = gr.Textbox(label="YouTube URL", placeholder="Enter YouTube URL here")
        resolution_input = gr.Dropdown(label="Resolution", choices=["2160p", "1440p", "1080p", "720p", "480p"], value="1080p")
        download_button_mp4 = gr.Button("Download MP4")
        output_mp4 = gr.Video(label="Downloaded Video")
        
        download_button_mp4.click(download_video, inputs=[url_input_mp4, resolution_input], outputs=output_mp4)
    
    with gr.Tab("MP3 Downloader"):
        gr.Markdown("### Download YouTube videos as MP3")
        url_input_mp3 = gr.Textbox(label="YouTube URL", placeholder="Enter YouTube URL here")
        download_button_mp3 = gr.Button("Download MP3")
        output_mp3 = gr.Audio(label="Downloaded Audio")
        
        download_button_mp3.click(download_audio, inputs=[url_input_mp3], outputs=output_mp3)

iface.launch()