import gradio as gr
from yt_dlp import YoutubeDL
import os
from pydub import AudioSegment

DOWNLOADS_FOLDER = "downloads"
os.makedirs(DOWNLOADS_FOLDER, exist_ok=True)

def download_audio(url, file_format, duration_sec):
    """
    Download audio from YouTube or SoundCloud and convert to desired format
    """
    try:
        # Download best audio available
        ydl_opts = {
            'format': 'bestaudio/best',
            'outtmpl': os.path.join(DOWNLOADS_FOLDER, '%(title)s.%(ext)s'),
            'quiet': True,
            'no_warnings': True,
            'extract_flat': False,
        }
        
        with YoutubeDL(ydl_opts) as ydl:
            info = ydl.extract_info(url, download=True)
        
        original_file = os.path.join(DOWNLOADS_FOLDER, f"{info['title']}.{info['ext']}")
        
        # Clean filename (remove special characters)
        safe_title = "".join(c for c in info['title'] if c.isalnum() or c in (' ', '-', '_')).rstrip()
        original_file_safe = os.path.join(DOWNLOADS_FOLDER, f"{safe_title}.{info['ext']}")
        
        # Rename if necessary
        if original_file != original_file_safe:
            os.rename(original_file, original_file_safe)
            original_file = original_file_safe
        
        # Load audio and trim to duration
        audio = AudioSegment.from_file(original_file)
        duration_ms = min(len(audio), int(duration_sec * 1000))
        trimmed_audio = audio[:duration_ms]
        
        # Determine output file path based on desired format
        base_name = os.path.splitext(original_file)[0]
        
        if file_format.lower() == "mp3":
            output_file = base_name + ".mp3"
            trimmed_audio.export(output_file, format="mp3", bitrate="192k")
        
        elif file_format.lower() == "opus":
            output_file = base_name + ".opus"
            trimmed_audio.export(output_file, format="opus", bitrate="128k")
        
        elif file_format.lower() == "wav":
            output_file = base_name + ".wav"
            trimmed_audio.export(output_file, format="wav")
        
        elif file_format.lower() == "m4a":
            output_file = base_name + ".m4a"
            trimmed_audio.export(output_file, format="ipod", codec="aac")
        
        else:  # keep original format
            output_file = original_file
        
        # Optional: Remove original downloaded file to save space
        if output_file != original_file and os.path.exists(original_file):
            os.remove(original_file)
        
        return output_file
    
    except Exception as e:
        raise gr.Error(f"Error downloading audio: {str(e)}")

# Create Gradio interface
with gr.Blocks(title="Audio Downloader", theme=gr.themes.Soft()) as iface:
    gr.Markdown("# 🎵 YouTube & SoundCloud Audio Downloader")
    gr.Markdown("Download audio from YouTube or SoundCloud and convert to your preferred format")
    
    with gr.Row():
        url_input = gr.Textbox(
            label="Enter URL",
            placeholder="Paste YouTube or SoundCloud URL here...",
            scale=4
        )
    
    with gr.Row():
        with gr.Column(scale=1):
            format_choice = gr.Dropdown(
                choices=["mp3", "m4a", "opus", "wav"],
                value="mp3",
                label="Output Format"
            )
        
        with gr.Column(scale=3):
            duration_slider = gr.Slider(
                minimum=5,
                maximum=600,  # Increased to 10 minutes
                value=60,
                step=5,
                label="Duration (seconds)",
                info="Maximum duration to download (0 = full track)"
            )
    
    with gr.Row():
        download_button = gr.Button("Download Audio", variant="primary", size="lg")
    
    with gr.Row():
        download_file = gr.File(label="Download your audio", interactive=False)
    
    # Example URLs
    with gr.Accordion("Example URLs (click to use)", open=False):
        gr.Examples(
            examples=[
                ["https://www.youtube.com/watch?v=dQw4w9WgXcQ"],
                ["https://soundcloud.com/antonio-antetomaso/mutiny-on-the-bounty-closing-titles-cover"],
                ["https://www.youtube.com/watch?v=JGwWNGJdvx8"],
                ["https://soundcloud.com/officialnikkig/kill-bill-sza-kill-bill-remix"]
            ],
            inputs=url_input,
            label=""
        )
    
    # Add instructions
    with gr.Accordion("Instructions", open=False):
        gr.Markdown("""
        ### How to use:
        1. Paste a YouTube or SoundCloud URL
        2. Choose your desired output format
        3. Set the maximum duration (or leave at 0 for full track)
        4. Click "Download Audio"
        
        ### Supported formats:
        - **MP3**: Most compatible, good quality
        - **M4A**: Apple format, smaller file size
        - **Opus**: Best quality at low bitrates
        - **WAV**: Lossless, large file size
        
        ### Notes:
        - Downloads are saved to a 'downloads' folder
        - Some tracks may have download restrictions
        - Very long videos may take time to process
        """)
    
    download_button.click(
        fn=download_audio,
        inputs=[url_input, format_choice, duration_slider],
        outputs=download_file,
        show_progress=True
    )

# Launch the interface
if __name__ == "__main__":
    iface.launch(
        show_error=True,
        share=False,  # Set to True to create a public link
        server_name="0.0.0.0" if os.getenv("GRADIO_SHARE") else "127.0.0.1",
        server_port=7860
    )