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

# Load .env
load_dotenv()
SUNO_KEY = os.getenv("SunoKey")

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

def generate_lyrics(prompt):
    # Step 1: Submit lyrics generation task
    submit_url = f"{API_BASE}/lyrics"
    headers = {
        "Authorization": f"Bearer {SUNO_KEY}",
        "Content-Type": "application/json"
    }
    data = {
        "prompt": prompt
    }

    submit_resp = requests.post(submit_url, json=data, headers=headers)
    if submit_resp.status_code != 200:
        return f"Submit failed: {submit_resp.status_code} {submit_resp.text}"

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

    # Step 2: Poll for lyrics
    details_url = f"{API_BASE}/lyrics/details?taskId={task_id}"
    for _ in range(30):  # poll up to ~30 times (~30*2 sec = 1 min)
        time.sleep(2)
        details_resp = requests.get(details_url, headers=headers)
        if details_resp.status_code != 200:
            continue
        details_json = details_resp.json()
        if details_json.get("data", {}).get("data"):
            lyrics_variants = details_json["data"]["data"]
            results = []
            for i, variant in enumerate(lyrics_variants):
                results.append(f"Variant {i+1}:\n{variant.get('text', '')}")
            return "\n\n".join(results)

    return "Lyrics generation timed out. Try again."

# Gradio interface
with gr.Blocks() as demo:
    gr.Markdown("## Suno Lyrics Generator")
    prompt_input = gr.Textbox(label="Enter your song prompt")
    output_box = gr.Textbox(label="Generated Lyrics", lines=15)
    generate_btn = gr.Button("Generate Lyrics")
    
    generate_btn.click(generate_lyrics, inputs=prompt_input, outputs=output_box)

demo.launch()