import argparse
import pathlib
import pandas
import wx

parser = argparse.ArgumentParser(description="邮件列表切片脚本")
parser.add_argument('--column-pays', type=str, nargs='?', default='国家地区')
parser.add_argument('--column-code', type=str, nargs='?', default='客户编号')
parser.add_argument('-d', '--outdir', type=str, nargs='?')
parser.add_argument('-s', '--size', type=int, nargs='?', default=1000)

args = parser.parse_args()

mappings = [
    {
        "locale": "fr",
        "predicate": lambda it: it in ['法国', 'FR', '比利时', 'BE']
    },
    {
        "locale": "de",
        "predicate": lambda it: it in ['德国', 'DE', '奥地利', 'AT', '瑞士', 'CH']
    },
    {
        "locale": "it",
        "predicate": lambda it: it in ['意大利', 'IT']
    },
    {
        "locale": "es",
        "predicate": lambda it: it in ['西班牙', 'ES']
    },
    {
        "locale": "pt",
        "predicate": lambda it: it in ['葡萄牙', 'PT']
    },
    {
        "locale": "en",
        "predicate": lambda it: bool(it)
    },
    {
        "locale": "un",
        "predicate": lambda it: True
    },
]

def main():
    try:
        print('[信息] 请选择数据源')
        app = wx.App(None)
        dialog = wx.FileDialog(None, 'Open', wildcard='*.xlsx', style=wx.FD_OPEN|wx.FD_FILE_MUST_EXIST)

        if dialog.ShowModal() == wx.ID_OK:
            filepath = dialog.GetPath()
        else:
            print('[警告] 操作取消')
            return 0
        dialog.Destroy()
        
        if args.outdir is None:
            print('[信息] 请选择保存位置')
            dialog = wx.DirDialog(None, 'Save to', style=wx.DD_DIR_MUST_EXIST)

            if dialog.ShowModal() == wx.ID_OK:
                args.outdir = dialog.GetPath()
            else:
                print('[警告] 操作取消')
                return 0
            dialog.Destroy()

        print(f'[信息] 正在读取数据：{filepath}')
        sheets = pandas.read_excel(filepath, sheet_name=None)
        names = sheets.keys()

        parts = pathlib.Path(filepath)
        prefix = parts.stem
    except Exception as e:
        print(f'[警告] 读取数据表失败：{e}')
        return 1

    if (length := len(sheets)) > 1:
        print(f'[信息] 已读取以下 ({length}) 工作表：', end='\n\n')
        for name in names: print(f'\t{name}')

        while True:
            if (key := input('\n[????] 请输入工作表名称：')) in names:
                sheet = sheets.get(key)
                break
            else:
                print(f'[警告] 名称无效，请重试')
    else:
        print(f'[信息] 已读取工作表：{names[0]}')
        sheet = sheets.get(names[0])

    data = sheet.to_dict(orient='records')
    primary = {}
    secondary = {}

    for record in data:
        pays = record.get(args.column_pays)

        for rule in mappings:
            locale = rule.get('locale')
            predicate = rule.get('predicate')

            if predicate(pays):
                group: list = primary.setdefault(locale, [])
                group.append(record)
                break

    for key, group in primary.items():
        factor = 1
        offset = 0
        cursor = 0

        while (limit := args.size * factor + offset) and (slice := group[cursor:limit]):
            last: dict = slice[-1:][0]
            next: dict = group[limit] if len(group) > limit else None

            if next is not None and (c1 := last.get(args.column_code)) and (c2 := next.get(args.column_code)) and c1 == c2:
                offset += 1
                continue

            name = '.'.join([key, str(factor)])
            secondary.setdefault(name, slice)
            factor += 1
            cursor = limit
            offset = 0

    for key, group in secondary.items():
        try:
            filepath = f'{args.outdir}/{prefix}.{key}.xlsx'
            df = pandas.DataFrame.from_records(group)

            print(f'[信息] 正在写入文件：{filepath}')
            df.to_excel(filepath, index=False, sheet_name='Sheet1')
        except Exception as e:
            print(f'[警告] 写入文件时发生了错误：{e}')
            continue

    return 0

try: status = main()
except KeyboardInterrupt: status = 144
exit(status)