project/mt_V2/area.py

"""
区域数据服务（基于 city.json 省/市/区三级数据）

数据格式：省(id/name/pid) → 市(id/name/pid/sons) → 区县名称字符串数组
直辖市特征：省名 == 市名（如"北京市"省下有唯一城市"北京市"）
"""
import json
import os
import logging
from collections import namedtuple
from typing import Optional, Tuple, Dict, Any, List, Set

log = logging.getLogger("area_service")

# ---------- 返回类型 ----------

AreaMatch = namedtuple('AreaMatch', ['province_id', 'city_id', 'province', 'city'])
AreaMatch.__new__.__annotations__ = {
    'province_id': int,
    'city_id': int,
    'province': str,
    'city': str,
}

# ---------- 常量 ----------

_PROVINCE_SUFFIXES = ['壮族自治区', '回族自治区', '维吾尔自治区', '自治区', '省', '市']
_CITY_SUFFIXES = [
    # 民族自治州复合后缀（按长度降序，确保最长匹配优先）
    '朝鲜族自治州', '土家族苗族自治州', '藏族羌族自治州',
    '哈尼族彝族自治州', '布依族苗族自治州', '苗族侗族自治州',
    '壮族苗族自治州', '傣族景颇族自治州', '蒙古族藏族自治州',
    '柯尔克孜自治州', '哈萨克自治州', '蒙古自治州',
    '藏族自治州', '彝族自治州', '白族自治州', '傣族自治州',
    '傈僳族自治州', '回族自治州',
    '自治州', '地区', '盟', '市',
]
_DISTRICT_SUFFIXES = [
    '自治县', '林区', '特区', '区', '县', '市', '旗',
    '朝鲜族自治县', '回族彝族自治县', '彝族回族苗族自治县',
    '哈尼族彝族傣族自治县',
]

# 虚拟城市名称别名（仅加入省级→市 Trie，不加入全局 Trie 以避免多省冲突）
_VIRTUAL_CITY_ALIASES = {
    '省直辖县级行政区划': ['省直辖县', '省辖县', '直辖县', '辖县'],
    '自治区直辖县级行政区划': ['兵团城市'],
}


# ---------- 工具函数 ----------

def _strip_suffix(name, suffixes):
    """尝试去掉后缀得到简称，去掉最长的匹配"""
    for s in sorted(suffixes, key=len, reverse=True):
        if name.endswith(s) and len(name) > len(s):
            return name[:-len(s)]
    return name


class _TrieNode:
    """前缀树节点（支持多值共存）"""
    __slots__ = ('children', 'values')

    def __init__(self) -> None:
        self.children: Dict[str, '_TrieNode'] = {}
        # 多值: [(type, id, city_id), ...]
        # city_id: province=0, city=c_id, district=所属city_id
        self.values: List[Tuple[str, int, int]] = []


class _Trie:
    """
    统一前缀树，省/市/区 逐字链式构造

    插入模式:
      - insert(word, type, id, city_id):            从 root 插入
      - insert_from(node, word, type, id, city_id): 从指定节点继续插入

    搜索规则:
      - 逐字遍历 addr，沿途收集最佳匹配
      - 停止位检查 city_id 唯一性：多种 city_id → 歧义 → None
    """

    _TYPE_PRIORITY = {'district': 3, 'city': 2, 'province': 1}

    def __init__(self) -> None:
        self.root = _TrieNode()

    # ---------- 插入 ----------

    @staticmethod
    def _add_value(node: _TrieNode, entry_type: str, entry_id: int, city_id: int) -> None:
        """在节点上添加值（仅去重，冲突由 search 停止位检测）"""
        for t, eid, cid in node.values:
            if t == entry_type and eid == entry_id:
                return  # 完全相同，忽略
        node.values.append((entry_type, entry_id, city_id))

    def insert(self, word: str, entry_type: str, entry_id: int,
               city_id: int = 0) -> List['_TrieNode']:
        """从 root 插入，返回叶节点列表"""
        node = self.root
        for ch in word:
            if ch not in node.children:
                node.children[ch] = _TrieNode()
            node = node.children[ch]
        _Trie._add_value(node, entry_type, entry_id, city_id)
        return [node]

    @staticmethod
    def insert_from(from_nodes: List['_TrieNode'],
                    word: str, entry_type: str, entry_id: int,
                    city_id: int = 0) -> List['_TrieNode']:
        """从指定节点列表继续逐字插入，返回新叶节点列表"""
        leaves = []
        for start in from_nodes:
            node = start
            for ch in word:
                if ch not in node.children:
                    node.children[ch] = _TrieNode()
                node = node.children[ch]
            _Trie._add_value(node, entry_type, entry_id, city_id)
            leaves.append(node)
        return leaves

    # ---------- 搜索 ----------

    def _best_value(self, node: _TrieNode) -> Optional[Tuple[str, int]]:
        """获取节点上最佳值（district > city > province）"""
        best = None
        best_pri = 0
        for t, eid, _ in node.values:
            pri = self._TYPE_PRIORITY.get(t, 0)
            if pri > best_pri:
                best_pri = pri
                best = (t, eid)
        return best

    def search(self, text: str) -> Optional[Tuple[str, int]]:
        """
        逐字遍历 text，返回沿途最佳匹配 (type, id)

        - 记录停止节点（沿途最后有值的节点）
        - 停止位检查：节点上所有 city_id 是否一致
          - 一致 → 返回停止节点最佳值
          - 不一致（多种 city_id）→ 歧义 → None
        """
        if not text:
            return None
        node = self.root
        stop_node = None
        for ch in text:
            if ch not in node.children:
                break  # 路径断了
            node = node.children[ch]
            if self._best_value(node) is not None:
                stop_node = node
        if stop_node is None:
            return None
        # 停止位 city_id 唯一性检查
        if len(stop_node.values) > 1:
            city_ids = set(cid for _, _, cid in stop_node.values)
            if len(city_ids) > 1:
                return None  # 多种 city_id → 歧义
        return self._best_value(stop_node)


