import os
import tempfile
import torch
import gradio as gr
from huggingface_hub import hf_hub_download, snapshot_download
import spaces
import requests
import json
from groq import Groq

# ============================================================
# 🎵 SOMA Music Studio - HeartMuLa Edition
# HeartMuLa-oss-3B + SOMA Multi-Agent + Comic Classic Theme
# ============================================================

# ============================================================
# 🎵 HeartMuLa 모델 다운로드 및 로딩
# ============================================================

def download_models():
    """Download all required model files from HuggingFace Hub."""
    cache_dir = os.environ.get("HF_HOME", os.path.expanduser("/tmp"))
    model_dir = os.path.join(cache_dir, "heartmula_models")

    if not os.path.exists(model_dir):
        os.makedirs(model_dir, exist_ok=True)

    # Download HeartMuLaGen (tokenizer and gen_config)
    print("Downloading HeartMuLaGen files...")
    for filename in ["tokenizer.json", "gen_config.json"]:
        hf_hub_download(
            repo_id="HeartMuLa/HeartMuLaGen",
            filename=filename,
            local_dir=model_dir,
        )

    # Download HeartMuLa-oss-3B
    print("Downloading HeartMuLa-oss-3B...")
    snapshot_download(
        repo_id="HeartMuLa/HeartMuLa-oss-3B",
        local_dir=os.path.join(model_dir, "HeartMuLa-oss-3B"),
    )

    # Download HeartCodec-oss
    print("Downloading HeartCodec-oss...")
    snapshot_download(
        repo_id="HeartMuLa/HeartCodec-oss",
        local_dir=os.path.join(model_dir, "HeartCodec-oss"),
    )

    print("All models downloaded successfully!")
    return model_dir

from heartlib import HeartMuLaGenPipeline

# 모델 다운로드 및 로딩
model_dir = download_models()

# Determine device and dtype
if torch.cuda.is_available():
    device = torch.device("cuda")
    dtype = torch.bfloat16
else:
    device = torch.device("cpu")
    dtype = torch.float32

print(f"Loading HeartMuLa pipeline on {device} with {dtype}...")
pipe = HeartMuLaGenPipeline.from_pretrained(
    model_dir,
    device=device,
    dtype=dtype,
    version="3B",
)
print("HeartMuLa Pipeline loaded successfully!")


# ============================================================
# 🎵 HeartMuLa 가이드 및 설정
# ============================================================

# HeartMuLa 권장 구조 태그 (공식 문서 기반)
STRUCTURE_TAGS = [
    "[Intro]", "[Verse]", "[Prechorus]", "[Chorus]", "[Bridge]",
    "[Interlude]", "[Hook]", "[Outro]", "[Inst]", "[Solo]"
]

# HeartMuLa 스타일 태그 가이드 (데모 페이지 기반)
HEARTMULA_TAG_GUIDE = """
## 🎼 HeartMuLa 입력 형식

### ⚠️ 중요: HeartMuLa는 2가지 입력만 받습니다!
1. **가사 (Lyrics)** - 구조 태그 포함
2. **스타일 태그 (Tags)** - 콤마 구분, 공백 없음

❌ 보컬 타입(듀엣, 그룹 등)은 지원하지 않습니다!

### 📝 구조 태그 (가사 내 사용):
`[Intro]`, `[Verse]`, `[Prechorus]`, `[Chorus]`, `[Bridge]`, `[Outro]`

### 🏷️ 스타일 태그 형식 (콤마 구분, 공백 없음):
```
piano,happy,wedding,synthesizer,romantic
energetic,drums,driving,rock,electric guitar
soft,Sad,Ballad,Longing
```

### 🎹 태그 카테고리 (실제 데모 기반):

**악기 (Instruments):**
Piano, Keyboard, synthesizer, drum machine, electric guitar, acoustic guitar, drums, Strings, bass

**장르 (Genre):**
pop, R&B, rock, Ballad, electronic, acoustic

**분위기/감정 (Mood/Emotion):**
warm, reflection, energetic, soft, Sad, Regret, Longing, meditation, faith, peaceful, Romantic, emotional, hope, hopeful, powerful, epic, driving, happy, uplifting, dreamy

**상황 (Situation):**
Cafe, wedding, Walking, self-discovery

### 🌍 지원 언어:
🇺🇸 English | 🇰🇷 Korean | 🇨🇳 Chinese | 🇯🇵 Japanese | 🇪🇸 Spanish
"""

