#!/usr/bin/env python3
"""
股票代码格式化工具
根据股票代码自动判断市场并添加后缀 (SZ/SH/BJ)
"""


def get_stock_exchange(stock_code: str) -> str:
    """
    根据股票代码判断所属交易所
    
    Args:
        stock_code: 纯数字股票代码，如 "000001", "600000"
    
    Returns:
        交易所缩写: "SZ" | "SH" | "BJ" | None
    """
    if not stock_code or not stock_code.isdigit():
        return None
    
    code = stock_code.strip()
    
    # 沪市 (SH)
    # - 600, 601, 603, 605: 主板
    # - 688, 689: 科创板
    # - 900: B股
    if code.startswith(('600', '601', '603', '605', '688', '689', '900')):
        return "SH"
    
    # 深市 (SZ)
    # - 000, 001: 主板
    # - 002, 003: 中小板/合并后主板
    # - 300: 创业板
    # - 200: B股
    if code.startswith(('000', '001', '002', '003', '300', '200')):
        return "SZ"
    
    # 北交所 (BJ)
    # - 8 开头: 43, 83, 87, 88 等
    if code.startswith('8') or code.startswith('4'):
        return "BJ"
    
    return None


def format_stock_code(stock_code: str) -> str:
    """
    将纯数字股票代码格式化为 代码.市场 格式
    
    Args:
        stock_code: 纯数字股票代码，如 "000001"
    
    Returns:
        格式化后的代码，如 "000001.SZ"
    
    Raises:
        ValueError: 无法识别市场时抛出
    """
    code = stock_code.strip()
    
    exchange = get_stock_exchange(code)
    if exchange is None:
        raise ValueError(f"无法识别的股票代码: {stock_code}")
    
    return f"{code}.{exchange}"


def format_stock_codes(stock_codes: list[str]) -> list[str]:
    """
    批量格式化股票代码
    
    Args:
        stock_codes: 股票代码列表
    
    Returns:
        格式化后的代码列表
    """
    results = []
    for code in stock_codes:
        try:
            formatted = format_stock_code(code)
            results.append(formatted)
        except ValueError as e:
            results.append(f"ERROR: {e}")
    return results


if __name__ == "__main__":
    # 测试示例
    test_codes = [
        "000001",   # 平安银行 - 深市
        "600000",   # 浦发银行 - 沪市
        "688981",   # 中芯国际 - 沪市(科创板)
        "300750",   # 宁德时代 - 深市(创业板)
        "835185",   # 贝特瑞 - 北交所
        "12345",    # 无效代码
    ]
    
    print("股票代码格式化测试:")
    print("-" * 30)
    for code in test_codes:
        try:
            result = format_stock_code(code)
            print(f"{code} -> {result}")
        except ValueError as e:
            print(f"{code} -> {e}")