import gradio as gr
import requests
import edge_tts
import os
import re
import scipy.io.wavfile
from transformers import pipeline

GENAI_KEY = os.getenv("GEMINI_KEY")

try:
    music_synthesiser = pipeline("text-to-audio", "facebook/musicgen-small")
except:
    music_synthesiser = None

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

async def band_consulting(user_input, member_name, consult_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(consult_lang_code, "Korean")
        
        system_instruction = f"""
        당신은 밴드 멤버 '{member_name}'입니다.
        1. 반드시 '{target_lang}'로만 답변하세요. 영어면 영어, 일본어면 일본어만 사용합니다.
        2. 음성 답변(TTS)은 5줄 이내로 핵심만 요약하세요.
        3. 모든 상세 설명, 화성학 조언, 텍스트 악보는 [TAB] 섹션에 넣으세요.
        4. 음악 생성용 영어 프롬프트는 마지막에 [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 "No Detailed Score"
        
        # ★ 악보 읽기 방지: [TAB] 이후 내용을 통째로 삭제하여 목소리용 텍스트 생성
        clean_text = re.sub(r'\[TAB\].*?(\[|$)', '', ai_text_raw, flags=re.DOTALL)
        clean_text = re.sub(r'\[MUSIC:.*?\]', '', clean_text).strip()
        tts_lines = clean_text.split('\n')
        tts_final = "\n".join(tts_lines[:5]) # 5줄 요약

        # 음성 생성 (특수기호 제거)
        tts_input = re.sub(r'[\*\#\-\_\~\|]', '', tts_final)
        voice_path = f"/tmp/v_{member_name.split()[0]}.mp3"
        await edge_tts.Communicate(tts_input, MEMBERS.get(member_name, "ko-KR-SunHiNeural")).save(voice_path)
        
        # 음악 생성 (512 tokens = 약 10~12초)
        music_path = None
        if music_match and music_synthesiser:
            p = music_match.group(1).strip() + ", high quality, studio recording"
            music_output = music_synthesiser(p, forward_params={"max_new_tokens": 512, "do_sample": True})
            music_path = f"/tmp/m_{member_name.split()[0]}.wav"
            scipy.io.wavfile.write(music_path, music_output["sampling_rate"], music_output["audio"][0].T)

        return tts_final, voice_path, music_path, tab_display
    except Exception as e:
        return 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()
    btn = gr.Button("GO"); btn.click(band_consulting, [i1, i2, i3], [o1, o2, o3, o4], api_name="predict")
demo.launch()