# HeartMuLa 권장 가사 구조
HEARTMULA_LYRICS_STRUCTURE = """
## 📝 HeartMuLa Recommended Lyrics Structure

### OPTIMAL FORMAT:
```
[Intro]

[Verse]
First verse lyrics here
Second line of first verse

[Prechorus]
Building tension here
Leading to the chorus

[Chorus]
Main hook and memorable melody
Most important part of song

[Verse]
Second verse develops story

[Bridge]
Contrast section here
Different melody or perspective

[Chorus]
Main hook repeated

[Outro]
Closing the song
```

### KEY RULES:
1. [Chorus] appears at least 2-3 times
2. [Prechorus] builds tension before [Chorus]
3. [Bridge] provides contrast before final [Chorus]
4. Use [Prechorus] NOT [Pre Chorus] (no space!)
"""

# 예제 가사
EXAMPLE_LYRICS = """[Intro]

[Verse]
The sun creeps in across the floor
I hear the traffic outside the door
The coffee pot begins to hiss
It is another morning just like this

[Prechorus]
The world keeps spinning round and round
Feet are planted on the ground
I find my rhythm in the sound

[Chorus]
Every day the light returns
Every day the fire burns
We keep on walking down this street
Moving to the same steady beat
It is the ordinary magic that we meet

[Verse]
The hours tick deeply into noon
Chasing shadows, chasing the moon
Work is done and the lights go low
Watching the city start to glow

[Bridge]
It is not always easy, not always bright
Sometimes we wrestle with the night
But we make it to the morning light

[Chorus]
Every day the light returns
Every day the fire burns
We keep on walking down this street
Moving to the same steady beat

[Outro]
Just another day
Every single day"""

EXAMPLE_TAGS = "piano,happy,uplifting,pop"

# 예제 프롬프트 (태그 세트) - HeartMuLa 데모 기반
EXAMPLE_TAG_PRESETS = {
    "☕ Cafe Warm Pop": "warm,reflection,pop,Cafe",
    "🎹 R&B Soft Regret": "R&B,Keyboard,Regret,drum machine,electric guitar,synthesizer,soft",
    "⚡ Electronic Energetic": "energetic,electronic,synthesizer,drum machine,self-discovery",
    "😢 Sad Ballad Longing": "soft,Sad,Ballad,Longing",
    "🧘 Meditation Peaceful": "meditation,faith,acoustic,peaceful",
    "💒 Wedding Romantic": "wedding,Piano,Strings,acoustic guitar,pop,Romantic",
    "💔 Emotional Piano": "Sad,Regret,electric guitar,drums,Piano,emotional",
    "🚶 Hopeful Acoustic": "Walking,acoustic guitar,Strings,hope,hopeful",
    "🎸 Epic Electronic": "synthesizer,drum machine,powerful,epic,Keyboard,electronic",
    "🔥 Driving Rock": "energetic,drums,driving,rock,electric guitar",
    "🎵 Happy Uplifting": "piano,happy,uplifting,pop",
    "🌙 Soft Dreamy": "soft,dreamy,piano,strings,peaceful"
}


# ============================================================
# 🧠 SOMA 에이전트 시스템 - 가사 생성
# ============================================================

