"""Geographic normalization helpers.

The platform stores "country" in two very different shapes:

  * ``identity_persons.country`` — an ISO-3166 alpha-2 code resolved from the
    SuccessFactors FoLocation master data (e.g. ``"CN"``, ``"US"``).
  * ``asset_office_addresses.country_or_region`` — free-text entered by an
    admin (e.g. ``"China"``, ``"中国"``, ``"United States"``, ``"USA"``).

The geographic-distribution maps (HR 人员分布 / 资产分布) key every datum by an
ISO alpha-2 code so the frontend world map (``world.zh.json``, matched on the
``iso_a2`` feature property) can colour each country with a single, stable key.
``normalize_country_code`` is the single chokepoint that collapses both shapes
into that canonical alpha-2 code.

It is deliberately conservative: anything it cannot confidently resolve returns
``None`` and the caller drops it (we never invent a country — same honesty rule
as the insights endpoints). Extend the lookup tables below as new
office/region spellings appear; do not add fuzzy matching.
"""
from __future__ import annotations

# Canonical ISO 3166-1 alpha-2 set. A raw value that is already a 2-letter
# code is accepted only if it is a real country code — otherwise garbage
# free-text like "XY" / "ZZ" would leak through and silently fail to render
# on the world map (no matching ``iso_a2`` feature). person.country (resolved
# from SuccessFactors) is always a valid code, so this drops only true noise.
_ISO_A2: frozenset[str] = frozenset(
    "AD AE AF AG AI AL AM AO AQ AR AS AT AU AW AX AZ "
    "BA BB BD BE BF BG BH BI BJ BL BM BN BO BQ BR BS BT BV BW BY BZ "
    "CA CC CD CF CG CH CI CK CL CM CN CO CR CU CV CW CX CY CZ "
    "DE DJ DK DM DO DZ EC EE EG EH ER ES ET FI FJ FK FM FO FR "
    "GA GB GD GE GF GG GH GI GL GM GN GP GQ GR GS GT GU GW GY "
    "HK HM HN HR HT HU ID IE IL IM IN IO IQ IR IS IT JE JM JO JP "
    "KE KG KH KI KM KN KP KR KW KY KZ LA LB LC LI LK LR LS LT LU LV LY "
    "MA MC MD ME MF MG MH MK ML MM MN MO MP MQ MR MS MT MU MV MW MX MY MZ "
    "NA NC NE NF NG NI NL NO NP NR NU NZ OM "
    "PA PE PF PG PH PK PL PM PN PR PS PT PW PY QA RE RO RS RU RW "
    "SA SB SC SD SE SG SH SI SJ SK SL SM SN SO SR SS ST SV SX SY SZ "
    "TC TD TF TG TH TJ TK TL TM TN TO TR TT TV TW TZ "
    "UA UG UM US UY UZ VA VC VE VG VI VN VU WF WS YE YT ZA ZM ZW".split()
)

# Common ISO alpha-3 → alpha-2 (some sources store alpha-3).
_ISO_A3_TO_A2: dict[str, str] = {
    "CHN": "CN", "HKG": "HK", "MAC": "MO", "TWN": "TW", "JPN": "JP", "KOR": "KR",
    "SGP": "SG", "MYS": "MY", "THA": "TH", "VNM": "VN", "IDN": "ID", "PHL": "PH",
    "IND": "IN", "AUS": "AU", "NZL": "NZ", "USA": "US", "CAN": "CA", "MEX": "MX",
    "BRA": "BR", "ARG": "AR", "GBR": "GB", "IRL": "IE", "FRA": "FR", "DEU": "DE",
    "NLD": "NL", "BEL": "BE", "CHE": "CH", "AUT": "AT", "ESP": "ES", "PRT": "PT",
    "ITA": "IT", "SWE": "SE", "NOR": "NO", "FIN": "FI", "DNK": "DK", "POL": "PL",
    "CZE": "CZ", "RUS": "RU", "UKR": "UA", "TUR": "TR", "ARE": "AE", "SAU": "SA",
    "ISR": "IL", "EGY": "EG", "ZAF": "ZA", "NGA": "NG", "KEN": "KE",
}

