import os
import hashlib
import json
import argparse
import sys
from tqdm import tqdm  # 导入进度条库

def calculate_file_hash(file_path, block_size=65536):
    """计算文件的 SHA-256 哈希值"""
    sha256 = hashlib.sha256()
    try:
        with open(file_path, 'rb') as f:
            for block in iter(lambda: f.read(block_size), b''):
                sha256.update(block)
        return sha256.hexdigest()
    except (PermissionError, FileNotFoundError):
        return None

def get_all_files(folder_path):
    """预扫描文件夹，获取所有文件列表用于初始化进度条"""
    all_files = []
    for root, _, files in os.walk(folder_path):
        for name in files:
            all_files.append(os.path.join(root, name))
    return all_files

def record_folder_state(folder_path, output_file):
    """记录文件夹状态"""
    if not os.path.isdir(folder_path):
        print(f"错误: 路径 '{folder_path}' 不是文件夹。")
        return

    print(f"正在预扫描目录结构...")
    file_list = get_all_files(folder_path)
    file_data = {}

    print(f"开始计算哈希值:")
    
    for full_path in tqdm(file_list, desc="进度", unit="file"):
        rel_path = os.path.relpath(full_path, folder_path)
        file_hash = calculate_file_hash(full_path)
        if file_hash:
            file_data[rel_path] = file_hash
    
    with open(output_file, 'w', encoding='utf-8') as f:
        json.dump(file_data, f, indent=4, ensure_ascii=False)
    print(f"\n完成！已记录 {len(file_data)} 个文件。")

def compare_folder_state(current_folder, record_file):
    """对比差异"""
    if not os.path.exists(record_file):
        print(f"错误: 找不到记录文件 '{record_file}'。")
        return

    with open(record_file, 'r', encoding='utf-8') as f:
        old_data = json.load(f)

    print(f"正在预扫描当前目录...")
    file_list = get_all_files(current_folder)
    current_data = {}

    print(f"正在分析当前文件内容:")
    for full_path in tqdm(file_list, desc="对比中", unit="file"):
        rel_path = os.path.relpath(full_path, current_folder)
        file_hash = calculate_file_hash(full_path)
        if file_hash:
            current_data[rel_path] = file_hash

    old_files = set(old_data.keys())
    new_files = set(current_data.keys())
    added = sorted(list(new_files - old_files))
    removed = sorted(list(old_files - new_files))
    common = old_files & new_files
    changed = sorted([f for f in common if old_data[f] != current_data[f]])

    print(f"\n{'='*20} 对比报告 {'='*20}")
    print(f"新增 [{len(added)}], 删除 [{len(removed)}], 变动 [{len(changed)}]")
    if added:
        print("\n新增文件:")
        for f in added: print(f"  [+] {f}")
    if removed:
        print("\n删除文件:")
        for f in removed: print(f"  [-] {f}")
    if changed:
        print("\n变动文件:")
        for f in changed: print(f"  [*] {f}")
    
    if not any([added, removed, changed]):
        print("\n结果: 没有任何变化。")
    print(f"{'='*49}\n")


def main():
    parser = argparse.ArgumentParser(description="文件夹变更检测工具")
    subparsers = parser.add_subparsers(dest="command")

    save_parser = subparsers.add_parser("save", help="记录状态")
    save_parser.add_argument("-p", "--path", required=True)
    save_parser.add_argument("-o", "--out", default="snapshot.json")

    check_parser = subparsers.add_parser("check", help="对比差异")
    check_parser.add_argument("-p", "--path", required=True)
    check_parser.add_argument("-d", "--db", required=True)

    args = parser.parse_args()
    if args.command == "save":
        record_folder_state(args.path, args.out)
    elif args.command == "check":
        compare_folder_state(args.path, args.db)
    else:
        parser.print_help()

if __name__ == "__main__":
    main()