LYRICS_AGENTS = {
    "lyricist": f"""You are a master lyricist optimized for HeartMuLa music generation.

{HEARTMULA_LYRICS_STRUCTURE}

Your task: Create powerful, memorable lyrics with OPTIMAL HeartMuLa tag placement.

CRITICAL RULES:
1. Use HeartMuLa structure tags: [Intro], [Verse], [Prechorus], [Chorus], [Bridge], [Interlude], [Hook], [Outro], [Inst], [Solo]
2. [Prechorus] (not [Pre Chorus]) - HeartMuLa format, NO SPACE!
3. Ensure [Chorus] is the most memorable, singable part
4. Include [Prechorus] to build tension before [Chorus]
5. Add [Bridge] for emotional contrast
6. Minimum 6-8 sections for quality generation
7. ⚠️ CRITICAL: Write lyrics in the EXACT language specified by the user!
   - If Korean is specified, write ALL lyrics in Korean (한국어)!
   - Do NOT translate or write in English unless explicitly asked!

Write lyrics that create the BEST POSSIBLE foundation for high-quality music generation.""",

    "producer": f"""You are a music producer specializing in song structure optimization for HeartMuLa.

{HEARTMULA_LYRICS_STRUCTURE}

Your task: Analyze and OPTIMIZE the tag structure for MAXIMUM musical impact.

CRITICAL OPTIMIZATION RULES:
1. Use HeartMuLa tags: [Intro], [Verse], [Prechorus], [Chorus], [Bridge], [Interlude], [Hook], [Outro]
2. IMPORTANT: Use [Prechorus] NOT [Pre Chorus] - no space!
3. Verify the structure follows genre conventions
4. Ensure proper tag sequence (Verse→Prechorus→Chorus)
5. Check that [Chorus] appears at least 2-3 times
6. Verify [Bridge] provides contrast before final chorus

Output the restructured lyrics with OPTIMAL HeartMuLa tag placement.""",

    "emotion_director": f"""You are an emotion director for music production.

Your task: Enhance emotional impact through STRATEGIC tag content.

EMOTIONAL MAPPING BY TAG:
- [Intro]: Intrigue, anticipation
- [Verse]: Storytelling, building connection  
- [Prechorus]: Rising tension, excitement
- [Chorus]: Peak emotion, catharsis
- [Bridge]: Vulnerability, reflection, contrast
- [Outro]: Resolution, lingering feeling

OPTIMIZATION RULES:
1. Each [Verse] should progress emotionally
2. [Chorus] must deliver the strongest emotional punch
3. [Bridge] should offer new emotional perspective
4. [Prechorus] should create anticipation for [Chorus]

Enhance the lyrics for MAXIMUM emotional resonance.""",

    "final_editor": f"""You are the final editor for HeartMuLa music production.

Your task: Output PERFECTLY FORMATTED, production-ready lyrics.

CRITICAL OUTPUT RULES:
1. Output ONLY the actual lyrics with structure tags
2. Use HeartMuLa tags EXACTLY:
   [Intro], [Verse], [Prechorus], [Chorus], [Bridge], [Interlude], [Hook], [Outro], [Inst], [Solo]
3. IMPORTANT: [Prechorus] NOT [Pre Chorus] - no space!
4. DO NOT include English translations in parentheses
5. DO NOT include explanations, descriptions, or markdown
6. DO NOT include lines starting with * or >
7. For [Inst] or [Solo] sections, write ONLY the tag
8. Ensure MINIMUM 6-8 different sections
9. Verify [Chorus] appears at least 2 times
10. PRESERVE THE ORIGINAL LANGUAGE - if lyrics are in Korean, output in Korean!

OUTPUT ONLY CLEAN LYRICS WITH OPTIMAL TAG STRUCTURE IN THE ORIGINAL LANGUAGE."""
}

# 태그 생성 에이전트 - HeartMuLa 데모 태그 기반
TAG_AGENT = f"""You are a style tag generator for HeartMuLa music generation.

HeartMuLa uses comma-separated tags WITHOUT spaces.

VALID TAG CATEGORIES (from official demo):

**Instruments:**
Piano, Keyboard, synthesizer, drum machine, electric guitar, acoustic guitar, drums, Strings, bass

**Genre:**
pop, R&B, rock, Ballad, electronic, acoustic

**Mood/Emotion:**
warm, reflection, energetic, soft, Sad, Regret, Longing, meditation, faith, peaceful, Romantic, emotional, hope, hopeful, powerful, epic, driving, happy, uplifting, dreamy

**Situation:**
Cafe, wedding, Walking, self-discovery

RULES:
1. Output 4-6 tags, comma-separated, NO SPACES between tags
2. Use ONLY tags from the categories above
3. Match the requested genre and mood

OUTPUT FORMAT (examples):
piano,happy,uplifting,pop
warm,reflection,pop,Cafe
energetic,drums,driving,rock,electric guitar
soft,Sad,Ballad,Longing

Output ONLY the comma-separated tags, nothing else."""


# ============================================================
# 🔧 유틸리티 함수
# ============================================================

def call_groq(api_key: str, system: str, user_prompt: str, context: str = "") -> str:
    """Groq API 호출"""
    try:
        client = Groq(api_key=api_key)
        
        messages = [{"role": "system", "content": system}]
        if context:
            messages.append({"role": "user", "content": f"Previous work:\n{context}\n\nTask: {user_prompt}"})
        else:
            messages.append({"role": "user", "content": user_prompt})
        
        completion = client.chat.completions.create(
            model="llama-3.3-70b-versatile",
            messages=messages,
            temperature=0.9,
            max_completion_tokens=4096,
            top_p=1,
            stream=False
        )
        
        if completion and completion.choices and completion.choices[0].message:
            return completion.choices[0].message.content or "Error: Empty response"
        return "Error: Invalid response"
        
    except Exception as e:
        return f"Error: {str(e)}"


