import os, re, uuid, torch, scipy.io.wavfile, edge_tts, asyncio, random
import numpy as np
import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
from groq import Groq

class ModelManager:
    _llm_pipeline = None
    _music_pipeline = None
    _groq_client = None

    @classmethod
    def get_qwen(cls):
        if cls._llm_pipeline is None:
            model_id = "Qwen/Qwen2.5-0.5B-Instruct"
            tokenizer = AutoTokenizer.from_pretrained(model_id)
            model = AutoModelForCausalLM.from_pretrained(model_id, device_map="cpu", torch_dtype=torch.float32)
            cls._llm_pipeline = pipeline("text-generation", model=model, tokenizer=tokenizer)
        return cls._llm_pipeline

    @classmethod
    def get_groq(cls):
        key = os.getenv("GROQ_API_KEY")
        if cls._groq_client is None and key:
            cls._groq_client = Groq(api_key=key)
        return cls._groq_client

    @classmethod
    def get_music(cls):
        if cls._music_pipeline is None:
            cls._music_pipeline = pipeline("text-to-audio", "facebook/musicgen-small", device="cpu")
        return cls._music_pipeline

# 8인 멤버 설정 (국가별 언어 및 보이스)
MEMBERS_CONFIG = {
    "서윤 (Korea)": {"voice": "ko-KR-SunHiNeural", "lang": "Korean"},
    "Chloe (USA)": {"voice": "en-US-AriaNeural", "lang": "English"},
    "Naomi (Japan)": {"voice": "ja-JP-NanamiNeural", "lang": "Japanese"},
    "Beatrice (Brazil)": {"voice": "pt-BR-FranciscaNeural", "lang": "Portuguese"},
    "Elena (Spain)": {"voice": "es-ES-ElviraNeural", "lang": "Spanish"},
    "Amira (Egypt)": {"voice": "ar-EG-SalmaNeural", "lang": "Arabic"},
    "Liwei (China)": {"voice": "zh-CN-XiaoxiaoNeural", "lang": "Chinese"},
    "Sophie (France)": {"voice": "fr-FR-DeniseNeural", "lang": "French"}
}

async def band_consulting(user_input, member_name, consult_lang, g_inst, b_inst, d_inst, chords):
    req_id = str(uuid.uuid4())[:8]
    voice_path = f"/tmp/v_{req_id}.mp3"
    music_path = f"/tmp/m_{req_id}.wav"
    
    # 1. 멤버 선정 로직: 영어가 선택되면 랜덤 멤버가 담당
    actual_member = member_name
    if consult_lang == "English":
        actual_member = random.choice(list(MEMBERS_CONFIG.keys()))
    
    target_lang = MEMBERS_CONFIG[actual_member]["lang"] if consult_lang == "Native" else consult_lang
    
    jam_context = f"Guitar: {g_inst}, Bass: {b_inst}, Drums: {d_inst}, Chords: {chords}"
    system_prompt = f"""You are {actual_member}. Respond ONLY in {target_lang}.
    Provide 5-7 lines of professional music advice. 
    [TAB] Section: Detailed chords/tabs.
    [MUSIC] Section: English prompt reflecting: {jam_context}"""
    
    ai_text_raw = ""
    groq_client = ModelManager.get_groq()
    if groq_client:
        try:
            res = groq_client.chat.completions.create(
                messages=[{"role": "system", "content": system_prompt}, {"role": "user", "content": user_input}],
                model="llama-3.3-70b-versatile"
            )
            ai_text_raw = res.choices[0].message.content
        except: pass

    if not ai_text_raw:
        qwen = ModelManager.get_qwen()
        out = qwen(f"<|im_start|>system\n{system_prompt}<|im_end|>\n<|im_start|>user\n{user_input}<|im_end|>\nassistant\n", max_new_tokens=1024)
        ai_text_raw = out[0]['generated_text'].split("assistant\n")[-1]

    # 파싱
    tab_match = re.search(r'\[TAB\](.*?)(\[|$)', ai_text_raw, re.DOTALL | re.IGNORECASE)
    music_match = re.search(r'\[MUSIC:(.*?)\]', ai_text_raw, re.IGNORECASE)
    tab_display = tab_match.group(1).strip() if tab_match else "No Data"
    clean_text = re.sub(r'\[TAB\].*?(\[|$)', '', ai_text_raw, flags=re.DOTALL | re.IGNORECASE)
    clean_text = re.sub(r'\[MUSIC:.*?\]', '', clean_text, flags=re.IGNORECASE).strip()
    
    # TTS
    voice_name = MEMBERS_CONFIG[actual_member]["voice"]
    communicate = edge_tts.Communicate(re.sub(r'[\*\#\-\_\~\|]', '', clean_text), voice_name)
    await communicate.save(voice_path)

    # MusicGen (12초)
    music_gen = ModelManager.get_music()
    music_p = music_match.group(1).strip() if music_match else "rock"
    music_output = music_gen(music_p, forward_params={"max_new_tokens": 512})
    audio_data = np.squeeze(music_output["audio"])
    audio_int16 = (audio_data * 32767).astype(np.int16)
    scipy.io.wavfile.write(music_path, music_output["sampling_rate"], audio_int16)

    return f"[{actual_member}] {clean_text}", voice_path, music_path, tab_display

with gr.Blocks() as demo:
    in_list = [gr.Textbox(visible=False) for _ in range(7)]
    out_list = [gr.Textbox(visible=False), gr.Audio(visible=False), gr.Audio(visible=False), gr.Textbox(visible=False)]
    btn = gr.Button("API", visible=False)
    btn.click(band_consulting, in_list, out_list, api_name="predict")

demo.queue().launch()