#!/usr/bin/env python3
"""
根据 classification.json 更新所有文档的 frontmatter
"""

import json
import os
import re
import sys
from pathlib import Path

def load_classification(vault_path):
    """加载分类结果"""
    with open(os.path.join(vault_path, 'classification.json'), 'r', encoding='utf-8') as f:
        return json.load(f)

def load_migration_plan(vault_path):
    """加载迁移计划"""
    with open(os.path.join(vault_path, 'migration-plan.json'), 'r', encoding='utf-8') as f:
        return json.load(f)

def update_frontmatter(file_path, platform, device):
    """更新文件的 frontmatter"""
    with open(file_path, 'r', encoding='utf-8') as f:
        content = f.read()
    
    # 检查是否有 frontmatter
    if not content.startswith('---'):
        # 添加新的 frontmatter
        new_content = f'---\nplatform: {platform}\ndevice: {device}\n---\n\n{content}'
        with open(file_path, 'w', encoding='utf-8') as f:
            f.write(new_content)
        return True
    
    # 解析现有 frontmatter
    parts = content.split('---', 2)
    if len(parts) < 3:
        return False
    
    frontmatter = parts[1]
    body = parts[2]
    
    # 更新或添加 platform
    if 'platform:' in frontmatter:
        frontmatter = re.sub(r'platform:.*', f'platform: {platform}', frontmatter)
    else:
        frontmatter += f'\nplatform: {platform}'
    
    # 更新或添加 device
    if 'device:' in frontmatter:
        frontmatter = re.sub(r'device:.*', f'device: {device}', frontmatter)
    else:
        frontmatter += f'\ndevice: {device}'
    
    # 写回文件
    new_content = f'---{frontmatter}---{body}'
    with open(file_path, 'w', encoding='utf-8') as f:
        f.write(new_content)
    
    return True

def main():
    vault_path = sys.argv[1] if len(sys.argv) > 1 else '.'
    
    # 加载分类结果和迁移计划
    classification = load_classification(vault_path)
    migration_plan = load_migration_plan(vault_path)
    
    # 创建源路径到分类的映射
    source_to_class = {}
    for item in classification['items']:
        source_to_class[item['path']] = {
            'platform': item['platform'],
            'device': item['device']
        }
    
    # 更新所有迁移的文件
    updated = 0
    skipped = 0
    
    for migration in migration_plan['migrations']:
        target_path = os.path.join(vault_path, migration['target'])
        source_path = migration['source']
        
        if not os.path.exists(target_path):
            print(f'跳过（文件不存在）: {target_path}')
            skipped += 1
            continue
        
        # 获取分类信息
        if source_path not in source_to_class:
            print(f'跳过（无分类信息）: {source_path}')
            skipped += 1
            continue
        
        class_info = source_to_class[source_path]
        
        # 更新 frontmatter
        if update_frontmatter(target_path, class_info['platform'], class_info['device']):
            print(f'已更新: {migration["target"]}')
            updated += 1
        else:
            print(f'更新失败: {migration["target"]}')
            skipped += 1
    
    print(f'\n更新完成: {updated} 个文件已更新, {skipped} 个文件跳过')

if __name__ == '__main__':
    main()