import spaces
import demucs.separate
import subprocess
import os
import uuid
import gradio as gr
from pydub import AudioSegment
import shutil

@spaces.GPU
def process_song(url):
    uuid_str = str(uuid.uuid4())
    cmd = f"spotdl {url} --output audio{uuid_str}"
    subprocess.run(cmd, shell=True)
    files = os.listdir(f"audio{uuid_str}")
    audio_files = [file for file in files if file.endswith((".mp3", ".wav", ".ogg"))]
    old_path = os.path.join(f"audio{uuid_str}", audio_files[0])
    new_path = os.path.join(f"audio{uuid_str}", "audio.mp3")
    os.rename(old_path, new_path)

    demucs.separate.main(
        [
            "--mp3",
            "--two-stems",
            "vocals",
            "-n",
            "htdemucs",
            new_path,
            "-o",
            f"output{uuid_str}",
        ]
    )

    sound = AudioSegment.from_mp3(f"./output{uuid_str}/htdemucs/audio/no_vocals.mp3")
    sound.export("output.mp3", format="mp3", bitrate="192k")

    shutil.rmtree(f"audio{uuid_str}")
    shutil.rmtree(f"output{uuid_str}")

    return "output.mp3"


iface = gr.Interface(
    fn=process_song,
    inputs=gr.Textbox(placeholder="Enter the Spotify URL", label="Song"),
    outputs=gr.Audio(label="Karaoke audio"),
)
iface.launch()