"""
   @Author      : CT
   @Description : 从 dwt_fb_base_report 解析卖家 business_addr 字段，匹配出中国省市，
                  写入 Doris dim层参考表 dim.dim_seller_address
                  这张表存储全部卖家（不分国家），只有 fb_country_name = 'CN' 的才解析省市，
                  非 CN 的 seller_province/seller_city 直接为 NULL，不做匹配尝试
                  匹配思路（只针对 fb_country_name = 'CN' 的行）：
                    1. 邮编精确匹配优先：business_addr 按双空格切分，倒数第2个 token 视为邮编，
                       去匹配 dim_china_zipcode 聚合出的"邮编→省市"表。只有一个邮编在
                       dim_china_zipcode 里对应唯一一组(province,city)才算命中；对应0组
                       （查不到/邮编本身有问题）或者多组（同一邮编跨了不同省市）统一都算
                       "邮编匹配不到"，不再像之前那样取出现次数最多的省市兜底——邮编数据本身
                       就模糊的情况，不应该被硬拍出一个可能错的答案
                    2. 邮编匹配不到的，退化用文本匹配兜底：不依赖固定 token 位置（实际地址里
                       省份不一定有、市和区的顺序不固定、市和下辖县级市可能连写、也可能只写到
                       区县压根没提市、还有拼音/英文混写的情况），改成扫描地址里除国家和邮编外
                       的所有 token，逐个跟 dim_china_zipcode 去重出来的 省份/城市/区县 字典
                       做包含关系匹配（中文名 + 拼音，拼音在这里现算，不额外维护拼音字段）：
                         a) 先扫是否有 token 命中某个省份，命中且唯一就把后续搜索范围缩小到
                            这个省份，减少同名区县跨省误判的概率
                         b) 扫 token 是否命中某个城市（比如"dongguanshi"包含"dongguan"、
                            "南阳市邓州市"包含"南阳市"），命中且唯一直接用
                         c) 城市没命中的，扫 token 是否命中某个区县（比如"仙游县"），命中且
                            唯一就反查这个区县归属的城市来用（解决地址只写到区县、没提市的情况）
                         d) 都没命中，最终 seller_province/seller_city 为 NULL
                       省/市/区县三级各自维护独立的中文+拼音后缀表（同一个"市"字省级/市级
                       含义不同，不能混用），"干净命中"只校验命中内容后面剩的部分（要么是空、
                       要么开头是同层级的一个后缀），前面剩的内容不做任何要求——这样"广东深圳
                       市"这种省市连写、"新郑市龙湖镇"这种区县+乡镇连写都能命中，不会因为命中
                       内容不在 token 开头就被拒绝
                    - fb_crawl_date 直接取 dwt_fb_base_report 里的抓取时间，作为 Doris
                      UNIQUE KEY 的 sequence_col，不用 current_timestamp() 现算写入时间；
                      该字段源表是 date_format 出来的字符串，写入前显式 to_timestamp 转成
                      timestamp，否则 Doris 连接器会报 string to timestamp 不能安全转换
                    - site_name 用脚本入参透传写入，Doris 表按 site_name LIST 自动分区，
                      多站点数据互不影响
   @Dependency  : pypinyin（只在 driver 端生成拼音字典再广播，executor 不需要装这个包）
   @SourceTable : dwt_fb_base_report（Hive）、dim_china_zipcode（Hive）
   @SinkTable   : dim.dim_seller_address（Doris）
   @Usage       : python dim_seller_address.py <site_name> <date_type> <date_info>
   @CreateTime  : 2026-07-15
   @UpdateTime  : 2026-07-20
"""

import os
import re
import sys

sys.path.append(os.path.dirname(sys.path[0]))

from pypinyin import lazy_pinyin
from utils.spark_util import SparkUtil
from utils.DorisHelper import DorisHelper
from pyspark.sql import functions as F
from pyspark.sql.types import StringType
from pyspark.sql.window import Window