class AreaService:
    """区域数据服务，从 city.json 加载并构建多级索引"""

    def __init__(self, area_data_json: str, addr_prefix_json: Optional[str] = None) -> None:
        area_data_json = os.path.normpath(area_data_json)
        log.info(f'area service loading from {area_data_json}')

        with open(area_data_json, 'r', encoding='utf-8') as f:
            data = json.load(f)

        # 数据容器
        self._province_map: Dict[int, Dict[str, Any]] = {}   # province_id -> {id, name, short_name}
        self._city_map: Dict[int, Dict[str, Any]] = {}       # city_id -> {id, name, short_name, province_id}

        # 层级索引
        self._addr_prefix_map: Dict[str, int] = {}            # 手工前缀 -> city_id
        self._area_trie: _Trie = _Trie()                      # 统一前缀树：省/市/区 -> id

        self._build_maps(data)
        self._build_indexes()

        # 加载可选的地址前缀映射
        if addr_prefix_json:
            self._load_addr_prefix_map(addr_prefix_json)

        log.info(
            f'area service loaded: {len(self._province_map)} provinces, '
            f'{len(self._city_map)} cities, {len(self._districts)} districts'
        )

    # ---------- 内部构建 ----------

    def _build_maps(self, data: List[Dict[str, Any]]) -> None:
        """从 JSON 构建省/市/区映射"""
        self._districts: List[Tuple[str, int]] = []  # (区县名, city_id)

        for prov in data:
            p_id = prov['id']
            p_name = prov['name']
            p_short = _strip_suffix(p_name, _PROVINCE_SUFFIXES)

            self._province_map[p_id] = {
                'id': p_id, 'name': p_name, 'short_name': p_short,
            }

            for city in prov.get('sons', []):
                c_id = city['id']
                c_name = city['name']
                c_short = _strip_suffix(c_name, _CITY_SUFFIXES)

                self._city_map[c_id] = {
                    'id': c_id, 'name': c_name, 'short_name': c_short,
                    'province_id': p_id,
                }

                # 收集区县
                for dist_name in city.get('sons', []):
                    self._districts.append((dist_name, c_id))

    def _build_indexes(self) -> None:
        """
        构建统一前缀树，三种插入模式：
          1. 省：从 root 插入（全称 + 简称）
          2. 市：从省叶节点继续插入（层级链）+ 从 root 直接插入（全局 fallback）
          3. 区：从 root 直接插入（全局 fallback）
        """
        # 预计算：直辖市的 p_id 集合（不插入省节点，直接插入市节点）
        municipality_pids: Set[int] = set()
        for c_id, city in self._city_map.items():
            p_id = city['province_id']
            if city['name'] == self._province_map[p_id]['name']:
                municipality_pids.add(p_id)

        # 1. 省级（全称 + 简称）
        # 直辖市不插入省节点，直接插入市节点（避免省/市同名导致 trie 停在 province）
        prov_leaves: Dict[int, List] = {}  # p_id → [叶节点]
        for p_id, prov in self._province_map.items():
            if p_id in municipality_pids:
                continue  # 直辖市跳过省插入
            leaves = self._area_trie.insert(prov['name'], 'province', p_id, 0)
            if prov['short_name'] and prov['short_name'] != prov['name'] and len(prov['short_name']) > 1:
                leaves += self._area_trie.insert(prov['short_name'], 'province', p_id, 0)
            prov_leaves[p_id] = leaves

        # 2. 市级
        city_leaves: Dict[int, List] = {}  # c_id → [叶节点]（用于区县层级链）
        for c_id, city in self._city_map.items():
            if city['name'] in ('市辖区',):
                continue
            p_id = city['province_id']
            c_name = city['name']
            c_short = city['short_name']

            # 直辖市：不经过省节点，直接从 root 插入市
            if p_id in municipality_pids:
                direct = self._area_trie.insert(c_name, 'city', c_id, c_id)
                city_leaves[c_id] = direct
                if c_short and c_short != c_name and len(c_short) > 1:
                    city_leaves[c_id] += self._area_trie.insert(c_short, 'city', c_id, c_id)
                # 虚拟城市别名
                if c_name in _VIRTUAL_CITY_ALIASES:
                    for alias in _VIRTUAL_CITY_ALIASES[c_name]:
                        city_leaves[c_id] += self._area_trie.insert(alias, 'city', c_id, c_id)
                continue

            # 普通城市：从省叶节点继续插入（层级链：省→市）
            if p_id in prov_leaves and c_name:
                chained = self._area_trie.insert_from(prov_leaves[p_id], c_name, 'city', c_id, c_id)
                city_leaves[c_id] = chained
                if c_short and c_short != c_name and len(c_short) > 1:
                    city_leaves[c_id] += self._area_trie.insert_from(prov_leaves[p_id], c_short, 'city', c_id, c_id)
                # 虚拟城市别名也从省叶节点链式插入
                if c_name in _VIRTUAL_CITY_ALIASES:
                    for alias in _VIRTUAL_CITY_ALIASES[c_name]:
                        city_leaves[c_id] += self._area_trie.insert_from(prov_leaves[p_id], alias, 'city', c_id, c_id)

            # 从 root 直接插入（全局 fallback）
            if c_name:
                direct = self._area_trie.insert(c_name, 'city', c_id, c_id)
                if c_id not in city_leaves:
                    city_leaves[c_id] = direct
                else:
                    city_leaves[c_id] += direct
            # 市简称（单字简称不插入，避免歧义：如"北林区"→"北"）
            if c_short and c_short != c_name and len(c_short) > 1:
                city_leaves[c_id] += self._area_trie.insert(c_short, 'city', c_id, c_id)

            # 虚拟城市别名
            if c_name in _VIRTUAL_CITY_ALIASES:
                for alias in _VIRTUAL_CITY_ALIASES[c_name]:
                    city_leaves[c_id] += self._area_trie.insert(alias, 'city', c_id, c_id)

        # 3. 区级（city_id = 所属市 c_id）
        # 3a. 从市叶节点继续插入（层级链：市→区）
        for dist_name, c_id in self._districts:
            if c_id in city_leaves:
                self._area_trie.insert_from(city_leaves[c_id], dist_name, 'district', c_id, c_id)
                d_short = _strip_suffix(dist_name, _DISTRICT_SUFFIXES)
                if d_short and d_short != dist_name and len(d_short) > 1:
                    self._area_trie.insert_from(city_leaves[c_id], d_short, 'district', c_id, c_id)
        # 3b. 从 root 直接插入（全局 fallback）
        for dist_name, c_id in self._districts:
            self._area_trie.insert(dist_name, 'district', c_id, c_id)
            d_short = _strip_suffix(dist_name, _DISTRICT_SUFFIXES)
            if d_short and d_short != dist_name and len(d_short) > 1:
                self._area_trie.insert(d_short, 'district', c_id, c_id)

    # ---------- 查询接口 ----------

    def _resolve(self, entry_type: str, entry_id: int) -> Optional[AreaMatch]:
        """将 trie 匹配结果转为 AreaMatch"""
        if entry_type == 'city':
            city = self._city_map[entry_id]
            prov = self._province_map.get(city['province_id'])
            return AreaMatch(
                prov['id'] if prov else 0, entry_id,
                prov['name'] if prov else '', city['name']
            )
        if entry_type == 'district':
            city = self._city_map.get(entry_id)
            if city:
                prov = self._province_map.get(city['province_id'])
                return AreaMatch(
                    prov['id'] if prov else 0, entry_id,
                    prov['name'] if prov else '', city['name']
                )
        return None

    def search_area(self, addr: str) -> Optional[AreaMatch]:
        """
        从地址中解析省市信息

        addr 进前缀树，出来就是结果。
        """
        if not addr:
            return None

        # 1. 手工配置的前缀（最高优先级）
        for prefix, city_id in self._addr_prefix_map.items():
            if addr.startswith(prefix):
                city = self._city_map.get(city_id)
                if city:
                    prov = self._province_map.get(city['province_id'])
                    return AreaMatch(
                        prov['id'] if prov else 0, city_id,
                        prov['name'] if prov else '', city['name']
                    )
                return AreaMatch(0, city_id, '', '')

        # 2. 统一前缀树匹配
        result = self._area_trie.search(addr)
        if not result:
            return None
        entry_type, entry_id = result
        # 仅匹配到省（非直辖市）→ 市未知
        if entry_type == 'province':
            return None
        return self._resolve(entry_type, entry_id)

    def get_area_info(self, area_id: int) -> Optional[Dict[str, Any]]:
        """
        根据区域 id 获取信息

        支持省级或市级 id。
        """
        area_id = int(area_id)
        city = self._city_map.get(area_id)
        if city:
            prov = self._province_map.get(city['province_id'])
            return {
                'id': area_id, 'name': city['name'], 'level': 'city',
                'province_id': city['province_id'],
                'province_name': prov['name'] if prov else '',
            }
        prov = self._province_map.get(area_id)
        if prov:
            return {'id': area_id, 'name': prov['name'], 'level': 'province'}
        return None

    def get_province_list(self) -> List[Dict[str, Any]]:
        """获取所有省份列表"""
        return [
            {'id': p['id'], 'name': p['name'], 'short_name': p['short_name']}
            for p in self._province_map.values()
        ]

    def get_city_list(self, province_id: Optional[int] = None) -> List[Dict[str, Any]]:
        """获取城市列表，可按省份过滤"""
        if province_id:
            return [
                {'id': c['id'], 'name': c['name'], 'short_name': c['short_name'],
                 'province_id': c['province_id']}
                for c in self._city_map.values() if c['province_id'] == province_id
            ]
        return [
            {'id': c['id'], 'name': c['name'], 'short_name': c['short_name'],
             'province_id': c['province_id']}
            for c in self._city_map.values()
        ]

    def get_district_list(self, city_id: Optional[int] = None) -> List[str]:
        """获取区县列表（字符串数组），可按城市 id 过滤"""
        if city_id is not None:
            return [name for name, cid in self._districts if cid == city_id]
        return [name for name, _ in self._districts]

    def batch_get_ids_by_names(self, pairs: List[Tuple[str, str]]) -> Dict[Tuple[str, str], Tuple[int, int]]:
        """
        批量通过 (province_name, city_name) 查询 (province_id, city_id)

        Args:
            pairs: [(province_name, city_name), ...]

        Returns:
            {(province_name, city_name): (province_id, city_id), ...}
            未找到的不会出现在结果中
        """
        if not pairs:
            return {}

        # 构建名称反查索引
        prov_name_to_id: Dict[str, int] = {
            prov['name']: pid for pid, prov in self._province_map.items()
        }
        city_name_to_id: Dict[Tuple[int, str], int] = {
            (city['province_id'], city['name']): cid
            for cid, city in self._city_map.items()
        }

        result: Dict[Tuple[str, str], Tuple[int, int]] = {}
        for p_name, c_name in pairs:
            if not p_name:
                continue
            p_id = prov_name_to_id.get(p_name)
            if p_id is None:
                continue
            c_id = city_name_to_id.get((p_id, c_name)) if c_name else 0
            result[(p_name, c_name)] = (p_id, c_id or 0)
        return result

    def _load_addr_prefix_map(self, json_path: str) -> None:
        """
        从 JSON 文件加载手工配置的地址前缀映射

        JSON 格式: [{"prefix": "xxx", "province": "省全称", "city": "市全称"}, ...]
        根据省全称和市全称查找对应的城市 id
        """
        json_path = os.path.normpath(json_path)

        with open(json_path, 'r', encoding='utf-8') as f:
            rules = json.load(f)

        # 构建名称反查 map
        prov_name_to_id: Dict[str, int] = {
            prov['name']: pid for pid, prov in self._province_map.items()
        }
        city_name_to_id: Dict[Tuple[int, str], int] = {
            (city['province_id'], city['name']): cid
            for cid, city in self._city_map.items()
        }

        self._addr_prefix_map = {}
        for rule in rules:
            prefix = rule.get('prefix', '')
            p_name = rule.get('province', '')
            c_name = rule.get('city', '')

            p_id = prov_name_to_id.get(p_name)
            if not p_id:
                log.warning(f'addr prefix rule: province not found: {p_name}')
                continue

            c_id = city_name_to_id.get((p_id, c_name))
            if not c_id:
                log.warning(f'addr prefix rule: city not found: {p_name} {c_name}')
                continue

            self._addr_prefix_map[prefix] = c_id

        log.info(f'loaded {len(self._addr_prefix_map)} addr prefix rules from {json_path}')