import streamlit as st
import requests
import os
import torchaudio
from pydub import AudioSegment
from audiocraft.models import MusicGen
from tempfile import NamedTemporaryFile

# CONFIGURATION
RESEMBLE_API_KEY = os.getenv("RESEMBLE_API_KEY", "your_resemble_api_key")
PROJECT_UUID = "96e7c9b4-voice"

st.set_page_config(page_title="Suno-Like Music App", layout="centered")
st.title("Suno-Like AI Music Generator")

# Lyrics Input
st.subheader("Step 1: Enter Lyrics")
lyrics = st.text_area("Write your lyrics here:", "We go hard, never stop, always chasing dreams")

# Prompt Input
st.subheader("Step 2: Describe Your Beat")
prompt = st.text_input("Beat style (e.g. dark trap, chill pop):", "dark emotional trap beat with heavy bass")

# Generate Instrumental
music_path = None
if st.button("Generate Music"):
    with st.spinner("Generating instrumental..."):
        model = MusicGen.get_pretrained("facebook/musicgen-small")
        model.set_generation_params(duration=10)
        output = model.generate([prompt])
        music_path = "output_music.wav"
        torchaudio.save(music_path, output[0].cpu(), 32000)
        st.audio(music_path, format="audio/wav")
        st.success("Instrumental ready!")

# Generate Voice with Resemble AI
def generate_voice_with_resemble(lyrics, project_uuid):
    url = f"https://app.resemble.ai/api/v2/projects/{project_uuid}/clips"
    headers = {
        "Authorization": f"Token {RESEMBLE_API_KEY}",
        "Content-Type": "application/json"
    }

    payload = {
        "title": "AI Generated Clip",
        "body": lyrics,
        "is_active": True
    }

    response = requests.post(url, headers=headers, json=payload)

    if response.status_code != 200:
        st.error(f"Resemble API Error: {response.status_code}")
        st.text(response.text)
        return None

    res_data = response.json()
    clip_url = res_data.get("item", {}).get("audio_src")
    return clip_url

voice_path = None
if st.button("Generate Voice with Resemble"):
    if not lyrics:
        st.warning("Please enter lyrics first.")
    else:
        with st.spinner("Generating voice..."):
            audio_url = generate_voice_with_resemble(lyrics, PROJECT_UUID)
            if audio_url:
                response = requests.get(audio_url)
                if response.status_code == 200:
                    with NamedTemporaryFile(delete=False, suffix=".mp3") as f:
                        f.write(response.content)
                        voice_path = f.name
                    st.audio(voice_path, format="audio/mp3")
                    st.success("Voice generation complete!")
                else:
                    st.error("Failed to download audio from Resemble.")
            else:
                st.error("Voice generation failed.")

# Combine Voice and Instrumental
if st.button("Combine Voice and Beat"):
    if not music_path or not voice_path:
        st.warning("Generate both voice and music before combining.")
    else:
        try:
            beat = AudioSegment.from_wav(music_path)
            voice = AudioSegment.from_mp3(voice_path)

            voice = voice - 2
            beat = beat - 4

            combined = beat.overlay(voice)
            combined_path = "combined_output.mp3"
            combined.export(combined_path, format="mp3")

            st.audio(combined_path, format="audio/mp3")
            st.success("Final output ready!")
        except Exception as e:
            st.error(f"Error while c
            ombining: {e}")