# Free-text names → alpha-2. English (incl. common aliases) + Chinese. Keys are
# lower-cased; the lookup lower-cases the input first.
_NAME_TO_A2: dict[str, str] = {
    # China & Greater China
    "china": "CN", "mainland china": "CN", "p.r. china": "CN", "prc": "CN",
    "中国": "CN", "中国大陆": "CN", "中华人民共和国": "CN",
    "hong kong": "HK", "hongkong": "HK", "香港": "HK", "中国香港": "HK",
    "macao": "MO", "macau": "MO", "澳门": "MO", "中国澳门": "MO",
    "taiwan": "TW", "台湾": "TW", "中国台湾": "TW",
    # APAC
    "japan": "JP", "日本": "JP",
    "korea": "KR", "south korea": "KR", "republic of korea": "KR", "韩国": "KR",
    "singapore": "SG", "新加坡": "SG",
    "malaysia": "MY", "马来西亚": "MY",
    "thailand": "TH", "泰国": "TH",
    "vietnam": "VN", "viet nam": "VN", "越南": "VN",
    "indonesia": "ID", "印度尼西亚": "ID", "印尼": "ID",
    "philippines": "PH", "菲律宾": "PH",
    "india": "IN", "印度": "IN",
    "australia": "AU", "澳大利亚": "AU", "澳洲": "AU",
    "new zealand": "NZ", "新西兰": "NZ",
    # Americas
    "united states": "US", "united states of america": "US", "usa": "US",
    "u.s.a.": "US", "u.s.": "US", "america": "US", "美国": "US",
    "canada": "CA", "加拿大": "CA",
    "mexico": "MX", "墨西哥": "MX",
    "brazil": "BR", "巴西": "BR",
    "argentina": "AR", "阿根廷": "AR",
    # Europe
    "united kingdom": "GB", "uk": "GB", "great britain": "GB", "england": "GB",
    "英国": "GB",
    "ireland": "IE", "爱尔兰": "IE",
    "france": "FR", "法国": "FR",
    "germany": "DE", "德国": "DE",
    "netherlands": "NL", "the netherlands": "NL", "荷兰": "NL",
    "belgium": "BE", "比利时": "BE",
    "switzerland": "CH", "瑞士": "CH",
    "austria": "AT", "奥地利": "AT",
    "spain": "ES", "西班牙": "ES",
    "portugal": "PT", "葡萄牙": "PT",
    "italy": "IT", "意大利": "IT",
    "sweden": "SE", "瑞典": "SE",
    "norway": "NO", "挪威": "NO",
    "finland": "FI", "芬兰": "FI",
    "denmark": "DK", "丹麦": "DK",
    "poland": "PL", "波兰": "PL",
    "czech republic": "CZ", "czechia": "CZ", "捷克": "CZ",
    "russia": "RU", "russian federation": "RU", "俄罗斯": "RU",
    "ukraine": "UA", "乌克兰": "UA",
    "turkey": "TR", "türkiye": "TR", "土耳其": "TR",
    # MEA
    "united arab emirates": "AE", "uae": "AE", "阿联酋": "AE",
    "saudi arabia": "SA", "沙特阿拉伯": "SA", "沙特": "SA",
    "israel": "IL", "以色列": "IL",
    "egypt": "EG", "埃及": "EG",
    "south africa": "ZA", "南非": "ZA",
    "nigeria": "NG", "尼日利亚": "NG",
    "kenya": "KE", "肯尼亚": "KE",
}


def normalize_country_code(raw: str | None) -> str | None:
    """Collapse a raw country value into an ISO-3166 alpha-2 code.

    Accepts an alpha-2 code, an alpha-3 code, or a free-text English/Chinese
    name. Returns the upper-case alpha-2 code, or ``None`` when the value is
    empty or cannot be confidently resolved (the caller drops it rather than
    guessing).
    """
    if not raw:
        return None
    value = raw.strip()
    if not value:
        return None
    upper = value.upper()
    # ASCII guard: CJK names like "中国" are 2 chars and ``str.isalpha()`` is
    # True for them, so without ``isascii()`` they'd be mistaken for a code.
    if upper.isascii() and upper.isalpha():
        if len(upper) == 2:
            return upper if upper in _ISO_A2 else None
        if len(upper) == 3:
            return _ISO_A3_TO_A2.get(upper)
    return _NAME_TO_A2.get(value.lower())