def clean_lyrics(text: str) -> str:
    """가사 후처리 - HeartMuLa 포맷 최적화"""
    import re
    
    if not text or not isinstance(text, str):
        return ""
    
    lines = text.split('\n')
    cleaned_lines = []
    
    for line in lines:
        if not line.strip():
            cleaned_lines.append('')
            continue
        
        # 스킵 패턴
        skip_patterns = [r'^\s*\*', r'^\s*>', r'^\s*---', r'^\s*###', r'^\s*\*\*.*\*\*\s*$']
        if any(re.match(p, line) for p in skip_patterns):
            continue
        
        # 정리
        line = re.sub(r'\s*\([A-Za-z].*?\)\s*$', '', line)
        line = re.sub(r'\*\*(.*?)\*\*', r'\1', line)
        
        # HeartMuLa 태그 정규화
        line = re.sub(r'\[Pre.?Chorus\]', '[Prechorus]', line, flags=re.IGNORECASE)
        line = re.sub(r'\[Post.?Chorus\]', '[Chorus]', line, flags=re.IGNORECASE)
        line = re.sub(r'\[Build.?Up\]', '[Prechorus]', line, flags=re.IGNORECASE)
        line = re.sub(r'\[Break\]', '[Interlude]', line, flags=re.IGNORECASE)
        
        if line.strip():
            cleaned_lines.append(line.strip())
    
    result = '\n'.join(cleaned_lines)
    result = re.sub(r'\n{3,}', '\n\n', result)
    
    return result.strip()


def clean_tags(tags: str) -> str:
    """태그 정리 - HeartMuLa 포맷 (콤마 구분, 공백 없음)"""
    if not tags:
        return "piano,happy,pop"
    
    tags = tags.lower().strip()
    tags = tags.replace(', ', ',').replace(' ,', ',').replace('  ', ' ')
    tags = tags.replace(' ', '_').replace(',,', ',')
    tags = tags.strip(',')
    
    tag_list = [t.strip() for t in tags.split(',') if t.strip()]
    unique_tags = list(dict.fromkeys(tag_list))
    
    return ','.join(unique_tags) if unique_tags else "piano,happy,pop"


# ============================================================
# 🎵 메인 함수들
# ============================================================

@spaces.GPU(duration=120)
def generate_lyrics_soma(
    api_key: str, theme: str, genre: str, mood: str, 
    language: str, additional: str, 
    progress=gr.Progress()
):
    """SOMA 가사 생성"""
    if not api_key or not api_key.strip():
        return "❌ Groq API Key 필요", "", "", "", ""
    if not theme or not theme.strip():
        return "❌ 주제를 입력하세요", "", "", "", ""
    
    base_prompt = f"""Create PROFESSIONAL lyrics optimized for HeartMuLa:
- Theme: {theme}
- Genre: {genre}
- Mood: {mood}
- Language: {language} ← CRITICAL: You MUST write ALL lyrics in {language} language!
{f'- Additional: {additional}' if additional else ''}

⚠️ LANGUAGE REQUIREMENT: Write the ENTIRE lyrics in {language}!
{"- 반드시 한국어로 가사를 작성하세요! 한글로만 작성!" if language == "Korean" else ""}
{"- 请用中文写所有歌词！必须全部用中文！" if language == "Chinese" else ""}
{"- 日本語で歌詞を書いてください！全て日本語で！" if language == "Japanese" else ""}
{"- ¡Escribe toda la letra en español! ¡Solo en español!" if language == "Spanish" else ""}

USE HeartMuLa STRUCTURE TAGS:
[Intro], [Verse], [Prechorus], [Chorus], [Bridge], [Interlude], [Hook], [Outro]

IMPORTANT: Use [Prechorus] NOT [Pre Chorus] - no space!

REQUIRED STRUCTURE:
1. [Intro] - Set the mood
2. [Verse] x2-3 - Tell the story
3. [Prechorus] - Build tension
4. [Chorus] x2-3 - Main hook (MOST important!)
5. [Bridge] - Emotional contrast
6. [Outro] - Conclusion

REMEMBER: ALL lyrics must be in {language}!"""

    try:
        progress(0.2, desc="🎤 작사가 - 초안 작성...")
        draft = call_groq(api_key, LYRICS_AGENTS["lyricist"], base_prompt)
        if draft.startswith("Error:"):
            return f"❌ {draft}", "", "", "", ""
        
        progress(0.4, desc="🎹 프로듀서 - 구조 최적화...")
        structured = call_groq(api_key, LYRICS_AGENTS["producer"], 
                              f"Optimize for {genre}. Use [Prechorus] not [Pre Chorus]! KEEP ALL LYRICS IN {language}!", draft)
        if structured.startswith("Error:"):
            return f"❌ {structured}", draft, "", "", ""
        
        progress(0.6, desc="💫 감성 디렉터 - 감정 강화...")
        emotional = call_groq(api_key, LYRICS_AGENTS["emotion_director"],
                             f"Enhance for {mood}. KEEP ALL LYRICS IN {language}!", structured)
        if emotional.startswith("Error:"):
            return f"❌ {emotional}", draft, structured, "", ""
        
        progress(0.8, desc="✨ 최종 편집...")
        lang_instruction = {
            "Korean": "한국어로만 출력하세요!",
            "Chinese": "用中文输出！",
            "Japanese": "日本語で出力してください！",
            "Spanish": "¡Escribe en español!"
        }.get(language, "")
        final = call_groq(api_key, LYRICS_AGENTS["final_editor"],
                         f"Output ONLY clean lyrics in {language}. Use [Prechorus] not [Pre Chorus]! {lang_instruction}", emotional)
        if final.startswith("Error:"):
            return f"❌ {final}", draft, structured, emotional, ""
        
        final_cleaned = clean_lyrics(final)
        
        progress(1.0, desc="✅ 완료!")
        return "✅ 가사 생성 완료!", draft, structured, emotional, final_cleaned
        
    except Exception as e:
        return f"❌ 예외: {str(e)}", "", "", "", ""


