import gradio as gr
import yt_dlp
import os
from pydub import AudioSegment
os.system("pip install yt_dlp pydub")

def download_media(url, media_type):
    ydl_opts = {
        'format': 'bestaudio/best' if media_type == 'audio' else 'bestvideo+bestaudio',
        'outtmpl': 'downloads/%(title)s.%(ext)s',
        'postprocessors': [{
            'key': 'FFmpegExtractAudio',
            'preferredcodec': 'mp3',
            'preferredquality': '192',
        }] if media_type == 'audio' else []
    }
    
    with yt_dlp.YoutubeDL(ydl_opts) as ydl:
        info = ydl.extract_info(url, download=True)
        file_path = ydl.prepare_filename(info)
        
        if media_type == 'audio':
            file_path = file_path.rsplit('.', 1)[0] + '.mp3'
    
    return file_path

def split_audio(file_path, start_time, end_time):
    audio = AudioSegment.from_file(file_path)
    split_audio = audio[start_time*1000:end_time*1000]
    split_file_path = file_path.rsplit('.', 1)[0] + f'_{start_time}-{end_time}.mp3'
    split_audio.export(split_file_path, format="mp3")
    return split_file_path

def show_media(url, media_type, start_time, end_time):
    file_path = download_media(url, media_type)
    if media_type == 'audio':
        split_file_path = split_audio(file_path, start_time, end_time)
        return None, split_file_path
    else:
        return file_path, None

with gr.Blocks() as demo:
    gr.Markdown("<h1><center>Media Downloader</center></h1>")
    gr.Markdown("please `❤` this space 🤗 if helpful")
    url = gr.Textbox(label="YouTube URL")
    media_type = gr.Radio(label="Media Type", choices=["audio", "video"])
    start_time = gr.Number(label="Start Time (seconds)", value=0)
    end_time = gr.Number(label="End Time (seconds)", value=10)
    download_button = gr.Button("Download")
    
    video_output = gr.Video()
    audio_output = gr.Audio()
    
    download_button.click(show_media, inputs=[url, media_type, start_time, end_time], outputs=[video_output, audio_output])

demo.launch()