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

if not os.path.exists("downloads"):
    os.makedirs("downloads")

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

def process_youtube_or_audio(url, uploaded_audio, start_time, end_time):
    try:
        filename = None
        song_name = None

        print(f"URL: {url}")  # Debug: Check URL input
        print(f"Uploaded audio: {uploaded_audio}")  # Debug: Check if audio is uploaded
        print(f"Start time: {start_time}, End time: {end_time}")  # Debug: Check times

        # Process YouTube URL
        if url:
            print("Processing YouTube URL...")
            ydl_opts = {
                'format': 'bestaudio/best',
                'outtmpl': 'downloads/%(id)s.%(ext)s',
                'cookiefile': 'cookies.txt'
            }
            with yt_dlp.YoutubeDL(ydl_opts) as ydl:
                info = ydl.extract_info(url, download=True)
                filename = os.path.join('downloads', f"{info['id']}.webm")
                song_name = sanitize_filename(info['title'])

        # Process uploaded audio file
        elif uploaded_audio:
            print("Processing uploaded audio...")
            filename = uploaded_audio.name  # Using the uploaded file's name
            song_name = "uploaded_audio"

        if not filename or not os.path.exists(filename):
            print(f"Error: No valid file found at {filename}")  # Debug
            return None, None

        # Load audio and trim it
        audio = AudioSegment.from_file(filename)
        start_time_ms = start_time * 1000
        end_time_ms = end_time * 1000
        start_time_ms = max(0, min(start_time_ms, len(audio)))
        end_time_ms = max(start_time_ms, min(end_time_ms, len(audio)))

        trimmed_audio = audio[start_time_ms:end_time_ms]

        # Export trimmed audio
        mp3_filename = f"downloads/{song_name}.mp3"
        trimmed_audio.export(mp3_filename, format="mp3")

        # Convert to m4r format for iPhone
        m4a_filename = f"downloads/{song_name}.m4a"
        try:
            subprocess.run([
                'ffmpeg', '-i', mp3_filename, '-vn', '-acodec', 'aac', '-b:a', '192k', m4a_filename
            ], stdout=subprocess.PIPE, stderr=subprocess.PIPE, timeout=30, check=True)
        except subprocess.TimeoutExpired:
            print("Error: ffmpeg took too long to process the file.")
            return None, None
        except subprocess.CalledProcessError as e:
            print(f"ffmpeg error: {e}")
            return None, None

        # Ensure the file exists before renaming
        if os.path.exists(m4a_filename):
            m4r_filename = f"downloads/{song_name}.m4r"
            os.rename(m4a_filename, m4r_filename)
            print(f"Files saved: {mp3_filename}, {m4r_filename}")  # Debug
            return os.path.abspath(mp3_filename), os.path.abspath(m4r_filename)
        else:
            print("Error: M4A file not created.")
            return None, None

    except Exception as e:
        print(f"Error: {e}")  # Debug: Print the error
        return None, None