import os
import time
import requests
from dotenv import load_dotenv
import gradio as gr

# Load .env
load_dotenv()
SUNO_KEY = os.getenv("SunoKey")
if not SUNO_KEY:
    raise ValueError("API key not found! Please set SunoKey in your .env file.")

API_BASE = "https://api.sunoapi.org/api/v1"

def generate_lyrics(prompt):
    """Submit lyrics generation task and poll for results."""
    # Step 1: Submit task
    submit_url = f"{API_BASE}/lyrics"
    headers = {
        "Authorization": f"Bearer {SUNO_KEY}",
        "Content-Type": "application/json"
    }
    data = {"prompt": prompt}
    
    try:
        submit_resp = requests.post(submit_url, json=data, headers=headers)
        submit_resp.raise_for_status()
    except requests.exceptions.HTTPError as e:
        return f"Submit failed: {submit_resp.status_code}\nResponse: {submit_resp.text}\nError: {e}"

    submit_json = submit_resp.json()
    task_id = submit_json.get("data", {}).get("taskId")
    if not task_id:
        return f"Failed to get taskId. Response:\n{submit_json}"

    # Step 2: Poll for lyrics
    details_url = f"{API_BASE}/lyrics/details?taskId={task_id}"
    for i in range(30):  # ~1 minute polling
        time.sleep(2)
        try:
            details_resp = requests.get(details_url, headers=headers)
        except Exception as e:
            return f"Error polling lyrics: {e}"
        
        if details_resp.status_code != 200:
            print(f"Polling attempt {i+1} failed: {details_resp.status_code} - {details_resp.text}")
            continue

        details_json = details_resp.json()
        lyrics_data = details_json.get("data", {}).get("data")
        if lyrics_data:
            results = []
            for idx, variant in enumerate(lyrics_data):
                results.append(f"Variant {idx+1} ({variant.get('title','')}):\n{variant.get('text','')}")
            return "\n\n".join(results)

    return "Lyrics generation timed out. Try again."

# Gradio UI
with gr.Blocks() as demo:
    gr.Markdown("## Suno API Lyrics Generator 🎵")
    prompt_input = gr.Textbox(label="Enter your song prompt", placeholder="A song about a peaceful night in the city")
    output_box = gr.Textbox(label="Generated Lyrics")
    generate_button = gr.Button("Generate Lyrics")
    
    generate_button.click(fn=generate_lyrics, inputs=prompt_input, outputs=output_box)

if __name__ == "__main__":
    demo.launch()