DORIS_DB = "dim"
DORIS_TABLE = "dim_seller_address"
DORIS_COLUMNS = "seller_id, account_name, business_addr, fb_country_name, seller_province, seller_city, fb_crawl_date, site_name"

# 省/市/区县三级分开维护后缀表：同一个"市"字在不同层级都会出现（直辖市是省级、
# 地级市是市级），必须按当前查的是哪一级字典选用对应的后缀表，不能混用一张表。
# 顺序要把更长/更具体的后缀排在前面（比如"自治县"要排在"县"前面），不然会被短
# 后缀先命中截断错。列表依据对 dim_china_zipcode 全量 distinct 值的实际核查确定。
PROVINCE_SUFFIXES = ["自治区", "特别行政区", "省", "市"]
PROVINCE_PINYIN_SUFFIXES = ["zizhiqu", "tebiexingzhengqu", "sheng", "shi", "province"]

CITY_SUFFIXES = ["自治州", "地区", "特别行政区", "盟", "市"]
CITY_PINYIN_SUFFIXES = ["zizhizhou", "diqu", "tebiexingzhengqu", "meng", "shi", "city"]

# 区县原来不剥后缀（怕剥完只剩1个字，比如"单县"->"单"太容易误匹配），现在改成剥，
# 靠 strip_suffix 里"剥完至少剩2字"的安全网挡住这种情况，"单县"这种还是不剥
DISTRICT_SUFFIXES = ["自治县", "自治旗", "林区", "矿区", "新城", "市", "区", "县", "旗"]
DISTRICT_PINYIN_SUFFIXES = [
    "zizhixian", "zizhiqi", "xincheng", "shi", "qu", "xian", "qi", "district", "county",
]

# city 字段这几个值是无效占位符，不是真实城市名，构建 city 字典前先过滤掉
INVALID_CITY_VALUES = ["县", "省直辖县级行政区划", "自治区直辖县级行政区划"]

# 新疆几个自治区直辖县级市（图木舒克/石河子/阿拉尔）在源数据里 city 字段是"占位前缀+
# 真实市名"拼在一起的脏文本（比如"自治区直辖县图木舒克市"），剥掉前缀还原真实市名，
# 不然这三个市永远匹配不到
INVALID_CITY_PREFIX = "自治区直辖县"

