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 download_handler(url, download_video):
    if url:
        output_file = download_media(url, download_video)
        return "Download completed!", output_file
    else:
        return "Please enter a YouTube URL", None

with gr.Blocks() as demo:
    gr.Markdown("# Download From YOUTUBE")
    gr.Markdown("## Download MP3/MP4 from YouTube URL")
    
    url = gr.Textbox(label="YouTube URL")
    download_video = gr.Checkbox(label="Download Video", value=False)
    download_btn = gr.Button("Download")
    
    output_text = gr.Markdown()
    output_file = gr.File()

    download_btn.click(download_handler, inputs=[url, download_video], outputs=[output_text, output_file])

    with gr.Sidebar():
        gr.Image("profile.jpg", label="Profile Image", width=200)
        gr.Markdown("### AI/ML Engineer")
        gr.Markdown("Build AI with Haseeb")
        gr.Markdown("### Made with ❤ by [Haseeb Ahmed](https://www.linkedin.com/in/muhammad-haseeb-ahmed-1954b5230/)")

demo.launch()