import random
import yaml
import os
from typing import Dict, List, Any

class MusicPromptGenerator:
    """
    音樂 AI Prompt 生成器 - ComfyUI 節點版本
    """
    
    def __init__(self):
        self.keyword_map = {
            "love": {"mood": "romantic", "genre": "pop", "theme": "love"},
            "rain": {"mood": "melancholic", "genre": "lofi", "theme": "rain"},
            "battle": {"mood": "energetic", "genre": "orchestral", "theme": "journey"},
            "dream": {"mood": "dreamy", "genre": "ambient", "theme": "fantasy"},
            "night": {"mood": "serene", "genre": "jazz", "theme": "night"},
            "space": {"mood": "mysterious", "genre": "electronic", "theme": "scifi"},
            "road": {"mood": "nostalgic", "genre": "rock", "theme": "journey"},
            "sunrise": {"mood": "happy", "genre": "ambient", "theme": "nature"},
            "sad": {"mood": "melancholic", "genre": "piano", "theme": "memory"}
        }
        
        self.default_data = {
            "moods": [
                {"en": "happy", "zh": "快樂"}, {"en": "dark", "zh": "陰鬱"},
                {"en": "nostalgic", "zh": "懷舊"}, {"en": "melancholic", "zh": "憂傷"},
                {"en": "euphoric", "zh": "狂喜"}, {"en": "mysterious", "zh": "神秘"},
                {"en": "serene", "zh": "寧靜"}, {"en": "tense", "zh": "緊張"},
                {"en": "romantic", "zh": "浪漫"}, {"en": "energetic", "zh": "活力"}
            ],
            "genres": [
                {"en": "pop", "zh": "流行"}, {"en": "jazz", "zh": "爵士"},
                {"en": "rock", "zh": "搖滾"}, {"en": "electronic", "zh": "電子"},
                {"en": "classical", "zh": "古典"}, {"en": "folk", "zh": "民謠"},
                {"en": "lofi", "zh": "低保真"}, {"en": "orchestral", "zh": "管弦樂"},
                {"en": "ambient", "zh": "環境音樂"}, {"en": "synthwave", "zh": "合成器浪潮"}
            ],
            "themes": [
                {"en": "love", "zh": "愛情"}, {"en": "rain", "zh": "雨天"},
                {"en": "city", "zh": "都市夜景"}, {"en": "ocean", "zh": "海洋"},
                {"en": "fantasy", "zh": "奇幻"}, {"en": "scifi", "zh": "科幻"},
                {"en": "journey", "zh": "旅程"}, {"en": "memory", "zh": "回憶"},
                {"en": "nature", "zh": "自然"}, {"en": "night", "zh": "夜晚"}
            ],
            "instruments": [
                {"en": "piano", "zh": "鋼琴"}, {"en": "guitar", "zh": "吉他"},
                {"en": "synth", "zh": "合成器"}, {"en": "strings", "zh": "弦樂"},
                {"en": "drums", "zh": "鼓組"}, {"en": "bass", "zh": "貝斯"},
                {"en": "violin", "zh": "小提琴"}, {"en": "saxophone", "zh": "薩克斯風"},
                {"en": "flute", "zh": "長笛"}, {"en": "harp", "zh": "豎琴"}
            ],
            "vocal_styles": [
                {"en": "female vocals", "zh": "女聲"}, {"en": "male vocals", "zh": "男聲"},
                {"en": "choir", "zh": "合唱"}, {"en": "layered vocals", "zh": "層疊人聲"},
                {"en": "whisper", "zh": "氣音呢喃"}, {"en": "falsetto", "zh": "假音"},
                {"en": "rap", "zh": "說唱"}, {"en": "spoken word", "zh": "念白"}
            ],
            "descriptions": [
                {"en": "cinematic", "zh": "電影感"}, {"en": "lofi", "zh": "輕鬆低保真"},
                {"en": "acoustic", "zh": "純淨原聲"}, {"en": "orchestral", "zh": "管弦編制"},
                {"en": "minimal", "zh": "極簡主義"}, {"en": "experimental", "zh": "實驗前衛"},
                {"en": "vintage", "zh": "復古風"}, {"en": "futuristic", "zh": "未來感"}
            ],
            "tempos": [
                {"en": "slow", "zh": "慢板"}, {"en": "medium", "zh": "中板"},
                {"en": "fast", "zh": "快板"}, {"en": "adagio", "zh": "柔板"},
                {"en": "allegro", "zh": "快板"}, {"en": "andante", "zh": "行板"}
            ]
        }

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "mode": (["random", "keyword"], {"default": "random"}),
                "keyword": ("STRING", {"default": "love", "multiline": False}),
                "num_prompts": ("INT", {"default": 1, "min": 1, "max": 10}),
                "language": (["en", "zh", "both"], {"default": "en"}),
                "template_type": (["default", "cinematic", "minimal", "custom"], {"default": "default"}),
                "custom_template": ("STRING", {
                    "default": "A {mood} {genre} track about {theme}, featuring {instrument} and {vocal_style}, {description} style, {tempo} tempo.",
                    "multiline": True
                }),
                "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
            },
            "optional": {
                "yaml_file": ("STRING", {"default": ""}),
            }
        }

    RETURN_TYPES = ("STRING",)
    RETURN_NAMES = ("music_prompt",)
    FUNCTION = "generate_prompt"
    CATEGORY = "AI Music"
    DESCRIPTION = "生成音樂 AI 提示詞"

    def load_yaml_data(self, file_path):
        """載入 YAML 詞庫檔案"""
        if not file_path or not os.path.exists(file_path):
            return self.default_data
        
        try:
            with open(file_path, 'r', encoding='utf-8') as f:
                return yaml.safe_load(f)
        except Exception as e:
            print(f"載入 YAML 檔案失敗: {e}")
            return self.default_data

    def fill_template(self, mood, genre, theme, instrument, vocal, desc, tempo, template, language):
        """填充模板"""
        en_template = template
        zh_template = "一首{mood}{genre}的音樂，主題是{theme}，包含{instrument}與{vocal_style}，風格{description}，節奏{tempo}。"
        
        # 確保所有參數都是字典格式
        mood_dict = mood if isinstance(mood, dict) else {"en": mood, "zh": mood}
        genre_dict = genre if isinstance(genre, dict) else {"en": genre, "zh": genre}
        theme_dict = theme if isinstance(theme, dict) else {"en": theme, "zh": theme}
        instrument_dict = instrument if isinstance(instrument, dict) else {"en": instrument, "zh": instrument}
        vocal_dict = vocal if isinstance(vocal, dict) else {"en": vocal, "zh": vocal}
        desc_dict = desc if isinstance(desc, dict) else {"en": desc, "zh": desc}
        tempo_dict = tempo if isinstance(tempo, dict) else {"en": tempo, "zh": tempo}
        
        en = en_template.format(
            mood=mood_dict["en"], genre=genre_dict["en"], theme=theme_dict["en"],
            instrument=instrument_dict["en"], vocal_style=vocal_dict["en"],
            description=desc_dict["en"], tempo=tempo_dict["en"]
        )
        zh = zh_template.format(
            mood=mood_dict["zh"], genre=genre_dict["zh"], theme=theme_dict["zh"],
            instrument=instrument_dict["zh"], vocal_style=vocal_dict["zh"],
            description=desc_dict["zh"], tempo=tempo_dict["zh"]
        )
        
        if language == "en":
            return en
        elif language == "zh":
            return zh
        else:
            return f"{en}\n{zh}"

    def get_template(self, template_type, custom_template):
        """取得模板"""
        templates = {
            "default": "A {mood} {genre} track about {theme}, featuring {instrument} and {vocal_style}, {description} style, {tempo} tempo.",
            "cinematic": "Cinematic {genre} music with {mood} atmosphere, theme: {theme}, featuring {instrument} and {vocal_style}, {tempo} tempo.",
            "minimal": "{mood} {genre} | {theme} | {instrument} | {vocal_style} | {description} | {tempo}",
        }
        return custom_template if template_type == "custom" else templates.get(template_type, templates["default"])

    def generate_prompt(self, mode, keyword, num_prompts, language, template_type, custom_template, seed, yaml_file=""):
        # 設定隨機種子
        if seed > 0:
            random.seed(seed)
        
        # 載入資料
        data = self.load_yaml_data(yaml_file) if yaml_file else self.default_data
        
        # 取得模板
        template = self.get_template(template_type, custom_template)
        
        results = []
        
        for _ in range(num_prompts):
            if mode == "keyword":
                # 關鍵字模式
                match = self.keyword_map.get(keyword.lower(), {})
                
                def find_item(category, default_value):
                    target = match.get(category, default_value)
                    for item in data[category + "s"]:
                        if isinstance(item, dict) and item.get("en") == target:
                            return item
                    return random.choice(data[category + "s"])
                
                mood = find_item("mood", "melancholic")
                genre = find_item("genre", "pop")
                theme = find_item("theme", "memory")
                instrument = random.choice(data["instruments"])
                vocal = random.choice(data["vocal_styles"])
                desc = random.choice(data["descriptions"])
                tempo = random.choice(data["tempos"])
                
            else:
                # 隨機模式
                mood = random.choice(data["moods"])
                genre = random.choice(data["genres"])
                theme = random.choice(data["themes"])
                instrument = random.choice(data["instruments"])
                vocal = random.choice(data["vocal_styles"])
                desc = random.choice(data["descriptions"])
                tempo = random.choice(data["tempos"])
            
            prompt = self.fill_template(mood, genre, theme, instrument, vocal, desc, tempo, template, language)
            results.append(prompt)
        
        # 合併多個提示詞
        final_prompt = "\n\n".join(results) if num_prompts > 1 else results[0]
        
        return (final_prompt,)

# 節點註冊
NODE_CLASS_MAPPINGS = {
    "MusicPromptGenerator": MusicPromptGenerator
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "MusicPromptGenerator": "🎵 Music AI Prompt Generator"
}