# pypinyin 对多音字地名默认读音有误，人工核实过的案例收在这份覆盖表里；城市级 key 剥后缀，区县级 key 带后缀
PINYIN_OVERRIDE = {
    # ===== 城市级（剥后缀后查表）=====
    "昌都": "changdu",  # 昌都市（西藏），"都"在此读 dū，pypinyin 默认读成 dōu
    "长治": "changzhi",  # 长治市（山西），"长"在此读 cháng，pypinyin 默认读成 zhǎng
    "朝阳": "chaoyang",  # 朝阳市（辽宁），"朝"在此读 cháo，pypinyin 默认读成 zhāo
    "阿勒泰": "aletai",  # 阿勒泰地区（新疆），"勒"在此读 lè，pypinyin 默认读成 lēi
    "锡林郭勒": "xilinguole",  # 锡林郭勒盟（内蒙古），"勒"读 lè，同上
    "克孜勒苏": "kezilesu",  # 克孜勒苏柯尔克孜自治州（新疆），"勒"读 lè，同上

    # ===== 区县级（原始值直接查，部分源数据本身就不带后缀）=====
    # "都"在此读 dū，pypinyin 对不常见的地名默认读成 dōu
    "三都": "sandu", "都兰": "dulan", "都安": "duan", "都匀": "duyun",
    "丰都县": "fengduxian", "于都县": "yuduxian", "商都县": "shangduxian", "宁都县": "ningduxian",
    "尧都区": "yaoduqu", "新都区": "xinduqu", "曾都区": "cengduqu", "望都县": "wangduxian",
    "武都区": "wuduqu", "殷都区": "yinduqu", "盐都区": "yanduqu", "秦都区": "qinduqu",
    "花都区": "huaduqu", "莲都区": "lianduqu", "都昌县": "duchangxian", "魏都区": "weiduqu",
    "昆都仑区": "kundulunqu",

    # "长"在此读 cháng，pypinyin 对不常见的地名默认读成 zhǎng
    "长岛": "changdao", "长白": "changbai", "长阳": "changyang", "长顺": "changshun", "南长": "nanchang",
    "长丰县": "changfengxian", "长宁区": "changningqu", "长宁县": "changningxian", "长岭县": "changlingxian",
    "长武县": "changwuxian", "长泰区": "changtaiqu", "长洲区": "changzhouqu", "长海县": "changhaixian",
    "长清区": "changqingqu", "天长市": "tianchangshi", "子长市": "zichangshi",
    "长葛市": "changgeshi", "长垣市": "changyuanshi",

    # "朝"在此读 cháo（"朝阳区"默认已经读对，不用加）
    "朝阳县": "chaoyangxian",

    # "什"在此读 shí，pypinyin 对不常见的地名默认读成 shén
    "阿图什": "atushi", "乌什县": "wushixian", "克什克腾旗": "keshiketengqi",
    "塔什库尔干塔吉克自治县": "tashikuergantajikezizhixian", "什邡市": "shifangshi",

    # 少数民族地名音译，"勒"本该读 lè，pypinyin 默认读成 lēi
    "库尔勒": "kuerle", "策勒县": "celexian", "阿勒泰市": "aletaishi", "霍林郭勒市": "huolinguoleshi",
    "尼勒克": "nileke",

    # 其他单独核实过的多音字地名
    "泊头市": "botoushi",  # "泊"在此读 bó（停泊），pypinyin 默认读成 pō（湖泊）
    "单县": "shanxian",  # "单"在此读 shàn，pypinyin 默认读成 dān
    "民乐县": "minlexian",  # 民乐县（甘肃），pypinyin 词组词典把它当成"音乐"词组读成了 mínyuè
    "蚌山区": "bengshanqu",  # 蚌山区（安徽蚌埠），"蚌"跟"蚌埠"一样读 bèng，pypinyin 默认读成 bàng
    "荥经县": "yingjingxian",  # 荥经县（四川），"荥"在此读 yíng；跟"荥阳"(xíng)是两个不同读音的地名

    # 以下几个源表 dim_china_zipcode 里暂时查不到，先保留人工核实过的读音
    "信都": "xindu", "襄都": "xiangdu", "郫都": "pidu",  # 邢台信都区/邢台襄都区/成都郫都区
    "长白朝鲜族": "changbaichaoxianzu", "长阳土家族": "changyangtujiazu",  # xx自治县
    "康巴什": "kangbashi",  # 鄂尔多斯康巴什区
    "龙圩": "longxu",  # 梧州龙圩区，"圩"在此读 xū（南方圩镇），pypinyin 默认读成 wéi（圩堤）
}

# dim_china_zipcode 里没有港澳台，手动补充：这三个地区不分省市，省份和城市都是自己本身
SPECIAL_REGIONS = [
    ("香港", ["xianggang", "hongkong"]),
    ("澳门", ["aomen", "macau", "macao"]),
    ("台湾", ["taiwan"]),
]

# 广西/宁夏/新疆这三个自治区官方全名带了民族限定词（壮族/回族/维吾尔），单纯剥掉"自治区"
# 后缀会剩"广西壮族"这种还带着民族限定词的名字，跟地址里几乎都直接写的简称"广西"对不上，
# 所以这三个单独覆盖成简称（内蒙古自治区、西藏自治区本身没有民族限定词，不受影响）
PROVINCE_SHORT_NAME_OVERRIDE = {
    "广西壮族自治区": "广西",
    "宁夏回族自治区": "宁夏",
    "新疆维吾尔自治区": "新疆",
}


def strip_suffix(name, suffixes):
    if not name:
        return name
    for suf in suffixes:
        if name.endswith(suf) and len(name) - len(suf) >= 2:
            return name[: -len(suf)]
    return name