@spaces.GPU(duration=60)
def quick_lyrics(api_key: str, theme: str, genre: str, mood: str, language: str, additional: str):
    """빠른 가사 생성"""
    if not api_key or not api_key.strip():
        return "❌ API Key 필요"
    if not theme or not theme.strip():
        return "❌ 주제를 입력하세요"
    
    prompt = f"""Create song lyrics for HeartMuLa:
- Theme: {theme}
- Genre: {genre}
- Mood: {mood}
- Language: {language} ← CRITICAL: Write ALL lyrics in {language}!
{f'- Special: {additional}' if additional else ''}

⚠️ LANGUAGE: You MUST write the ENTIRE lyrics in {language}!
{"반드시 한국어로 작성하세요! 한글 가사만!" if language == "Korean" else ""}
{"请用中文写歌词！" if language == "Chinese" else ""}
{"日本語で書いてください！" if language == "Japanese" else ""}
{"¡Escribe en español!" if language == "Spanish" else ""}

USE STRUCTURE (8-10 sections):
[Intro] → [Verse] → [Prechorus] → [Chorus] → [Verse] → [Prechorus] → [Chorus] → [Bridge] → [Chorus] → [Outro]

IMPORTANT: Use [Prechorus] NOT [Pre Chorus]!

OUTPUT ONLY lyrics with tags in {language}. NO explanations."""

    try:
        result = call_groq(api_key, f"""Professional songwriter for HeartMuLa.
{HEARTMULA_LYRICS_STRUCTURE}
Output ONLY clean lyrics.""", prompt)
        
        if result.startswith("Error:"):
            return f"❌ {result}"
        
        return clean_lyrics(result)
    except Exception as e:
        return f"❌ {str(e)}"


@spaces.GPU(duration=60)
def generate_tags_ai(api_key: str, genre: str, mood: str, current_tags: str):
    """AI 태그 생성 - HeartMuLa 형식"""
    if not api_key:
        return "piano,happy,pop"
    
    prompt = f"""Generate HeartMuLa style tags based on:
- Genre: {genre}
- Mood: {mood}
- Current hint: {current_tags}

Use ONLY these tag categories (from HeartMuLa demo):
- Instruments: Piano, Keyboard, synthesizer, drum machine, electric guitar, acoustic guitar, drums, Strings, bass
- Genre: pop, R&B, rock, Ballad, electronic, acoustic
- Mood: warm, reflection, energetic, soft, Sad, Regret, Longing, meditation, faith, peaceful, Romantic, emotional, hope, hopeful, powerful, epic, driving, happy, uplifting, dreamy
- Situation: Cafe, wedding, Walking, self-discovery

Output 4-6 comma-separated tags WITHOUT spaces between tags.
Example: piano,happy,uplifting,pop,Romantic"""

    try:
        result = call_groq(api_key, TAG_AGENT, prompt)
        if result.startswith("Error:"):
            return "piano,happy,pop"
        return clean_tags(result)
    except:
        return "piano,happy,pop"


@spaces.GPU(duration=180)
def generate_music_heartmula(
    lyrics: str,
    tags: str,
    max_duration_seconds: int,
    temperature: float,
    topk: int,
    cfg_scale: float,
    progress=gr.Progress(track_tqdm=True),
):
    """HeartMuLa로 음악 생성"""
    if not lyrics or not lyrics.strip():
        raise gr.Error("⚠️ 가사를 입력하세요!")

    if not tags or not tags.strip():
        raise gr.Error("⚠️ 태그를 입력하세요!")

    # 태그 정리
    tags = clean_tags(tags)
    
    # 임시 파일 생성
    with tempfile.NamedTemporaryFile(suffix=".mp3", delete=False) as f:
        output_path = f.name

    max_audio_length_ms = max_duration_seconds * 1000

    try:
        with torch.no_grad():
            pipe(
                {
                    "lyrics": lyrics,
                    "tags": tags,
                },
                max_audio_length_ms=max_audio_length_ms,
                save_path=output_path,
                topk=topk,
                temperature=temperature,
                cfg_scale=cfg_scale,
            )
        
        return output_path, f"✅ 음악 생성 완료! ({max_duration_seconds}초, 태그: {tags})"
    
    except Exception as e:
        return None, f"❌ 생성 실패: {str(e)}"


