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

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

def sanitize_filename(filename):
    """Remove invalid characters from filename"""
    return re.sub(r'[<>:"/\\|?*]', '', filename)

def download_audio(url, file_format):
    try:
        if not url or not url.strip():
            return None, "Error: Please enter a URL"
        
        # Set download options
        ydl_opts = {
            'format': 'bestaudio/best',
            'outtmpl': os.path.join(DOWNLOADS_FOLDER, '%(title)s.%(ext)s'),
            'quiet': True,
            'no_warnings': True
        }
        
        with YoutubeDL(ydl_opts) as ydl:
            info = ydl.extract_info(url, download=True)
            original_filename = ydl.prepare_filename(info)
        
        # Get the actual downloaded filename
        if os.path.exists(original_filename):
            downloaded_file = original_filename
        else:
            # Sometimes the extension might be different
            base_name = os.path.splitext(original_filename)[0]
            for ext in ['.m4a', '.webm', '.mp3', '.opus', '.flac', '.wav']:
                possible_file = base_name + ext
                if os.path.exists(possible_file):
                    downloaded_file = possible_file
                    break
            else:
                downloaded_file = original_filename
        
        # If user wants specific format, convert if needed
        if file_format.lower() != "original":
            target_ext = f".{file_format.lower()}"
            current_ext = os.path.splitext(downloaded_file)[1].lower()
            
            if current_ext != target_ext:
                # Create sanitized output filename
                base_name = os.path.splitext(downloaded_file)[0]
                sanitized_base = sanitize_filename(os.path.basename(base_name))
                output_file = os.path.join(DOWNLOADS_FOLDER, sanitized_base + target_ext)
                
                try:
                    # Convert audio format
                    audio = AudioSegment.from_file(downloaded_file)
                    audio.export(output_file, format=file_format.lower())
                    # Optionally remove original file after conversion
                    # os.remove(downloaded_file)
                    downloaded_file = output_file
                except Exception as e:
                    print(f"Conversion failed: {e}")
                    # Return original file if conversion fails
                    pass
        
        return downloaded_file, "Download completed successfully!"
    
    except Exception as e:
        return None, f"Error: {str(e)}"

# Create Gradio interface
iface = gr.Interface(
    fn=download_audio,
    inputs=[
        gr.Textbox(label="YouTube or SoundCloud URL", placeholder="Enter URL here..."),
        gr.Dropdown(["original", "mp3", "m4a", "wav"], value="mp3", label="Output Format")
    ],
    outputs=[
        gr.File(label="Download Audio"),
        gr.Textbox(label="Status")
    ],
    title="🎵 Audio Downloader",
    description="Download audio from YouTube or SoundCloud. Supports MP3, M4A, WAV formats.",
    examples=[
        ["https://www.youtube.com/watch?v=dQw4w9WgXcQ", "mp3"],
        ["https://soundcloud.com/aviciiofficial/levels", "mp3"],
    ]
)

# Launch the interface
if __name__ == "__main__":
    iface.launch(share=False, show_error=True)