import os
import shutil
import glob
from nudenet import NudeDetector
from PIL import Image, UnidentifiedImageError

def main():
    current_dir = os.getcwd()
    nsfw_dir = os.path.join(current_dir, "NSFW")
    
    # 1. 创建 NSFW 文件夹
    if not os.path.exists(nsfw_dir):
        os.makedirs(nsfw_dir)
        print(f"📁 已创建隔离文件夹: {nsfw_dir}")

    # 2. 初始化检测器 (首次运行会自动下载模型)
    print("⏳ 正在加载 NSFW 检测模型，首次运行可能需要下载...")
    detector = NudeDetector()
    
    # 定义属于 NSFW 的标签 (根据 nudenet 的标签库)
    NSFW_LABELS = {
        "FEMALE_BREAST_EXPOSED",
        "FEMALE_GENITALIA_EXPOSED",
        "MALE_GENITALIA_EXPOSED",
        "ANUS_EXPOSED",
        "BUTTOCKS_EXPOSED",
        # 加入以下擦边标签：
        "MALE_BREAST_EXPOSED",      # 男士露胸（有时候会被误判）
    }

    # 3. 查找当前目录下所有的 jpg
    jpg_files = glob.glob(os.path.join(current_dir, "*.jpg")) + glob.glob(os.path.join(current_dir, "*.jpeg"))
    
    if not jpg_files:
        print("❌ 当前目录下没有找到 .jpg 或 .jpeg 文件。")
        return

    print(f"🔍 共找到 {len(jpg_files)} 张图片，开始扫描...\n")

    nsfw_count = 0
    
    for img_path in jpg_files:
        filename = os.path.basename(img_path)
        base_name = os.path.splitext(filename)[0]
        txt_path = os.path.join(current_dir, f"{base_name}.txt")

        try:
            # 检测图片
            detections = detector.detect(img_path)
            
            is_nsfw = False
            found_labels = set()
            
            for det in detections:
                # 如果检测到的标签在 NSFW 列表中，且置信度 > 0.6
                if det["class"] in NSFW_LABELS and det["score"] > 0.6:
                    is_nsfw = True
                    found_labels.add(det["class"])

            if is_nsfw:
                print(f"🚫 发现 NSFW: {filename} -> 特征: {', '.join(found_labels)}")
                
                # 移动图片
                shutil.move(img_path, os.path.join(nsfw_dir, filename))
                
                # 移动同名 TXT (如果有的话)
                if os.path.exists(txt_path):
                    shutil.move(txt_path, os.path.join(nsfw_dir, f"{base_name}.txt"))
                else:
                    print(f"   ⚠️ 警告: {filename} 没有对应的 txt 文件")
                
                nsfw_count += 1
                
        except UnidentifiedImageError:
            print(f"⚠️ 跳过损坏的图片: {filename}")
        except Exception as e:
            print(f"⚠️ 检测 {filename} 时出错: {e}")

    print("\n" + "="*40)
    print(f"🎉 扫描完成！共处理 {len(jpg_files)} 张图，成功隔离 {nsfw_count} 张 NSFW 图片及其文本。")

if __name__ == "__main__":
    main()