def clean_city_name(raw_city):
    if raw_city and raw_city.startswith(INVALID_CITY_PREFIX) and len(raw_city) > len(INVALID_CITY_PREFIX):
        return raw_city[len(INVALID_CITY_PREFIX):]
    return raw_city


def to_pinyin(word):
    if not word:
        return None
    return PINYIN_OVERRIDE.get(word) or "".join(lazy_pinyin(word))


def lookup_backfill_city(backfill_map, province, district):
    if not province or not district:
        return None
    # 回填字典的 key 同时来自 city 侧（剥 CITY_SUFFIXES）和 district 侧（剥
    # DISTRICT_SUFFIXES），两者唯一的公共后缀是"市"，用 DISTRICT_SUFFIXES 剥这里
    # 传入的查询值即可覆盖实际会出现的场景（district 字段本身的后缀只可能是市/区/县/旗）
    return backfill_map.get((province, strip_suffix(district, DISTRICT_SUFFIXES)))


def match_text_city(addr, province_lookup, city_lookup, district_lookup):
    """
    在 business_addr 里（去掉国家、邮编两个 token 后）扫描剩余 token，只认"干净命中"：
    命中后面剩下的部分要么是空的、要么开头是一个行政后缀（比如"广东省"命中"广东"剩
    "省"、"南阳市邓州市"命中"南阳"剩"市邓州市"、"jiangsushengsuzhoushi"命中
    "suzhou"剩"shi"）——前缀不做任何校验，因为"广东深圳市"这种省市连写、
    "新郑市龙湖镇"这种区县+乡镇连写都是合法地址，不能要求命中的内容必须是token开头。
    校验后缀时按当前查的是省/市/区哪一级字典，选用对应层级的后缀表（同一个"市"字
    在省级/市级都合法，必须区分开，不能混用一张表）。
    "黑龙江北路"命中"黑龙江"剩"北路"（不以任何省级后缀开头）不算命中——不然满大街
    带省市名的路名/街道名都会被误当成地址真实归属。
    先定省份缩小范围，再依次尝试命中城市/区县，返回 "province|-|city" 或 None
    """
    if not addr:
        return None
    tokens = [t.strip() for t in addr.split("  ") if t.strip()]
    if len(tokens) < 2:
        return None
    tokens = tokens[:-1]  # 去掉倒数第1个token(国家)
    # 倒数第2个token长得像邮编（纯数字）才一并去掉；有些地址本身就没填邮编，
    # 这种情况倒数第2个token其实是真实地名（比如"...顺德区  佛山市  CN"），
    # 误当邮编砍掉会把仅有的城市信息也搜没了，所以不像邮编就保留它参与匹配
    if tokens and re.fullmatch(r"\d{4,10}", tokens[-1]):
        tokens = tokens[:-1]
    if not tokens:
        return None
    # 去空格再比较拼音：不然"Guang Dong"这种按音节分开写的拼音永远匹配不上"guangdong"
    tokens_low = [re.sub(r"\s+", "", t.lower()) for t in tokens]

    def is_clean(base, text, suffixes):
        if not base:
            return False
        idx = text.find(base)
        if idx == -1:
            return False
        leftover = text[idx + len(base):]
        return leftover == "" or any(leftover.startswith(suf) for suf in suffixes)

    def is_hit(base_cn, base_py, cn_suffixes, py_suffixes):
        for tok, tok_low in zip(tokens, tokens_low):
            if is_clean(base_cn, tok, cn_suffixes):
                return True
            if is_clean(base_py, tok_low, py_suffixes):
                return True
        return False

    def unique_match(lookup, key_of, restrict_province, cn_suffixes, py_suffixes):
        # lookup 里每行最后两列固定是 (base_cn, base_py)，前面几列由 key_of 决定取哪些做候选key
        matched = set()
        for row in lookup:
            if restrict_province and row[0] != restrict_province:
                continue
            if is_hit(row[-2], row[-1], cn_suffixes, py_suffixes):
                matched.add(key_of(row))
        return next(iter(matched)) if len(matched) == 1 else None

    matched_province = unique_match(
        province_lookup, key_of=lambda row: row[0], restrict_province=None,
        cn_suffixes=PROVINCE_SUFFIXES, py_suffixes=PROVINCE_PINYIN_SUFFIXES,
    )

    picked = unique_match(
        city_lookup, key_of=lambda row: (row[0], row[1]), restrict_province=matched_province,
        cn_suffixes=CITY_SUFFIXES, py_suffixes=CITY_PINYIN_SUFFIXES,
    )
    if picked:
        return f"{picked[0]}|-|{picked[1]}"

    picked = unique_match(
        district_lookup, key_of=lambda row: (row[0], row[1]), restrict_province=matched_province,
        cn_suffixes=DISTRICT_SUFFIXES, py_suffixes=DISTRICT_PINYIN_SUFFIXES,
    )
    if picked:
        return f"{picked[0]}|-|{picked[1]}"

    return None


