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

# Ensure downloads directory exists
os.makedirs("downloads", exist_ok=True)

def sanitize_filename(filename):
    """Sanitize filenames by replacing special characters."""
    return re.sub(r'[^a-zA-Z0-9_-]', '_', filename)

def process_youtube_or_audio(url, uploaded_audio, start_time, end_time):
    """Downloads or processes audio, trims it, and exports ringtones."""
    try:
        filename, song_name = None, None

        if url:
            ydl_opts = {
                'format': 'm4a/bestaudio',
                'outtmpl': 'downloads/%(title)s.%(ext)s',
                'cookiefile': 'cookies.txt'
            }
            with yt_dlp.YoutubeDL(ydl_opts) as ydl:
                info = ydl.extract_info(url, download=True)
                filename = ydl.prepare_filename(info)
                song_name = sanitize_filename(info['title'])

                # Ensure correct extension
                if not filename.endswith(".m4a"):
                    possible_file = f"downloads/{song_name}.m4a"
                    if os.path.exists(possible_file):
                        filename = possible_file

        elif uploaded_audio is not None:
            filename = uploaded_audio
            song_name = sanitize_filename(os.path.splitext(os.path.basename(uploaded_audio))[0])

        # Validate file existence
        if not filename or not os.path.exists(filename):
            print(f"File {filename} not found!")
            return None, None

        # Process audio
        audio = AudioSegment.from_file(filename)
        start_ms, end_ms = int(start_time * 1000), int(end_time * 1000)
        start_ms = max(0, min(start_ms, len(audio)))
        end_ms = max(start_ms, min(end_ms, len(audio)))

        trimmed_audio = audio[start_ms:end_ms]

        # Save MP3
        mp3_filename = f"downloads/{song_name}.mp3"
        trimmed_audio.export(mp3_filename, format="mp3")

        # Rename M4A to M4R for iPhone
        m4a_filename = filename
        m4r_filename = f"downloads/{song_name}.m4r"
        os.rename(m4a_filename, m4r_filename)

        return mp3_filename, m4r_filename

    except Exception as e:
        print(f"Error: {e}")
        return None, None

# Gradio Interface
with gr.Blocks() as interface:
    gr.HTML("""
        <h1 style="font-size: 2rem; text-align: center;"><i class="fas fa-music"></i>&nbsp;PYTR - Python YouTube Ringtones</h1>
        <p style="text-align: center;">Enter a YouTube URL or upload an audio file to create ringtones.</p>
        <p style="text-align: center;"><a href="https://ringtones.jessejesse.xyz">Ringtones</a> | <a href="https://github.com/sudo-self/pytr">GitHub</a></p>
    """)

    with gr.Row():
        youtube_url = gr.Textbox(placeholder="Enter the URL here...", label="YouTube URL")
        uploaded_audio = gr.File(label="Upload Audio File", type="filepath")

    gr.HTML("<h3 style='text-align: center;'>Trim Audio</h3>")

    with gr.Row():
        start_time = gr.Slider(0, 20, value=0, label="Start Time (seconds)")
        end_time = gr.Slider(1, 20, value=20, label="End Time (seconds)")

    process_button = gr.Button("Create Ringtones")

    with gr.Row():
        mp3_download = gr.File(label="Android Ringtone")
        iphone_ringtone = gr.File(label="iPhone Ringtone")

    with gr.Row():
        instructions = gr.Textbox(
            value="Android: move the download to the 'Ringtones' folder.\nApple: import download with 'Garage Band' and export to ringtones.",
            label="Install Ringtones",
            interactive=False
        )

    process_button.click(process_youtube_or_audio, inputs=[youtube_url, uploaded_audio, start_time, end_time], outputs=[mp3_download, iphone_ringtone])

interface.launch(share=True)