from pathlib import Path

import pandas
import pandas as pd

from utils import base_country_code, base_mysql
from utils.base_country_code import format_sql_value
from utils.log import log

city_code_map = {
    "石家庄市": "130100",
    "唐山市": "130200",
    "秦皇岛市": "130300",
    "邯郸市": "130400",
    "邢台市": "130500",
    "保定市": "130600",
    "张家口市": "130700",
    "承德市": "130800",
    "沧州市": "130900",
    "廊坊市": "131000",
    "衡水市": "131100",
}

def get_df(path):
    file_paths = list(Path(path).glob('*'))
    if not file_paths:
        log.info("未找到任何文件")
        return None
    for file in file_paths:
        if "地市" in file.name:
            log.info(f"处理多文件: {file.name}")
            file_path = Path(path) / file
            return pd.read_excel(file_path, header=None).iloc[5:]
    return None

def process_folder(path):
    year, month = base_country_code.extract_year_month_from_path(path)
    year_month = f'{year}-{month:02d}'

    df = get_df(path)
    if df is None:
        log.info("未找到任何文件")
        return None

    if year == 2025 and month >= 3:
        col_total_index, col_monthly_export_index, col_monthly_import_index = 2, 10, 18
    elif year_month in ['2023-02', '2025-01', '2024-01']:
        col_total_index, col_monthly_export_index, col_monthly_import_index = 1, 5, 9
    else:
        col_total_index, col_monthly_export_index, col_monthly_import_index = 1, 9, 17
    country_name_index = 1 if year == 2025 and month >= 3 else 0
    sql_arr = []
    sql_arr_copy = []
    for index, row in df.iterrows():

        city_name = str(row.values[country_name_index]).strip()
        if city_name.startswith('河北省'):
            city_name = city_name.lstrip('河北省')

        city_code = city_code_map.get(city_name)
        if not city_code:
            log.info(f"未找到省 '{city_name}' 对应市编码")
            continue

        monthly_export, monthly_import, monthly_total = value_row(row, col_total_index, col_monthly_export_index, col_monthly_import_index)
        yoy_export, yoy_import, yoy_import_export = 0, 0, 0
        if year_month == '2023-02':
            # 所有总额除2
            monthly_import = round(float(monthly_import) / 2, 4)
            monthly_export = round(float(monthly_export) / 2, 4)
            monthly_total = round(float(monthly_total) / 2, 4)
            sql_1 = (f"INSERT INTO t_yujin_crossborder_prov_region_trade "
                   f"(crossborder_year, crossborder_year_month, prov_code, prov_name, city_code, city_name, monthly_total, monthly_export, monthly_import,yoy_import_export, yoy_import, yoy_export, create_time) VALUES "
                   f"('2023', '2023-01', '130000', '河北省', '{city_code}', '{city_name}', {format_sql_value(monthly_total)}, {format_sql_value(monthly_export)}, {format_sql_value(monthly_import)}, '{yoy_import_export}', '{yoy_import}', '{yoy_export}', now()) ON DUPLICATE KEY UPDATE create_time = now() ;\n")
            sql_arr_copy.append(sql_1)

        # 组装 SQL 语句
        sql = (f"INSERT INTO t_yujin_crossborder_prov_region_trade "
               f"(crossborder_year, crossborder_year_month, prov_code, prov_name, city_code, city_name, monthly_total, monthly_export, monthly_import,yoy_import_export, yoy_import, yoy_export, create_time) VALUES "
               f"('{year}', '{year_month}', '130000', '河北省', '{city_code}', '{city_name}', {format_sql_value(monthly_total)}, {format_sql_value(monthly_export)}, {format_sql_value(monthly_import)}, '{yoy_import_export}', '{yoy_import}', '{yoy_export}', now()) ON DUPLICATE KEY UPDATE create_time = now() ;\n")
        sql_arr.append(sql)

    log.info(f"√ {year_month} prov_region_trade 成功生成 SQL 文件 size {len(sql_arr)} ")
    # 解析完后生成sql文件批量入库
    base_mysql.bulk_insert(sql_arr)
    if year_month == '2023-02':
        log.info(f"√ {year_month} sql_arr_copy 成功生成 SQL 文件 size {len(sql_arr_copy)} ")
        base_mysql.bulk_insert(sql_arr_copy)
    log.info(f"√ {year_month} prov_region_trade SQL 存表完成!")


def value_row(row, col_total_index, col_monthly_export_index, col_monthly_import_index):
    monthly_total = str(row.values[col_total_index]).strip()
    monthly_export = str(row.values[col_monthly_export_index]).strip()
    monthly_import = str(row.values[col_monthly_import_index]).strip()
    return monthly_export, monthly_import, monthly_total

def value_special_handler(value):
    if pandas.isna(value) or value == "--":
        return "0"
    else:
        return value

def hierarchical_traversal(root_path):
    root = Path(root_path)
    year_dirs = [
        item for item in root.iterdir()
        if item.is_dir() and base_country_code.YEAR_PATTERN.match(item.name)
    ]

    for year_dir in sorted(year_dirs, key=lambda x: x.name, reverse=True):
        log.info(f"\n年份：{year_dir.name} | 省份：hebei")

        month_dirs = []
        for item in year_dir.iterdir():
            if item.is_dir() and base_country_code.MONTH_PATTERN.match(item.name):
                month_dirs.append({"path": item, "month": int(item.name)})

        if month_dirs:
            for md in sorted(month_dirs, key=lambda x: x["month"], reverse=True):
                log.info(f"  月份：{md['month']:02d} | 路径：{md['path']}")
                process_folder(md['path'])

if __name__ == '__main__':
    hierarchical_traversal(base_country_code.download_dir)
    log.info(f"河北石家庄海关城市所有文件处理完成！")