class DimSellerAddress(object):

    def __init__(self, site_name, date_type, date_info):
        self.site_name = site_name
        self.date_type = date_type
        self.date_info = date_info
        app_name = f"{self.__class__.__name__}:{site_name}:{date_info}"
        self.spark = SparkUtil.get_spark_session(app_name)
        self.df_seller = None
        self.df_zip_map = None
        self.province_lookup = None
        self.city_lookup = None
        self.district_lookup = None
        self.city_backfill_map = None
        self.province_lookup_bc = None
        self.city_lookup_bc = None
        self.district_lookup_bc = None
        self.city_backfill_bc = None
        self.df_save = None

    def read_data(self):
        print("读取 dwt_fb_base_report")
        sql = f"""
        select seller_id, account_name, business_addr, fb_country_name, fb_crawl_date
        from dwt_fb_base_report
        where site_name = '{self.site_name}'
          and date_type = '{self.date_type}'
          and date_info = '{self.date_info}'
        """
        self.df_seller = self.spark.sql(sqlQuery=sql)
        print(sql)

        print("构建省份/城市/区县文本匹配字典 + city缺失回填字典（driver端生成拼音，广播给executor）")
        self._build_text_lookups()

        print("读取 dim_china_zipcode 全量原始行，用回填字典补全 city 缺失的行，再聚合出邮编->省市归属"
              "（同邮编对应不止一组省市的，视为匹配不到）")
        self._build_zip_map()

    def _build_text_lookups(self):
        provinces = [
            r["province"] for r in
            self.spark.sql("select distinct province from dim_china_zipcode where province is not null").collect()
        ]
        def province_base_name(p):
            return PROVINCE_SHORT_NAME_OVERRIDE.get(p) or strip_suffix(p, PROVINCE_SUFFIXES)

        self.province_lookup = []
        for p in provinces:
            base = province_base_name(p)
            self.province_lookup.append((p, base, to_pinyin(base)))

        invalid_city_clause = ", ".join(f"'{v}'" for v in INVALID_CITY_VALUES)
        city_rows = self.spark.sql(
            "select distinct province, city from dim_china_zipcode "
            "where province is not null and city is not null "
            f"and city not in ({invalid_city_clause})"
        ).collect()
        self.city_lookup = []
        for r in city_rows:
            city_name = clean_city_name(r["city"])
            base = strip_suffix(city_name, CITY_SUFFIXES)
            self.city_lookup.append((r["province"], city_name, base, to_pinyin(base)))

        # 同一个省份+剥后缀名字，如果同时存在"XX地区"和"XX市"两个原始城市名（比如"铜仁地区"/
        # "铜仁市"，2011年撤地设市遗留的新旧两个名字都留在源数据里），文本匹配时两个候选长度
        # 一样、谁都选不出来，所以只保留现在还在用的"市"，把过时的"地区"去掉
        by_stripped = {}
        for row in self.city_lookup:
            by_stripped.setdefault((row[0], row[2]), []).append(row)
        deduped_city_lookup = []
        for (province, stripped), rows in by_stripped.items():
            if len(rows) > 1:
                shi_rows = [r for r in rows if r[1].endswith("市")]
                rows = shi_rows or rows
            deduped_city_lookup.extend(rows)
        self.city_lookup = deduped_city_lookup

        # 港澳台不分省市，省份和城市都是自己本身
        for name, pinyins in SPECIAL_REGIONS:
            for py in pinyins:
                self.province_lookup.append((name, name, py))
                self.city_lookup.append((name, name, name, py))

        # length(district)>=2：源数据里有极少数区县字段是单字垃圾值（比如上海某几条记录
        # district 字段直接是"区"这一个字），单字后面又刚好是行政后缀本身，会跟全国任何
        # 以"区"结尾的 token 撞车，真实的区县名没有单字的，直接过滤掉
        district_rows = self.spark.sql(
            "select distinct province, city, district from dim_china_zipcode "
            "where province is not null and city is not null and district is not null and district != '' "
            "and length(district) >= 2"
        ).collect()
        self.district_lookup = []
        for r in district_rows:
            base = strip_suffix(r["district"], DISTRICT_SUFFIXES)
            self.district_lookup.append((r["province"], r["city"], base, to_pinyin(base)))

        # 同一个(省份,区县)如果同时归属不止一个城市（比如"公主岭市"曾归四平市、2014年后
        # 改由长春市代管；"江口县"曾归"铜仁地区"、现归"铜仁市"），能靠"只有一个以市结尾"
        # 消歧的就消歧（这种是地区改市/地区代管变更，取现在这个）；两个都以市结尾的说明是
        # 更复杂的历史沿革（比如"襄樊市"改名"襄阳市"）或者单纯两个市恰好有同名区县，没法
        # 安全判断，这条区县整个丢弃，不瞎猜
        by_province_district = {}
        for row in self.district_lookup:
            by_province_district.setdefault((row[0], row[2]), []).append(row)
        deduped_district_lookup = []
        for (province, district), rows in by_province_district.items():
            if len({r[1] for r in rows}) > 1:
                shi_rows = [r for r in rows if r[1].endswith("市")]
                if len(shi_rows) != 1:
                    continue
                rows = shi_rows
            deduped_district_lookup.extend(rows)
        self.district_lookup = deduped_district_lookup

        # 上面"地区/市"去重把过时的城市名（比如"铜仁地区"）从 city_lookup 删掉了；区县
        # 字典里如果还有挂在这些过时城市名下面的区县（比如撤地设市前，"铜仁地区"下面
        # 有个同名的县级市"铜仁市"，这条归属数据现在已经陈旧、不可信），先清掉，避免
        # 拿这种陈旧归属去跟当前有效的城市名做下面的跨级去重判断
        valid_cities_by_province = {}
        for row in self.city_lookup:
            valid_cities_by_province.setdefault(row[0], set()).add(row[1])
        self.district_lookup = [
            row for row in self.district_lookup
            if row[1] in valid_cities_by_province.get(row[0], set())
        ]

        # "省直辖县级市"这种（比如新郑市实际由郑州市代管）在源数据里经常同时出现在两个
        # 层级：city_lookup 里是独立一条"新郑市"，district_lookup 里又挂在"郑州市"下面
        # 当区县，两边都留着会导致城市这一级判不出唯一结果（"新郑市"/"郑州市"两个候选
        # 打平）。以 district 层级的归属为准：city_lookup 里的城市名如果在同一个省份的
        # district_lookup 里也作为"挂在别的城市下面的区县"出现，就不再当独立城市，从
        # city_lookup 剔除，只留 district 那条，靠区县反查城市的机制处理。
        # 必须限定"别的城市"：很多地级市自己就辖有同名的县（长沙市/长沙县、衡阳市/衡阳
        # 县这种极常见的市县同名），这种归属城市就是自己本身，不能算冲突去掉自己
        # district_lookup 的 row[2] 现在是剥过 DISTRICT_SUFFIXES 的区县名，city_lookup
        # 的 row[1] 还是原始城市名，两边唯一的公共后缀是"市"，按 CITY_SUFFIXES 剥一下
        # city 名字再比较，才能对上（"新郑市"剥完是"新郑"，跟区县那边剥完的"新郑"一致）
        district_names_by_province = {}
        for row in self.district_lookup:
            district_names_by_province.setdefault(row[0], set()).add((row[2], row[1]))
        self.city_lookup = [
            row for row in self.city_lookup
            if not any(
                district_name == strip_suffix(row[1], CITY_SUFFIXES) and governing_city != row[1]
                for district_name, governing_city in district_names_by_province.get(row[0], set())
            )
        ]

        # 新版邮编源数据里有一批行 city 是空的，实际城市名被错放进了 district 字段（比如
        # 省=江苏省 市="" 区=苏州），或者 district 是真区县、只是这一行恰好没填 city（比如
        # 省=广东省 市="" 区=宝安）。用"city不为空"的正常数据反过来建一个
        # (省份, 剥后缀的名字) -> 城市 的回填字典：
        #   1) 名字剥完后缀等于某个城市剥后缀的名字，回填这个城市（"苏州"->"苏州市"）
        #   2) 名字是某个区县、且这个区县在正常数据里归属唯一，回填归属的城市（"宝安"->"深圳市"）
        # 同一个 (省份,名字) 两种来源指向不同城市的，说明有歧义，不回填，city 继续留空不参与匹配
        backfill, conflict = {}, set()

        def add_candidate(province, name, city, suffixes):
            key = (province, strip_suffix(name, suffixes))
            if key in backfill and backfill[key] != city:
                conflict.add(key)
            else:
                backfill[key] = city

        for r in city_rows:
            clean_city = clean_city_name(r["city"])
            add_candidate(r["province"], clean_city, clean_city, CITY_SUFFIXES)
        for r in district_rows:
            add_candidate(r["province"], r["district"], r["city"], DISTRICT_SUFFIXES)
        self.city_backfill_map = {k: v for k, v in backfill.items() if k not in conflict}

        # city 缺失但能靠回填字典唯一推出城市的行，额外补进区县字典，让文本匹配也能兜住
        # （地址里正好写了这种不带后缀的区县名，比如"宝安"而不是"宝安区"）
        backfill_district_rows = self.spark.sql(
            "select distinct province, district from dim_china_zipcode "
            "where province is not null and city is null and district is not null and district != '' "
            "and length(district) >= 2"
        ).collect()
        for r in backfill_district_rows:
            city = lookup_backfill_city(self.city_backfill_map, r["province"], r["district"])
            if city:
                self.district_lookup.append((r["province"], city, r["district"], to_pinyin(r["district"])))

        print(
            f"字典规模：province={len(self.province_lookup)}，city={len(self.city_lookup)}，"
            f"district={len(self.district_lookup)}，city回填={len(self.city_backfill_map)}"
        )

        # 显式广播，避免字典跟着 UDF 闭包在每个 task 里重复序列化
        self.province_lookup_bc = self.spark.sparkContext.broadcast(self.province_lookup)
        self.city_lookup_bc = self.spark.sparkContext.broadcast(self.city_lookup)
        self.district_lookup_bc = self.spark.sparkContext.broadcast(self.district_lookup)
        self.city_backfill_bc = self.spark.sparkContext.broadcast(self.city_backfill_map)

    def _build_zip_map(self):
        df_raw = self.spark.sql(
            "select province, city, district, zipcode from dim_china_zipcode "
            "where province is not null and zipcode is not null"
        )

        city_backfill_bc = self.city_backfill_bc

        def backfill_city(province, city, district):
            if city:
                return clean_city_name(city)
            return lookup_backfill_city(city_backfill_bc.value, province, district)

        u_backfill_city = F.udf(backfill_city, StringType())

        df_filled = df_raw.withColumn(
            "city_filled", u_backfill_city(F.col("province"), F.col("city"), F.col("district"))
        ).filter(F.col("city_filled").isNotNull())

        df_distinct = df_filled.select(
            F.col("province"), F.col("city_filled").alias("city"), F.col("zipcode")
        ).distinct()

        combo_cnt = F.count(F.lit(1)).over(Window.partitionBy("zipcode"))
        self.df_zip_map = df_distinct.withColumn("combo_cnt", combo_cnt).filter(F.col("combo_cnt") == 1).select(
            F.col("zipcode"),
            F.col("province").alias("map_province"),
            F.col("city").alias("map_city"),
        )

    def handle_data(self):
        print("解析 business_addr 匹配省市（仅 fb_country_name = 'CN' 的行）")
        addr_split = F.split(F.col("business_addr"), "  ")
        # 邮编 token 里偶尔会混入零宽字符等不可见字符（复制粘贴带过来的），精确比较会
        # 匹配不上字典里干净的邮编，这里只保留数字部分再去 join，不影响正常邮编
        addr_zip_clean = F.regexp_replace(F.element_at(addr_split, -2), r"[^0-9]", "")
        df = self.df_seller.withColumn(
            "addr_zip",
            F.when(F.col("fb_country_name") == "CN", addr_zip_clean).otherwise(F.lit(None)),
        )

        df = df.join(self.df_zip_map, df["addr_zip"] == self.df_zip_map["zipcode"], how="left")

        province_lookup_bc, city_lookup_bc, district_lookup_bc = (
            self.province_lookup_bc, self.city_lookup_bc, self.district_lookup_bc,
        )

        # 短路判断必须写在函数体内部：F.when 不会让 Spark 跳过 UDF 调用本身（实测过）
        def udf_match_text(addr, country, zip_matched):
            if zip_matched or country != "CN":
                return None
            return match_text_city(addr, province_lookup_bc.value, city_lookup_bc.value, district_lookup_bc.value)

        u_match_text = F.udf(udf_match_text, StringType())

        df = df.withColumn(
            "text_match",
            u_match_text(F.col("business_addr"), F.col("fb_country_name"), F.col("map_province").isNotNull()),
        )
        text_split = F.split(F.col("text_match"), r"\|-\|")

        self.df_save = df.select(
            F.col("seller_id"),
            F.col("account_name"),
            F.col("business_addr"),
            F.col("fb_country_name"),
            F.coalesce(F.col("map_province"), text_split.getItem(0)).alias("seller_province"),
            F.coalesce(F.col("map_city"), text_split.getItem(1)).alias("seller_city"),
            # dwt_fb_base_report.fb_crawl_date 是 date_format 出来的字符串，Doris 表定义的是
            # DATETIME，写入前要显式转成 timestamp，不然 Doris 连接器会报
            # "Cannot safely cast 'fb_crawl_date': string to timestamp"
            F.to_timestamp(F.col("fb_crawl_date"), "yyyy-MM-dd HH:mm:ss").alias("fb_crawl_date"),
            F.lit(self.site_name).alias("site_name"),
        ).cache()

    def save_data(self):
        total = self.df_save.count()
        print(f"写入 Doris {DORIS_DB}.{DORIS_TABLE}，共 {total} 条")
        DorisHelper.spark_export_with_columns(
            df_save=self.df_save,
            db_name=DORIS_DB,
            table_name=DORIS_TABLE,
            table_columns=DORIS_COLUMNS,
            use_type="selection",
        )
        self.df_save.unpersist()
        print("写入完毕")

    def run(self):
        self.read_data()
        self.handle_data()
        self.save_data()


if __name__ == "__main__":
    site_name = sys.argv[1] if len(sys.argv) > 1 else "us"
    date_type = sys.argv[2] if len(sys.argv) > 2 else "month"
    date_info = sys.argv[3] if len(sys.argv) > 3 else None

    if not date_info:
        print("必须传入 date_info 参数")
        sys.exit(1)

    DimSellerAddress(site_name=site_name, date_type=date_type, date_info=date_info).run()