def load_tag_preset(preset_name):
    """태그 프리셋 로드"""
    if preset_name in EXAMPLE_TAG_PRESETS:
        return EXAMPLE_TAG_PRESETS[preset_name]
    return ""


# ============================================================
# 🎨 Comic Classic Theme CSS
# ============================================================

css = """
@import url('https://fonts.googleapis.com/css2?family=Bangers&family=Comic+Neue:wght@400;700&display=swap');

.gradio-container {
    background-color: #FEF9C3 !important;
    background-image: radial-gradient(#1F2937 1px, transparent 1px) !important;
    background-size: 20px 20px !important;
    font-family: 'Comic Neue', cursive, sans-serif !important;
}

.huggingface-space-header, footer, .footer { display: none !important; }

.header-title h1 {
    font-family: 'Bangers', cursive !important;
    color: #1F2937 !important;
    font-size: 3rem !important;
    text-align: center !important;
    text-shadow: 4px 4px 0px #FACC15, 6px 6px 0px #1F2937 !important;
    letter-spacing: 3px !important;
    -webkit-text-stroke: 2px #1F2937 !important;
}

.section-title {
    font-family: 'Bangers', cursive !important;
    color: #1F2937 !important;
    font-size: 1.6rem !important;
    border-bottom: 4px solid #3B82F6 !important;
    padding-bottom: 8px !important;
    text-shadow: 2px 2px 0px #FACC15 !important;
}

.gr-panel, .gr-box, .block {
    background: #FFFFFF !important;
    border: 3px solid #1F2937 !important;
    border-radius: 8px !important;
    box-shadow: 6px 6px 0px #1F2937 !important;
}

textarea, input[type="text"], input[type="password"] {
    background: #FFFFFF !important;
    border: 3px solid #1F2937 !important;
    border-radius: 8px !important;
    font-family: 'Comic Neue', cursive !important;
    font-weight: 700 !important;
}

textarea:focus, input:focus {
    border-color: #3B82F6 !important;
    box-shadow: 4px 4px 0px #3B82F6 !important;
}

.gr-button-primary, button.primary {
    background: #3B82F6 !important;
    border: 3px solid #1F2937 !important;
    border-radius: 8px !important;
    color: #FFFFFF !important;
    font-family: 'Bangers', cursive !important;
    font-size: 1.2rem !important;
    letter-spacing: 2px !important;
    box-shadow: 5px 5px 0px #1F2937 !important;
    text-shadow: 1px 1px 0px #1F2937 !important;
}

.gr-button-primary:hover { 
    background: #2563EB !important;
}

.gr-button-secondary, button.secondary {
    background: #EF4444 !important;
    border: 3px solid #1F2937 !important;
    color: #FFFFFF !important;
    font-family: 'Bangers', cursive !important;
    box-shadow: 4px 4px 0px #1F2937 !important;
}

.generate-btn {
    background: #10B981 !important;
    border: 3px solid #1F2937 !important;
    color: #FFFFFF !important;
    font-family: 'Bangers', cursive !important;
    font-size: 1.4rem !important;
    box-shadow: 5px 5px 0px #1F2937 !important;
}

.generate-btn:hover {
    background: #059669 !important;
}

.gr-accordion {
    background: #FACC15 !important;
    border: 3px solid #1F2937 !important;
    box-shadow: 4px 4px 0px #1F2937 !important;
}

.tag-input textarea {
    background: #FEF3C7 !important;
    border: 3px dashed #F59E0B !important;
    font-family: 'Courier New', monospace !important;
}

.status-box textarea {
    background: #1F2937 !important;
    color: #10B981 !important;
    font-family: 'Courier New', monospace !important;
    border: 3px solid #10B981 !important;
}

label {
    color: #1F2937 !important;
    font-family: 'Comic Neue', cursive !important;
    font-weight: 700 !important;
}

::-webkit-scrollbar { width: 12px; }
::-webkit-scrollbar-track { background: #FEF9C3; border: 2px solid #1F2937; }
::-webkit-scrollbar-thumb { background: #3B82F6; border: 2px solid #1F2937; }
::selection { background: #FACC15; color: #1F2937; }
"""


# ============================================================
# 🎨 Gradio UI
# ============================================================

