import pandas as pd
from typing import List, Optional, Dict, Tuple
import os

class AreaCodeTool:
    def __init__(self, csv_path: str = None):
        """
        初始化行政区划代码工具
        
        Args:
            csv_path: CSV文件路径，如果为None则使用默认路径
        """
        if csv_path is None:
            # 获取当前文件所在目录
            current_dir = os.path.dirname(os.path.abspath(__file__))
            csv_path = os.path.join(current_dir, "area_code.csv")
        
        # 读取CSV文件
        self.df = pd.read_csv(csv_path, dtype={'code': str})
        # 确保code列为字符串类型
        self.df['code'] = self.df['code'].astype(str)
        # 构建区域名称到代码的映射
        self._build_name_maps()

    def _build_name_maps(self):
        """构建区域名称到代码的映射"""
        self.full_name_map = dict(zip(self.df['name'], self.df['code']))
        
        # 构建省级映射
        self.province_map = {}
        # 构建市级映射
        self.city_map = {}
        # 构建区县级映射
        self.district_map = {}
        
        for _, row in self.df.iterrows():
            name = row['name'].strip()
            code = row['code']
            parts = name.split('省' if '省' in name else '市')
            
            if '省' in name:
                province = parts[0] + '省'
                self.province_map[province] = code
                
                if len(parts) > 1 and parts[1]:
                    city_parts = parts[1].split('市')
                    if city_parts[0]:
                        city = city_parts[0] + '市'
                        self.city_map[city] = code
                        
                        if len(city_parts) > 1 and city_parts[1]:
                            district = city_parts[1]
                            self.district_map[district] = code
            else:
                # 处理直辖市等特殊情况
                if parts[0]:
                    self.city_map[parts[0] + '市'] = code

    def find_code(self, area_name: str) -> List[Tuple[str, str]]:
        """
        查找区域代码
        
        Args:
            area_name: 区域名称，可以是完整或部分名称
            
        Returns:
            List[Tuple[str, str]]: 返回匹配的(区域名称, 代码)列表
        """
        results = []
        
        # 尝试完整匹配
        if area_name in self.full_name_map:
            results.append((area_name, self.full_name_map[area_name]))
            return results
            
        # 尝试省级匹配
        if area_name.endswith('省') and area_name in self.province_map:
            results.append((area_name, self.province_map[area_name]))
        
        # 尝试市级匹配
        if area_name.endswith('市') and area_name in self.city_map:
            results.append((area_name, self.city_map[area_name]))
            
        # 尝试区县级匹配
        if area_name in self.district_map:
            results.append((area_name, self.district_map[area_name]))
            
        # 模糊匹配
        if not results:
            mask = self.df['name'].str.contains(area_name, na=False)
            matches = self.df[mask]
            results.extend([(row['name'], row['code']) for _, row in matches.iterrows()])
            
        return results

    def get_full_name(self, code: str) -> Optional[str]:
        """
        根据代码获取完整的区域名称
        
        Args:
            code: 区域代码
            
        Returns:
            Optional[str]: 完整的区域名称，如果未找到则返回None
        """
        mask = self.df['code'] == code
        matches = self.df[mask]
        if not matches.empty:
            return matches.iloc[0]['name']
        return None

# 使用示例
def example_usage():
    tool = AreaCodeTool()
    
    # 测试不同类型的查询
    test_cases = [
        "安徽省",
        "安庆市",
        "迎江区",
        "安徽省安庆市",
        "安徽省安庆市迎江区",
        "安庆"  # 模糊查询
    ]
    
    for query in test_cases:
        print(f"\n查询: {query}")
        results = tool.find_code(query)
        for name, code in results:
            print(f"匹配结果: {name} -> {code}")
            
    # 测试代码反查
    code = "340802"  # 安徽省安庆市迎江区
    full_name = tool.get_full_name(code)
    if full_name:
        print(f"\n代码反查: {code} -> {full_name}")

if __name__ == "__main__":
    example_usage()