import os
import sys
import re
import gradio as gr

# ── Setup: add ACE-Step-1.5 to sys.path so acestep.* imports work ──
ACE_ROOT = os.path.join(os.path.dirname(__file__), "ACE-Step-1.5")
sys.path.insert(0, ACE_ROOT)
os.environ["ACESTEP_GENERATION_TIMEOUT"] = "1200"

from acestep.handler import AceStepHandler
from acestep.inference import GenerationParams, GenerationConfig, generate_music
# ── Globals ──
DIT_HANDLER = None

def initialize_model():
    global DIT_HANDLER
    print("[System] Loading ACE-Step 1.5 model...")
    DIT_HANDLER = AceStepHandler()
    status_msg, success = DIT_HANDLER.initialize_service(
        project_root=ACE_ROOT,
        config_path="acestep-v15-turbo",
        device="auto",
        offload_to_cpu=False,
    )
    if not success:
        raise RuntimeError(f"Model initialization failed: {status_msg}")
    print("[System] Model loaded successfully!")

# Load model on startup (first request will wait if not ready)
try:
    initialize_model()
except Exception as e:
    print(f"[ERROR] Failed to load model at startup: {e}")
    print("[INFO] Model will be loaded on first request instead")

# ── Groq LLM ──
def generate_lyrics_and_tags(story, genre, mood, tempo, vocals_enabled):
    from groq import Groq

    api_key = os.environ.get("GROQ_API_KEY")
    if not api_key:
        raise ValueError("GROQ_API_KEY not set. Add it in HF Space Secrets.")

    client = Groq(api_key=api_key)

    system_prompt = """You are an expert music producer and songwriter. Based on a user's story and preferences, create compelling song lyrics and musical style tags."""

    user_prompt = f"""Story: {story}

Musical Preferences:
- Genre: {genre}
- Mood: {mood}
- Tempo: {tempo}
- Vocals: {'Yes - write lyrics to be sung' if vocals_enabled else 'Instrumental only'}

Generate:
1. Full song lyrics with [Verse], [Chorus], [Bridge] sections.
2. Musical style tags: genre, BPM, key, instruments, production style (comma-separated).

Format:
=== LYRICS ===
[lyrics]
=== STYLE TAGS ===
[comma-separated style tags]"""

    response = client.chat.completions.create(
        model="llama-3.3-70b-versatile",
        messages=[
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_prompt},
        ],
        temperature=0.7,
        max_tokens=2000,
    )

    content = response.choices[0].message.content

    lyrics = ""
    style_tags = "pop, 120bpm"

    lyrics_match = re.search(r'=== LYRICS ===\s*(.*?)(?:=== STYLE TAGS ===|$)', content, re.DOTALL)
    if lyrics_match:
        lyrics = lyrics_match.group(1).strip()

    tags_match = re.search(r'=== STYLE TAGS ===\s*(.*?)$', content, re.DOTALL)
    if tags_match:
        style_tags = tags_match.group(1).strip()

    return lyrics, style_tags

# ── ACE-Step Generation ──
def generate_audio(lyrics, style_tags):
    global DIT_HANDLER

    if DIT_HANDLER is None:
        print("[System] Loading model on demand...")
        initialize_model()

    save_dir = os.path.join(os.path.dirname(__file__), "generated_songs")
    os.makedirs(save_dir, exist_ok=True)

    params = GenerationParams(
        task_type="text2music",
        thinking=False,
        caption=style_tags,
        lyrics=lyrics,
        duration=60,
        inference_steps=8,
    )
    gen_config = GenerationConfig(
        batch_size=1,
        audio_format="wav",
    )

    result = generate_music(
        DIT_HANDLER, None, params=params, config=gen_config, save_dir=save_dir
    )

    if result.success and result.audios:
        return result.audios[0]["path"]
    else:
        raise RuntimeError(f"Generation failed: {result.status_message}")

# ── Gradio Interface ──
def generate_song(story, genre, mood, tempo, vocals_enabled, progress=gr.Progress()):
    try:
        progress(0.1, desc="Analyzing story and generating lyrics...")
        lyrics, style_tags = generate_lyrics_and_tags(
            story, genre, mood, tempo, vocals_enabled
        )

        print(f"\n--- Generated Lyrics ---\n{lyrics}")
        print(f"\n--- Style Tags ---\n{style_tags}")

        progress(0.3, desc="Creating music (1-3 minutes on GPU)...")
        audio_path = generate_audio(lyrics, style_tags)

        progress(1.0, desc="Complete!")
        return audio_path, lyrics, style_tags
    except Exception as e:
        raise gr.Error(str(e))

DESCRIPTION = """
# 🎵 AI Song Generator
Tell your story and get a custom song generated by AI in under 3 minutes.
"""

with gr.Blocks(title="AI Song Generator", theme=gr.themes.Soft()) as demo:
    gr.Markdown(DESCRIPTION)

    with gr.Row():
        with gr.Column(scale=2):
            story = gr.Textbox(
                label="Your Story",
                placeholder="Tell me about your experience, moment, or idea for a song...",
                lines=5,
            )

            with gr.Row():
                genre = gr.Dropdown(
                    label="Genre",
                    choices=["Pop", "Rock", "Hip-Hop", "Electronic", "R&B", "Lo-fi", "Jazz", "Classical", "Folk", "Indie"],
                    value="Pop",
                )
                mood = gr.Dropdown(
                    label="Mood",
                    choices=["Happy", "Energetic", "Triumphant", "Romantic", "Dreamy", "Melancholic", "Calm", "Dark"],
                    value="Happy",
                )

            with gr.Row():
                tempo = gr.Dropdown(
                    label="Tempo",
                    choices=["Slow (60-80 BPM)", "Medium (90-120 BPM)", "Fast (130-160 BPM)"],
                    value="Medium (90-120 BPM)",
                )
                vocals = gr.Checkbox(
                    label="Include vocals",
                    value=True,
                )

            generate_btn = gr.Button("🎵 Generate Song", variant="primary", size="lg")

        with gr.Column(scale=2):
            audio_output = gr.Audio(label="Your Song", type="filepath")

            with gr.Accordion("Show Lyrics", open=False):
                lyrics_output = gr.Textbox(label="Generated Lyrics", lines=10)

            with gr.Accordion("Show Style Tags", open=False):
                tags_output = gr.Textbox(label="Style Tags", lines=3)

    generate_btn.click(
        fn=generate_song,
        inputs=[story, genre, mood, tempo, vocals],
        outputs=[audio_output, lyrics_output, tags_output],
    )

    gr.Markdown("---\n⚡ Powered by ACE-Step 1.5 Turbo · 🧠 Groq · Built with Gradio")

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