with gr.Blocks(css=css, title="🎵 SOMA Music Studio - HeartMuLa") as demo:
    
    # Header - HOME Badge
    gr.HTML("""
        <div style="text-align: center; margin: 20px 0 10px 0;">
            <a href="https://www.humangen.ai" target="_blank" style="text-decoration: none;">
                <img src="https://img.shields.io/static/v1?label=🏠 HOME&message=HUMANGEN.AI&color=0000ff&labelColor=ffcc00&style=for-the-badge" alt="HOME">
            </a>
        </div>
    """)
    
    gr.Markdown("# 🎵 SOMA MUSIC STUDIO 🎶", elem_classes="header-title")
    
    gr.Markdown("""
    <p style="text-align: center; font-family: 'Comic Neue', cursive; font-size: 1.1rem; font-weight: 700; color: #1F2937;">
        💫 HeartMuLa-oss-3B + SOMA Multi-Agent = 최고 품질 AI 음악 생성 💫
    </p>
    """)
    
    # API Key
    with gr.Accordion("🔑 Groq API Key (가사 생성용)", open=False):
        GROQ_KEY = os.environ.get("GROQ_API_KEY", "")
        groq_key = gr.Textbox(
            label="Groq API Key",
            type="password",
            value=GROQ_KEY,
            placeholder="gsk_..." if not GROQ_KEY else "✅ Secret 로드됨",
            interactive=not bool(GROQ_KEY)
        )
    
    with gr.Row(equal_height=False):
        # ========== 좌측: 가사 생성 ==========
        with gr.Column(scale=1, min_width=400):
            gr.Markdown("## 📝 LYRICS GENERATOR", elem_classes="section-title")
            
            theme_input = gr.Textbox(
                label="🎯 노래 주제",
                placeholder="예: 이별 후 성장, 꿈을 향한 도전, 사랑의 고백...",
                lines=2
            )
            
            with gr.Row():
                lyrics_genre = gr.Dropdown(
                    label="🎸 장르",
                    choices=["K-Pop", "Pop", "R&B", "Hip-Hop", "Ballad", "Rock", 
                            "EDM", "Jazz", "Folk", "Disco", "Cinematic"],
                    value="K-Pop"
                )
                lyrics_mood = gr.Dropdown(
                    label="💫 분위기",
                    choices=["Empowering", "Melancholic", "Joyful", "Romantic",
                            "Energetic", "Dreamy", "Nostalgic", "Peaceful", "Confident"],
                    value="Empowering"
                )
            
            with gr.Row():
                lyrics_language = gr.Dropdown(
                    label="🌍 언어 (Language)",
                    choices=["English", "Korean", "Chinese", "Japanese", "Spanish"],
                    value="English",
                    info="HeartMuLa 공식 지원 5개 언어"
                )
            
            lyrics_additional = gr.Textbox(
                label="✨ 추가 지시 (선택)",
                placeholder="특별 요청: 특정 감정, 스타일 등...",
                lines=1
            )
            
            with gr.Row():
                quick_btn = gr.Button("⚡ QUICK", variant="secondary")
                soma_lyrics_btn = gr.Button("🧠 SOMA GENERATE", variant="primary")
            
            lyrics_status = gr.Textbox(label="📊 상태", interactive=False, max_lines=1, elem_classes="status-box")
            
            with gr.Accordion("🔍 SOMA 작업 과정", open=False):
                step1_out = gr.Textbox(label="1️⃣ 작사가", lines=3, interactive=False)
                step2_out = gr.Textbox(label="2️⃣ 프로듀서", lines=3, interactive=False)
                step3_out = gr.Textbox(label="3️⃣ 감성", lines=3, interactive=False)
                step4_out = gr.Textbox(label="4️⃣ 최종", lines=3, interactive=False)
            
            final_lyrics = gr.Textbox(
                label="✏️ 가사 (편집 가능)",
                lines=16,
                value=EXAMPLE_LYRICS,
                placeholder="가사가 여기 표시됩니다..."
            )
        
        # ========== 우측: 음악 생성 ==========
        with gr.Column(scale=1, min_width=400):
            gr.Markdown("## 🎵 MUSIC GENERATOR", elem_classes="section-title")
            
            # 태그 프리셋
            gr.Markdown("### 🏷️ 태그 프리셋 (클릭하면 자동 입력)")
            tag_preset = gr.Dropdown(
                label="프리셋 선택",
                choices=list(EXAMPLE_TAG_PRESETS.keys()),
                value=None
            )
            
            tags_input = gr.Textbox(
                label="🎹 스타일 태그 (콤마 구분, 공백 없음)",
                value=EXAMPLE_TAGS,
                placeholder="piano,happy,uplifting,pop",
                lines=1,
                elem_classes="tag-input"
            )
            
            with gr.Row():
                ai_tag_btn = gr.Button("🤖 AI 태그 생성", variant="secondary", size="sm")
            
            with gr.Accordion("⚙️ 생성 설정", open=True):
                max_duration = gr.Slider(
                    minimum=30, maximum=240, value=120, step=10,
                    label="⏱️ 최대 길이 (초)",
                    info="30~240초"
                )
                
                with gr.Row():
                    temperature = gr.Slider(
                        minimum=0.1, maximum=2.0, value=1.0, step=0.1,
                        label="🌡️ Temperature",
                        info="높을수록 창의적"
                    )
                    topk = gr.Slider(
                        minimum=1, maximum=100, value=50, step=1,
                        label="🎯 Top-K"
                    )
                
                cfg_scale = gr.Slider(
                    minimum=1.0, maximum=3.0, value=1.5, step=0.1,
                    label="📐 CFG Scale",
                    info="Classifier-free guidance"
                )
            
            generate_btn = gr.Button(
                "🎶 GENERATE MUSIC!", 
                variant="primary", 
                size="lg",
                elem_classes="generate-btn"
            )
            
            music_status = gr.Textbox(label="📊 상태", interactive=False, max_lines=1, elem_classes="status-box")
            audio_output = gr.Audio(label="🎧 생성된 음악", type="filepath")
            
            gr.Markdown("""
            ### 💡 Tips
            - **구조 태그:** [Intro], [Verse], [Prechorus], [Chorus], [Bridge], [Outro]
            - **스타일 태그:** piano,happy,uplifting,pop (콤마 구분, 공백 없음!)
            - **지원 언어:** 🇺🇸 EN | 🇰🇷 KO | 🇨🇳 ZH | 🇯🇵 JA | 🇪🇸 ES
            - ⚠️ **보컬 타입(듀엣, 그룹 등)은 지원하지 않습니다!**
            - **Temperature:** 높으면 창의적, 낮으면 일관성
            """)
    
    # 가이드
    with gr.Accordion("📖 HeartMuLa 가이드", open=False):
        gr.Markdown(f"""
## ⚠️ HeartMuLa 입력 형식 (중요!)

HeartMuLa는 **2가지 입력만** 받습니다:
1. **가사 (Lyrics)** - 구조 태그 포함
2. **스타일 태그 (Tags)** - 콤마 구분

❌ **지원하지 않는 것:**
- 보컬 타입 선택 (듀엣, 그룹, 코러스 등)
- 레퍼런스 오디오 (아직 미지원)
- BPM/템포 직접 지정

---

## 🌍 지원 언어 (5개)
| 언어 | 코드 | 상태 |
|------|------|------|
| 🇺🇸 English | EN | ✅ |
| 🇰🇷 Korean | KO | ✅ |
| 🇨🇳 Chinese | ZH | ✅ |
| 🇯🇵 Japanese | JA | ✅ |
| 🇪🇸 Spanish | ES | ✅ |

---

{HEARTMULA_TAG_GUIDE}

---

{HEARTMULA_LYRICS_STRUCTURE}

---

**Model:** [HeartMuLa-oss-3B](https://huggingface.co/HeartMuLa/HeartMuLa-oss-3B) | 
**Paper:** [arXiv](https://arxiv.org/abs/2601.10547) | 
**Code:** [GitHub](https://github.com/HeartMuLa/heartlib)

*Licensed under CC BY-NC 4.0 (Non-Commercial)*
        """)
    
    # ========== Event Handlers ==========
    
    # 태그 프리셋 로드
    tag_preset.change(
        fn=load_tag_preset,
        inputs=[tag_preset],
        outputs=[tags_input]
    )
    
    # AI 태그 생성
    ai_tag_btn.click(
        fn=generate_tags_ai,
        inputs=[groq_key, lyrics_genre, lyrics_mood, tags_input],
        outputs=[tags_input]
    )
    
    # 빠른 가사 생성
    quick_btn.click(
        fn=quick_lyrics,
        inputs=[groq_key, theme_input, lyrics_genre, lyrics_mood, lyrics_language, lyrics_additional],
        outputs=[final_lyrics]
    )
    
    # SOMA 가사 생성
    soma_lyrics_btn.click(
        fn=generate_lyrics_soma,
        inputs=[groq_key, theme_input, lyrics_genre, lyrics_mood, lyrics_language, lyrics_additional],
        outputs=[lyrics_status, step1_out, step2_out, step3_out, step4_out]
    ).then(
        fn=lambda x: x,
        inputs=[step4_out],
        outputs=[final_lyrics]
    )
    
    # 음악 생성 (HeartMuLa)
    generate_btn.click(
        fn=generate_music_heartmula,
        inputs=[final_lyrics, tags_input, max_duration, temperature, topk, cfg_scale],
        outputs=[audio_output, music_status]
    )


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