"""Locale-aware main agent prompt loading."""

from __future__ import annotations

from functools import lru_cache
from pathlib import Path

DEFAULT_MAIN_AGENT_PROMPT_LOCALE = "zh-Hans"

_PROMPT_FILES = {
    "zh-Hans": "zh-Hans.md",
    "zh-Hant": "zh-Hant.md",
    "en": "en.md",
}

_LOCALE_ALIASES = {
    "zh": "zh-Hans",
    "zh-cn": "zh-Hans",
    "zh-hans": "zh-Hans",
    "zh-sg": "zh-Hans",
    "zh-hant": "zh-Hant",
    "zh-tw": "zh-Hant",
    "zh-hk": "zh-Hant",
    "zh-mo": "zh-Hant",
    "en": "en",
    "en-us": "en",
    "en-gb": "en",
}


def get_main_agent_prompt(locale: str | None = None) -> str:
    """Return the main-agent identity prompt for a prompt locale."""

    prompt_locale = normalize_main_agent_prompt_locale(locale)
    return _load_main_agent_prompt(prompt_locale)


def normalize_main_agent_prompt_locale(locale: str | None = None) -> str:
    cleaned = (locale or DEFAULT_MAIN_AGENT_PROMPT_LOCALE).strip()
    if not cleaned:
        return DEFAULT_MAIN_AGENT_PROMPT_LOCALE
    normalized = _LOCALE_ALIASES.get(cleaned.lower())
    if normalized:
        return normalized
    return cleaned if cleaned in _PROMPT_FILES else DEFAULT_MAIN_AGENT_PROMPT_LOCALE


@lru_cache(maxsize=len(_PROMPT_FILES))
def _load_main_agent_prompt(locale: str) -> str:
    filename = _PROMPT_FILES.get(locale, _PROMPT_FILES[DEFAULT_MAIN_AGENT_PROMPT_LOCALE])
    path = Path(__file__).with_name("main_agent") / filename
    if not path.exists():
        fallback_path = Path(__file__).with_name("main_agent") / _PROMPT_FILES[DEFAULT_MAIN_AGENT_PROMPT_LOCALE]
        return fallback_path.read_text(encoding="utf-8").strip()
    return path.read_text(encoding="utf-8").strip()