import time
from pathlib import Path

import pandas as pd

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

city_code_map = {
    "南京市": "3201",
    "无锡市": "3202",
    "徐州市": "3203",
    "常州市": "3204",
    "苏州市": "3205",
    "南通市": "3206",
    "连云港市": "3207",
    "淮安市": "3208",
    "盐城市": "3209",
    "扬州市": "3210",
    "镇江市": "3211",
    "泰州市": "3212",
    "宿迁市": "3213"
}

ignore_city_code_arr = ['江阴市','宜兴市','常熟市','张家港市','昆山市','吴江市','太仓市','启东市','东台市','仪征市','丹阳市','兴化市']

def get_df(path):
    global df,  df_type
    file_paths = list(Path(path).glob('*'))
    if not file_paths:
        print("未找到任何文件")
        return
    if len(file_paths) == 1:
        file_path = file_paths[0]
        print(f"处理单文件: {file_path.name}")
        xls = pd.ExcelFile(file_path)

        sheet_name = base_country_code.find_sheet_by_keyword(file_path, "地")
        if not sheet_name:
            print(f"{file_path} 未找到包含 地市 sheet")
            return None
        df = pd.read_excel(xls, sheet_name=sheet_name, header=None).iloc[5:]
        df_type = 0

    else:
        for file in file_paths:
            if "地区" in file.name:
                print(f"处理多文件: {file.name}")
                file_path = Path(path) / file
                df = pd.read_excel(file_path, header=None).iloc[6:]
                df_type = 1
                break
    return df, df_type

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

    sql_arr = []
    res = get_df(path)
    if res is None:
        print(f"{year_month} prov_region_trade 未找到包含 地市 sheet")
        return
    df, df_type = res
    if df_type == 0:
        country_name_index = 0
        col_total_index = 1
    else:
        country_name_index = 1
        col_total_index  = 2

    for index, row in df.iterrows():
        city_name = str(row.values[country_name_index]).strip()
        flag = False
        for ignore_city_code in ignore_city_code_arr:
            if city_name.startswith('其中') or ignore_city_code.endswith(city_name):
                flag = True
                break
        if flag:
            print(f"忽略 {city_name}")
            continue
        city_code = city_code_map.get(city_name)
        if not city_code:
            print(f"未找到省 '{city_name}' 对应市编码")
            continue

        monthly_export, monthly_import, monthly_total = value_row(row, col_total_index)
        if df_type == 0:
            monthly_export, monthly_import, monthly_total = round(float(monthly_export) * 10000, 4), round(float(monthly_import) * 10000, 4), round(float(monthly_total) * 10000, 4)
        yoy_export, yoy_import, yoy_import_export = 0, 0, 0

        # 组装 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}', '320000', '江苏省', '{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());\n")
        sql_arr.append(sql)

    print(f"√ {year_month} prov_region_trade 成功生成 SQL 文件 size {len(sql_arr)} ")
    # 解析完后生成sql文件批量入库
    base_mysql.bulk_insert(sql_arr)
    print(f"√ {year_month} prov_region_trade SQL 存表完成!")


def value_row(row,col_total_index):
    monthly_total = str(row.values[col_total_index]).strip()
    monthly_export = str(row.values[col_total_index + 2]).strip()
    monthly_import = str(row.values[col_total_index + 4]).strip()
    return monthly_export, monthly_import, monthly_total

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):
        print(f"\n年份：{year_dir.name} | 省份：jiangsu")

        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):
                print(f"  月份：{md['month']:02d} | 路径：{md['path']}")
                process_folder(md['path'])

if __name__ == '__main__':
    hierarchical_traversal(base_country_code.download_dir)
    print(f"江苏南京海关城市所有文件处理完成！")
    time.sleep(5)
    base_mysql.update_january_yoy('江苏省')
    base_mysql.update_shandong_yoy('江苏省')
    print("江苏南京同比sql处理完成")