#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import argparse
import sys
from pathlib import Path
from typing import Set, Dict, Tuple 

# <--- 新增: 导入 Pillow 库用于图像处理 ---
try:
    from PIL import Image
    # 禁用解压缩炸弹检查，以防万一标签图非常大且简单
    Image.MAX_IMAGE_PIXELS = None
except ImportError:
    print("错误: 本脚本需要 'Pillow' 库来检查图像尺寸。", file=sys.stderr)
    print("请运行: pip install Pillow", file=sys.stderr)
    sys.exit(1)
# --- 新增结束 ---


# 定义支持的文件扩展名
VALID_EXTENSIONS = {'.png', '.jpg', '.jpeg'}

def process_directory(path: Path, 
                      prefix: str = "", 
                      suffix: str = "") -> Tuple[Set[str], Dict[str, Path], Dict[Tuple[int, int], int]]: # <--- 修改：更新了返回类型
    """
    处理单个目录，提取所有有效文件的文件名（stem），并根据前缀和后缀进行规范化。
    同时统计所有图像的尺寸。

    返回:
        normalized_stems (Set[str]): 规范化处理后的文件名集合。
        normalized_to_full_path (Dict[str, Path]): 规范化文件名 -> 原始完整路径 的映射。
        size_counts (Dict[Tuple[int, int], int]): (宽度, 高度) -> 数量 的映射。 # <--- 新增
    """
    if not path.is_dir():
        print(f"错误: 路径 '{path}' 不是一个有效的目录。", file=sys.stderr)
        return set(), {}, {} # <--- 修改

    normalized_stems: Set[str] = set()
    normalized_to_full_path: Dict[str, Path] = {} 
    size_counts: Dict[Tuple[int, int], int] = {} # <--- 新增：用于统计尺寸

    for file_path in path.glob('*'):
        # 确保是文件，并且扩展名在我们的有效列表中
        if file_path.is_file() and file_path.suffix.lower() in VALID_EXTENSIONS:
            
            # <--- 新增: 尝试获取图像尺寸 ---
            try:
                # 使用 with...as... 确保文件被正确关闭
                with Image.open(file_path) as img:
                    size: Tuple[int, int] = img.size
                    size_counts[size] = size_counts.get(size, 0) + 1
            except Exception as e:
                print(f"警告: 无法读取图像 '{file_path.name}' 的尺寸. 错误: {e}", file=sys.stderr)
            # --- 尺寸获取结束 ---

            original_stem = file_path.stem
            normalized_stem = original_stem

            # 仅当提供了前缀/后缀时才进行处理
            if prefix and normalized_stem.startswith(prefix):
                normalized_stem = normalized_stem[len(prefix):]
            
            if suffix and normalized_stem.endswith(suffix):
                normalized_stem = normalized_stem[:-len(suffix)]

            # 检查处理后是否重名，如果重名则发出警告
            if normalized_stem in normalized_to_full_path: 
                print(f"警告: 规范化后文件名发生冲突。")
                print(f"  '{file_path.name}' 和 '{normalized_to_full_path[normalized_stem].name}' 都变成了 '{normalized_stem}'") 
            
            normalized_stems.add(normalized_stem)
            normalized_to_full_path[normalized_stem] = file_path 

    return normalized_stems, normalized_to_full_path, size_counts # <--- 修改


# <--- 新增: 打印尺寸统计的辅助函数 ---
def print_size_report(title: str, size_counts: Dict[Tuple[int, int], int]):
    """
    格式化并打印尺寸统计报告。
    """
    print(f"\n{title}")
    if not size_counts:
        print("  (未找到或无法读取任何图像文件)")
        return
    
    total_files = sum(size_counts.values())
    print(f"  总共 {total_files} 个文件，分布如下:")
    
    # 按尺寸 (宽, 高) 排序输出
    for size, count in sorted(size_counts.items()):
        width, height = size
        print(f"  - 尺寸 (宽, 高) {width}x{height}: {count} 个文件")
# --- 新增结束 ---


