import os
import argparse
import frontmatter
from frontmatter.default_handlers import TOMLHandler

def process_files(content_dir, action, target, new_name=None):
    handler = TOMLHandler()
    count = 0

    for root, dirs, files in os.walk(content_dir):
        for file in files:
            if file.endswith('.md'):
                path = os.path.join(root, file)
                try:
                    post = frontmatter.load(path, handler=handler)
                    modified = False

                    # 获取 categories（通常是字符串）
                    current_cat = post.get('categories')

                    if action == 'delete' and current_cat == target:
                        # 直接从字典中删除该键
                        del post['categories']
                        modified = True

                    elif action == 'rename' and current_cat == target:
                        # 更新为新的字符串
                        post['categories'] = new_name
                        modified = True

                    if modified:
                        with open(path, 'wb') as f:
                            frontmatter.dump(post, f, handler=handler)
                        print(f"✅ 已处理: {path}")
                        count += 1

                except Exception as e:
                    print(f"⚠️ 跳过/错误 {path}: {e}")

    print(f"\n✨ 完成！修改了 {count} 个文件。")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Hugo TOML 单分类管理工具")
    parser.add_argument("--dir", default="content", help="Hugo content 目录")

    group = parser.add_mutually_exclusive_group(required=True)
    group.add_argument("--delete", metavar="分类名")
    group.add_argument("--rename", nargs=2, metavar=("旧分类", "新分类"))

    args = parser.parse_args()

    if args.delete:
        process_files(args.dir, 'delete', args.delete)
    elif args.rename:
        process_files(args.dir, 'rename', args.rename[0], args.rename[1])