from pathlib import Path

import pandas
import pandas as pd

from com.zf.crawl import base_country_code
from com.zf.crawl import base_mysql

city_code_map = {
    "安徽省合肥市": "3401",
    "安徽省芜湖市": "3402",
    "安徽省蚌埠市": "3403",
    "安徽省淮南市": "3404",
    "安徽省马鞍山市": "3405",
    "安徽省淮北市": "3406",
    "安徽省铜陵市": "3407",
    "安徽省安庆市": "3408",
    "安徽省黄山市": "3410",
    "安徽省滁州市": "3411",
    "安徽省阜阳市": "3412",
    "安徽省宿州市": "3413",
    "安徽省六安市": "3415",
    "安徽省亳州市": "3416",
    "安徽省池州市": "3417",
    "安徽省宣城市": "3418"
}

def get_df(path):
    global df,  df_type
    file_paths = list(Path(path).glob('*'))
    if not file_paths:
        print("未找到任何文件")
        return
    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[5:]
            break
    return df

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

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

    for index, row in df.iterrows():

        city_name = str(row.values[country_name_index]).strip()
        city_code = city_code_map.get(city_name)
        if not city_code:
            print(f"未找到省 '{city_name}' 对应市编码")
            continue

        monthly_export, monthly_import, monthly_total, yoy_export, yoy_import, yoy_import_export = value_row(row, month)

        if month == 2:
            # 所有总额除2
            year_month_2 = f'{year}-01'
            monthly_import = round(float(monthly_import) / 2, 4)
            monthly_export = round(float(monthly_export) / 2, 4)
            monthly_total = round(float(monthly_total) / 2, 4)
            yoy_import_export, yoy_import, yoy_export = 0, 0, 0
            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_2}', '340000', '安徽省', '{city_code}', '{city_name}', '{monthly_total}', '{monthly_export}', '{monthly_import}', '{yoy_import_export}', '{yoy_import}', '{yoy_export}', now());\n")
            sql_arr_copy.append(sql)

        # 组装 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}', '340000', '安徽省', '{city_code}', '{city_name}', '{monthly_total}', '{monthly_export}', '{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)
    if month == 2:
        print(f"√ {year_month} prov_region_trade 成功生成 SQL 文件 size {len(sql_arr_copy)} ")
        base_mysql.bulk_insert(sql_arr_copy)
    print(f"√ {year_month} prov_region_trade SQL 存表完成!")


def value_row(row, month):
    if month == 2:
        col_total_index, col_export_index,  col_import_index, yoy_import_export_index, yoy_export_index, yoy_import_index = 1, 3, 5 ,2, 4, 6
    else:
        col_total_index, col_export_index,  col_import_index, yoy_import_export_index, yoy_export_index, yoy_import_index  = 3, 7, 11, 4, 8, 12
    monthly_total = round(float(row.values[col_total_index]) * 10000, 4)
    yoy_import_export = float(str(row.values[yoy_import_export_index]).strip())
    monthly_export = round(float(row.values[col_export_index]) * 10000, 4)
    yoy_export = float(str(row.values[yoy_export_index]).strip())
    monthly_import = round(float(row.values[col_import_index]) * 10000, 4)
    yoy_import = float(str(row.values[yoy_import_index]).strip())
    return monthly_export, monthly_import, monthly_total, yoy_export, yoy_import, yoy_import_export


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"安徽合肥海关城市所有文件处理完成！")