import gradio as gr
import requests

# -----------------------------
# CONFIG
# -----------------------------
SUNO_API_BASE = "https://api.sunoapi.org"
import os

#API_KEY = os.getenv("SunoKey")
from dotenv import load_dotenv

# Load environment variables from the .env file
load_dotenv()

# Access the environment variables
suno_api_key = os.getenv('SUNO_API_KEY')

if not API_KEY:
    raise RuntimeError("Missing SunoKey environment variable")

HEADERS = {
    "Authorization": f"Bearer {API_KEY}",
    "Content-Type": "application/json"
}

# -----------------------------
# API HELPERS (minimal, clear)
# -----------------------------

def generate_lyrics(theme, style):
    payload = {
        "prompt": f"Theme: {theme}. Style: {style}."
    }
    r = requests.post(f"{SUNO_API_BASE}/lyrics/generate", json=payload, headers=HEADERS)
    r.raise_for_status()
    return r.json().get("lyrics", "")


def generate_song(prompt, lyrics, model="v5"):
    payload = {
        "model": model,
        "prompt": prompt,
        "lyrics": lyrics
    }
    r = requests.post(f"{SUNO_API_BASE}/music/generate", json=payload, headers=HEADERS)
    r.raise_for_status()
    return r.json()


def separate_vocals(song_id):
    payload = {"song_id": song_id}
    r = requests.post(f"{SUNO_API_BASE}/audio/separate", json=payload, headers=HEADERS)
    r.raise_for_status()
    return r.json()

# -----------------------------
# GRADIO LOGIC
# -----------------------------

def run_pipeline(style, theme, custom_lyrics):
    # Step 1: lyrics
    if not custom_lyrics.strip():
        lyrics = generate_lyrics(theme, style)
    else:
        lyrics = custom_lyrics

    # Step 2: song
    song = generate_song(style, lyrics)
    song_id = song.get("id")
    audio_url = song.get("audio_url")

    return lyrics, audio_url, song_id


def run_separation(song_id):
    result = separate_vocals(song_id)
    return result.get("vocals_url"), result.get("instrumental_url")

# -----------------------------
# UI
# -----------------------------

with gr.Blocks(title="AI Music Instrument – Suno") as app:
    gr.Markdown("# 🎵 AI Music Instrument (Suno API)\nTreat AI like a studio instrument, not a platform.")

    with gr.Row():
        style = gr.Textbox(label="Style / Genre", placeholder="old ballad, lute, choir, dark medieval")
        theme = gr.Textbox(label="Theme", placeholder="fallen king, fate, prophecy")

    lyrics_input = gr.Textbox(label="Custom Lyrics (optional)", lines=6)

    run_btn = gr.Button("Generate Song")

    lyrics_out = gr.Textbox(label="Lyrics Used", lines=8)
    audio_out = gr.Audio(label="Generated Song", type="filepath")
    song_id_out = gr.Textbox(label="Song ID")

    gr.Markdown("## 🎚 Stem Separation")
    sep_btn = gr.Button("Separate Vocals / Instrumental")
    vocals_out = gr.Audio(label="Vocals Stem", type="filepath")
    inst_out = gr.Audio(label="Instrumental Stem", type="filepath")

    run_btn.click(run_pipeline, [style, theme, lyrics_input], [lyrics_out, audio_out, song_id_out])
    sep_btn.click(run_separation, song_id_out, [vocals_out, inst_out])

app.launch()