import os
import shutil
import argparse
import re

def flatten_directory(source_dir):
    """
    将指定目录下的所有txt文件移动到当前目录并重命名

    Args:
        source_dir: 源目录
    """

    for root, dirs, files in os.walk(source_dir):
        for file in files:
            if file.endswith('.txt'):
                src_file = os.path.join(root, file)
                dst_file = os.path.join(source_dir, f"[{os.path.basename(root)}]_{file}")
                print(f'move {src_file} {dst_file}')
                shutil.move(src_file, dst_file)

def unflatten_directory(source_dir):
    """
    将当前目录下的txt文件按照文件名中的目录信息进行分类

    Args:
        source_dir: 源目录
    """

    for file in os.listdir(source_dir):
        if file.endswith('.txt'):
            #dir_name, filename = file.split(']', 1)[0][1:], file.split(']', 1)[1][1:]
            # 方法二：使用正则表达式
            match = re.match(r"\[(.*)]_(.*)", file)
            if match:
                dir_name, filename = match.groups()

            dst_dir = os.path.join(source_dir, dir_name)
            dst_file = os.path.join(dst_dir, filename)
            src_file = os.path.join(source_dir, file)

            # 创建目标目录
            os.makedirs(dst_dir, exist_ok=True)

            # 移动文件
            print(f'move {src_file} {dst_file}')
            shutil.move(src_file, dst_file)

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Flatten or unflatten a directory of txt files')
    parser.add_argument('-f', '--flatten', action='store_true', help='Flatten the directory')
    parser.add_argument('-u', '--unflatten', action='store_true', help='Unflatten the directory')
    parser.add_argument('directory', help='The directory to process')
    args = parser.parse_args()

    if args.flatten and args.unflatten:
        print("Please choose either --flatten or --unflatten, not both.")
        exit(1)

    if args.flatten:
        flatten_directory(args.directory)
    elif args.unflatten:
        unflatten_directory(args.directory)
    else:
        print("Please specify either --flatten or --unflatten.")
        exit(1)