import gradio as gr
import requests
import edge_tts
import os
import re
import scipy.io.wavfile
from transformers import pipeline
# ★ GPU 지원을 위한 라이브러리 (Hugging Face 환경인 경우)
try:
    import spaces
    has_spaces = True
except:
    has_spaces = False

GENAI_KEY = os.getenv("GEMINI_KEY")

# 모델 로드 (Hugging Face 최적화)
try:
    music_synthesiser = pipeline("text-to-audio", "facebook/musicgen-small", device=0 if has_spaces else -1)
except:
    music_synthesiser = None

# 8인 멤버 및 데이터 (삭제/축소 없음)
MEMBERS = { "서윤 (Korea)": "ko-KR-SunHiNeural", "Chloe (USA)": "en-US-AriaNeural", "Naomi (Japan)": "ja-JP-NanamiNeural", "Beatrice (Brazil)": "pt-BR-FranciscaNeural", "Elena (Spain)": "es-ES-ElviraNeural", "Amira (Egypt)": "ar-EG-SalmaNeural", "Liwei (China)": "zh-CN-XiaoxiaoNeural", "Sophie (France)": "fr-FR-DeniseNeural" }

# ★ GPU 가속 데코레이터 (Hugging Face 무료 GPU 활성화)
def generate_music_safe(prompt):
    if not music_synthesiser: return None, None
    try:
        # 리더님 요청대로 512(약 12초) 풀 파워!
        output = music_synthesiser(prompt, forward_params={"max_new_tokens": 512, "do_sample": True})
        return output["sampling_rate"], output["audio"][0].T
    except:
        return None, None

async def band_consulting(user_input, member_name, lang_code):
    try:
        lang_map = {"ko":"Korean","en":"English","ja":"Japanese","pt":"Portuguese","es":"Spanish","ar":"Arabic","zh":"Chinese","fr":"French"}
        target_lang = lang_map.get(lang_code, "Korean")
        
        system_instruction = f"당신은 {member_name}입니다. 반드시 {target_lang}로만 답변하세요. 음성은 5줄 이내, 상세 설명은 [TAB]에, 음악 프롬프트는 [MUSIC: 프롬프트]에 넣으세요."
        
        url = f"https://generativelanguage.googleapis.com/v1beta/models/gemini-1.5-flash:generateContent?key={GENAI_KEY}"
        payload = {"contents": [{"parts": [{"text": f"{system_instruction}\n질문: {user_input}"}]}]}
        
        res = requests.post(url, json=payload).json()
        ai_text_raw = res['candidates'][0]['content']['parts'][0]['text']
        
        tab_match = re.search(r'\[TAB\](.*?)(\[|$)', ai_text_raw, re.DOTALL)
        music_match = re.search(r'\[MUSIC:(.*?)\]', ai_text_raw, re.IGNORECASE)
        tab_display = tab_match.group(1).strip() if tab_match else "상세 악보 없음"
        
        clean_text = re.sub(r'\[TAB\].*?(\[|$)', '', ai_text_raw, flags=re.DOTALL)
        clean_text = re.sub(r'\[MUSIC:.*?\]', '', clean_text).strip()
        tts_final = "\n".join(clean_text.split('\n')[:5])

        # 음성 생성
        voice_path = f"/tmp/v_{member_name.split()[0]}.mp3"
        await edge_tts.Communicate(re.sub(r'[\*\#\-\_\~\|]', '', tts_final), MEMBERS.get(member_name, "ko-KR-SunHiNeural")).save(voice_path)
        
        # 음악 생성 (타임아웃 방지 로직 포함)
        music_path = None
        if music_match:
            sr, audio = generate_music_safe(music_match.group(1).strip())
            if sr is not None:
                music_path = f"/tmp/m_{member_name.split()[0]}.wav"
                scipy.io.wavfile.write(music_path, sr, audio)

        return tts_final, voice_path, music_path, tab_display
    except Exception as e:
        return f"System Error: {str(e)}", None, None, "Error"

with gr.Blocks() as demo:
    i1 = gr.Textbox(); i2 = gr.Textbox(); i3 = gr.Textbox()
    o1 = gr.Textbox(); o2 = gr.Audio(); o3 = gr.Audio(); o4 = gr.Textbox()
    # ★ 큐(Queue)를 활성화하여 서버가 끊기지 않게 함
    btn = gr.Button("GO"); btn.click(band_consulting, [i1, i2, i3], [o1, o2, o3, o4], api_name="predict")

# ★ queue()를 반드시 붙여야 긴 작업을 브라우저가 기다려줍니다.
demo.queue().launch()