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

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

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

def process_directory(path: Path, 
                      prefix: str = "", 
                      suffix: str = "") -> (Set[str], Dict[str, str]):
    """
    处理单个目录，提取所有有效文件的文件名（stem），并根据前缀和后缀进行规范化。

    返回:
        normalized_stems (Set[str]): 规范化处理后的文件名集合。
        original_to_normalized_map (Dict[str, str]): 原始文件名到规范化文件名的映射。
                      （这里我们反过来，用 规范化 -> 原始，更方便后续查找）
    """
    if not path.is_dir():
        print(f"错误: 路径 '{path}' 不是一个有效的目录。", file=sys.stderr)
        return set(), {}

    normalized_stems: Set[str] = set()
    normalized_to_original_map: Dict[str, str] = {}

    for file_path in path.glob('*'):
        # 确保是文件，并且扩展名在我们的有效列表中
        if file_path.is_file() and file_path.suffix.lower() in VALID_EXTENSIONS:
            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_original_map:
                print(f"警告: 规范化后文件名发生冲突。")
                print(f"  '{original_stem}' 和 '{normalized_to_original_map[normalized_stem]}' 都变成了 '{normalized_stem}'")
            
            normalized_stems.add(normalized_stem)
            normalized_to_original_map[normalized_stem] = original_stem

    return normalized_stems, normalized_to_original_map


def main():
    parser = argparse.ArgumentParser(description="比较两个文件夹中的文件名是否匹配。")
    parser.add_argument("-i", "--image", 
                        type=Path, 
                        required=True, 
                        help="Image 文件夹路径")
    parser.add_argument("-l", "--label", 
                        type=Path, 
                        required=True, 
                        help="Label 文件夹路径")
    parser.add_argument("-p", "--prefix", 
                        type=str, 
                        default="", 
                        help="在 Label 文件名中要忽略的前缀")
    parser.add_argument("-s", "--suffix", 
                        type=str, 
                        default="", 
                        help="在 Label 文件名中要忽略的后缀")

    args = parser.parse_args()

    # 1. 处理 Image 文件夹 (不需要前缀后缀)
    print(f"--- 正在处理 Image 文件夹: {args.image} ---")
    image_stems, _ = process_directory(args.image)
    if not image_stems:
        print(f"未在 Image 文件夹中找到任何 .png 或 .jpg 文件。")

    # 2. 处理 Label 文件夹 (需要前缀后缀)
    print(f"\n--- 正在处理 Label 文件夹: {args.label} ---")
    print(f"(忽略前缀: '{args.prefix}', 忽略后缀: '{args.suffix}')")
    label_stems_normalized, label_norm_to_orig_map = process_directory(
        args.label, 
        args.prefix, 
        args.suffix
    )
    if not label_stems_normalized:
        print(f"未在 Label 文件夹中找到任何 .png 或 .jpg 文件。")

    # 3. 执行比较 (使用集合运算)
    matching_stems = image_stems.intersection(label_stems_normalized)
    
    # Image 文件夹中多余的 (在 Image 中，但不在 Label 中)
    extra_in_image = image_stems.difference(label_stems_normalized)
    
    # Label 文件夹中多余的 (在 Label 中，但不在 Image 中)
    extra_in_label_normalized = label_stems_normalized.difference(image_stems)
    
    # 将 Label 中多余的文件名转换回原始名称
    extra_in_label_original = {label_norm_to_orig_map[stem] for stem in extra_in_label_normalized}

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

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

    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)))
    print("(显示的是原始文件名，非规范化名称)")
    if not extra_in_label_original:
        print("(无)")
    else:
        for file_stem in sorted(extra_in_label_original):
            print(file_stem)
    
    print("\n" + "="*30)

if __name__ == "__main__":
    main()