# ling-space/i18n/core.py
from typing import Dict
from . import common, tab_chat, tab_code, tab_smart_writer, tab_welcome, recommended_inputs

MODULES = [common, tab_chat, tab_code, tab_smart_writer, tab_welcome, recommended_inputs]

class I18nManager:
    def __init__(self):
        self.ui_translations: Dict[str, Dict[str, str]] = {}
        self.prompt_fragments: Dict[str, Dict[str, str]] = {}
        self._load_translations()

    def _load_translations(self):
        """
        Load and merge translations from all registered modules.
        Structure:
        {
            "key_name": {
                "en": "English Text",
                "zh": "中文文本"
            },
            ...
        }
        """
        for module in MODULES:
            # Merge UI Translations
            if hasattr(module, "ui_translations"):
                for key, trans_map in module.ui_translations.items():
                    if key in self.ui_translations:
                        print(f"[I18n] Warning: Duplicate UI key '{key}' in {module}. Overwriting.")
                    self.ui_translations[key] = trans_map
            
            # Merge Prompt Fragments
            if hasattr(module, "prompt_fragments"):
                for key, trans_map in module.prompt_fragments.items():
                    if key in self.prompt_fragments:
                        print(f"[I18n] Warning: Duplicate prompt key '{key}' in {module}. Overwriting.")
                    self.prompt_fragments[key] = trans_map

    def get_all_ui_translations(self) -> Dict[str, Dict[str, str]]:
        """Returns the complete dictionary for UI translations."""
        return self.ui_translations

    def get_all_prompt_translations(self) -> Dict[str, Dict[str, str]]:
        """Returns the complete dictionary for prompt fragments."""
        return self.prompt_fragments

    def get(self, key: str, lang: str, default: str) -> str:
        """
        Get a specific translation.
        Checks UI translations first, then prompt fragments.
        """
        # Check UI dict
        if key in self.ui_translations:
            return self.ui_translations[key].get(lang, default or self.ui_translations[key].get("en", key))
        
        # Check Prompt dict
        if key in self.prompt_fragments:
            return self.prompt_fragments[key].get(lang, default or self.prompt_fragments[key].get("en", key))
            
        return default or key