def main():
    parser = argparse.ArgumentParser(description="比较两个文件夹中的文件名是否匹配，并可选择删除不匹配项。") 
    parser.add_argument("-i", "--image", 
                        type=Path, 
                        default="./ORI", 
                        help="Image 文件夹路径")
    parser.add_argument("-l", "--label", 
                        type=Path, 
                        default="./Label", 
                        help="Label 文件夹路径")
    parser.add_argument("-p", "--prefix", 
                        type=str, 
                        default="", 
                        help="在 Label 文件名中要忽略的前缀")
    parser.add_argument("-s", "--suffix", 
                        type=str, 
                        default="", 
                        help="在 Label 文件名中要忽略的后缀")
    parser.add_argument("-y", "--yes",
                        action="store_true",
                        help="自动确认删除所有不匹配的文件，跳过交互式提示。")

    args = parser.parse_args()

    # 1. 处理 Image 文件夹
    print(f"--- 正在处理 Image 文件夹: {args.image} ---")
    image_stems, image_norm_to_path, image_size_counts = process_directory(args.image) # <--- 修改
    if not image_stems:
        print(f"未在 Image 文件夹中找到任何 .png 或 .jpg 文件。")
    # <--- 新增: 打印 Image 尺寸报告 ---
    print_size_report("--- Image 文件夹尺寸统计 ---", image_size_counts)

    # 2. 处理 Label 文件夹
    print(f"\n--- 正在处理 Label 文件夹: {args.label} ---")
    print(f"(忽略前缀: '{args.prefix}', 忽略后缀: '{args.suffix}')")
    label_stems_normalized, label_norm_to_path, label_size_counts = process_directory( # <--- 修改
        args.label, 
        args.prefix, 
        args.suffix
    )
    if not label_stems_normalized:
        print(f"未在 Label 文件夹中找到任何 .png 或 .jpg 文件。")
    # <--- 新增: 打印 Label 尺寸报告 ---
    print_size_report("--- Label 文件夹尺寸统计 ---", label_size_counts)


    # 3. 执行比较 (使用集合运算)
    matching_stems = image_stems.intersection(label_stems_normalized)
    extra_in_image = image_stems.difference(label_stems_normalized)
    extra_in_label_normalized = label_stems_normalized.difference(image_stems)
    extra_in_label_original_stems = {label_norm_to_path[stem].stem for stem in extra_in_label_normalized} 

    # 4. 输出结果
    print("\n" + "="*30)
    print("       匹配结果报告")
    print("="*30)

    print(f"\n匹配的文件总数: {len(matching_stems)}")

    # <--- 新增: 检查匹配项的尺寸是否一致 ---
    print("\n--- 正在检查匹配文件的尺寸 ---")
    mismatched_size_pairs = []
    if not matching_stems:
        print("(无匹配文件，跳过尺寸检查)")
    else:
        print(f"正在检查 {len(matching_stems)} 对文件...")
        for stem in sorted(matching_stems):
            try:
                img_path = image_norm_to_path[stem]
                lbl_path = label_norm_to_path[stem]
                
                # 再次打开以比较尺寸
                with Image.open(img_path) as img:
                    img_size = img.size
                with Image.open(lbl_path) as lbl:
                    lbl_size = lbl.size
                
                if img_size != lbl_size:
                    mismatched_size_pairs.append((img_path, img_size, lbl_path, lbl_size))
                    
            except Exception as e:
                print(f"  错误: 无法比较 '{stem}' 的尺寸. 错误: {e}", file=sys.stderr)

        if not mismatched_size_pairs:
            print(f"√ 成功: 所有 {len(matching_stems)} 对匹配文件均具有相同的尺寸。")
        else:
            print(f"\n!!! 警告: 发现 {len(mismatched_size_pairs)} 对文件尺寸不匹配:")
            for img_path, img_size, lbl_path, lbl_size in mismatched_size_pairs:
                print(f"  - [Image] {img_path.name} {img_size} != [Label] {lbl_path.name} {lbl_size}")
    # --- 尺寸检查结束 ---

    print("\n--- Image 文件夹中多余的文件 (共 {} 个) ---".format(len(extra_in_image)))
    if not extra_in_image:
        print("(无)")
    else:
        for file_stem in sorted(extra_in_image):
            print(file_stem)

    print("\n--- Label 文件夹中多余的文件 (共 {} 个) ---".format(len(extra_in_label_original_stems))) 
    print("(显示的是原始文件名，非规范化名称)")
    if not extra_in_label_original_stems: 
        print("(无)")
    else:
        for file_stem in sorted(extra_in_label_original_stems): 
            print(file_stem)
    
    print("\n" + "="*30)

    # 5. 删除交互部分 (无修改)
    if not extra_in_image and not extra_in_label_normalized:
        print("\n所有文件均完美匹配，无需删除。")
        # <--- 新增: 如果文件名匹配，但尺寸不匹配，也在这里提醒一下 ---
        if mismatched_size_pairs:
             print("注意：虽然文件名匹配，但有 {} 对文件的尺寸不一致，请检查上面的报告。".format(len(mismatched_size_pairs)))
        return

    # ... (后续的删除逻辑保持不变) ...
    confirm_delete = False
    if args.yes:
        print("\n检测到 -y/--yes 参数，将自动删除不匹配的文件。")
        confirm_delete = True
    else:
        try:
            choice = input("\n警告: 是否要永久删除所有上述 '多余' 的文件? (y/n): ").strip().lower()
            if choice in ['y', 'yes']:
                confirm_delete = True
        except (KeyboardInterrupt, EOFError):
            print("\n操作被用户中断。")
            confirm_delete = False

    if confirm_delete:
        print("\n--- 正在删除 Image 文件夹中的多余文件 ---")
        deleted_count = 0
        for stem in sorted(extra_in_image):
            file_path = image_norm_to_path[stem]
            try:
                file_path.unlink()
                print(f"  已删除: {file_path.name}")
                deleted_count += 1
            except OSError as e:
                print(f"  删除失败: {file_path.name} (错误: {e})", file=sys.stderr)
        print(f"--- Image 文件夹共删除 {deleted_count} 个文件 ---")

        print("\n--- 正在删除 Label 文件夹中的多余文件 ---")
        deleted_count = 0
        for norm_stem in sorted(extra_in_label_normalized):
            file_path = label_norm_to_path[norm_stem]
            try:
                file_path.unlink()
                print(f"  已删除: {file_path.name} (原始文件名)")
                deleted_count += 1
            except OSError as e:
                print(f"  删除失败: {file_path.name} (错误: {e})", file=sys.stderr)
        print(f"--- Label 文件夹共删除 {deleted_count} 个文件 ---")
        
        print("\n删除操作完成。")
    
    else:
        print("\n操作已取消。未删除任何文件。")


if __name__ == "__main__":
    main()