#!/usr/bin/env python3
"""
Subtitle Safe Area QC · 字幕安全区域质检器
==========================​===================
自动检查字幕有没有遮脸、遮身体表演、贴边、看不清。

依赖：
  pip install opencv-python pillow numpy
  # 可选：pip install face-recognition dlib  # 人脸识别
  # 可选：pip install ultralytics  # YOLO 人体检测

用法：
  # 检查单张帧
  python subtitle-safe-area-qc.py --frame frame.png --subtitle-box x,y,w,h --output qc-report.json

  # 批量检查视频帧
  python subtitle-safe-area-qc.py --video input.mp4 --subtitle-ass subtitles.ass --output qc-report.json

  # 交互式标注模式（手动框选人脸/身体区域）
  python subtitle-safe-area-qc.py --frame frame.png --interactive

  # 生成质检报告（包含问题帧截图）
  python subtitle-safe-area-qc.py --video input.mp4 --subtitle-ass subtitles.ass --generate-report --output-dir ./qc-report/

输出 JSON 格式：
  {
    "frame": "frame_00123.png",
    "timestamp": 12.34,
    "issues": [
      {"type": "face_occlusion", "severity": "high", "bbox": [x,y,w,h]},
      {"type": "body_occlusion", "severity": "medium", "bbox": [x,y,w,h]},
      {"type": "edge_too_close", "severity": "low", "distance": 5},
      {"type": "low_contrast", "severity": "medium", "contrast_ratio": 2.1}
    ],
    "safe": false
  }

路径：
  video-ai-system/engines/subtitle-pipeline/qc-tools/subtitle-safe-area-qc.py
"""

import argparse
import json
import os
import sys
from pathlib import Path

try:
    import cv2
    import numpy as np
    from PIL import Image
except ImportError as e:
    print(f"[ERROR] 缺少依赖：{e}")
    print("请先安装：pip install opencv-python pillow numpy")
    sys.exit(1)


# 尝试导入可选依赖
try:
    import face_recognition
    FACE_RECOGNITION_AVAILABLE = True
except ImportError:
    FACE_RECOGNITION_AVAILABLE = False
    print("[WARN] 未安装 face-recognition，将使用简化人脸检测")

try:
    from ultralytics import YOLO
    YOLO_AVAILABLE = True
except ImportError:
    YOLO_AVAILABLE = False
    print("[WARN] 未安装 ultralytics，将使用简化身体检测")


def detect_faces(image: np.ndarray) -> list:
    """
    检测人脸区域
    
    :param image: 图片数组
    :return: 人脸边界框列表 [[x,y,w,h], ...]
    """
    if FACE_RECOGNITION_AVAILABLE:
        # 使用 face_recognition 库（基于 dlib）
        rgb_image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        face_locations = face_recognition.face_locations(rgb_image)
        
        # 转换为 [x,y,w,h] 格式
        boxes = []
        for top, right, bottom, left in face_locations:
            x = left
            y = top
            w = right - left
            h = bottom - top
            boxes.append([x, y, w, h])
        
        return boxes
    else:
        # 简化方案：使用 OpenCV Haar Cascade
        cascade_path = cv2.data.haarcascades + "haarcascade_frontalface_default.xml"
        face_cascade = cv2.CascadeClassifier(cascade_path)
        
        gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
        faces = face_cascade.detectMultiScale(gray, 1.1, 4)
        
        return faces.tolist() if len(faces) > 0 else []


def detect_body(image: np.ndarray) -> list:
    """
    检测身体区域
    
    :param image: 图片数组
    :return: 身体边界框列表 [[x,y,w,h], ...]
    """
    if YOLO_AVAILABLE:
        # 使用 YOLO 检测人体
        model = YOLO("yolov8n.pt")  # 自动下载
        results = model(image)
        
        boxes = []
        for result in results:
            for box in result.boxes:
                cls = int(box.cls[0])
                # COCO 数据集中，person 的类别 ID 是 0
                if cls == 0:
                    x1, y1, x2, y2 = box.xyxy[0].tolist()
                    boxes.append([int(x1), int(y1), int(x2-x1), int(y2-y1)])
        
        return boxes
    else:
        # 简化方案：使用背景减除或轮廓检测
        print("[WARN] 未安装 YOLO，使用简化身体检测（可能不准确）")
        
        # 策略：检测图像中的大轮廓（假设身体是大区域）
        gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
        _, binary = cv2.threshold(gray, 128, 255, cv2.THRESH_BINARY_INV)
        
        contours, _ = cv2.findContours(binary, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        
        boxes = []
        for cnt in contours:
            area = cv2.contourArea(cnt)
            if area > image.shape[0] * image.shape[1] * 0.1:  # 面积 > 10%
                x, y, w, h = cv2.boundingRect(cnt)
                boxes.append([x, y, w, h])
        
        return boxes


def detect_subtitle_box(image: np.ndarray) -> list:
    """
    检测字幕框位置
    
    :param image: 图片数组
    :return: 字幕边界框 [x,y,w,h]
    """
    # 策略：检测底部区域的白色文本
    height, width = image.shape[:2]
    
    # 裁剪底部区域
    bottom_region = image[int(height * 0.85):, :]
    
    # 转为灰度图
    gray = cv2.cvtColor(bottom_region, cv2.COLOR_BGR2GRAY)
    
    # 二值化（检测白色文本）
    _, binary = cv2.threshold(gray, 200, 255, cv2.THRESH_BINARY)
    
    # 查找轮廓
    contours, _ = cv2.findContours(binary, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    
    if not contours:
        return [0, int(height * 0.85), width, int(height * 0.12)]  # 默认底部区域
    
    # 收集所有轮廓的边界框
    all_x = []
    all_y = []
    all_w = []
    all_h = []
    
    for cnt in contours:
        x, y, w, h = cv2.boundingRect(cnt)
        all_x.append(x)
        all_y.append(y)
        all_w.append(w)
        all_h.append(h)
    
    # 合并所有轮廓的边界框
    # 先在 bottom_region 坐标系里算完，最后再加偏移
    x_bottom = min(all_x)
    y_bottom = min(all_y)
    max_x_bottom = max([all_x[i] + all_w[i] for i in range(len(all_x))])
    max_y_bottom = max([all_y[i] + all_h[i] for i in range(len(all_y))])
    w = max_x_bottom - x_bottom
    h = max_y_bottom - y_bottom
    
    # 最后统一加偏移
    x = x_bottom
    y = y_bottom + int(height * 0.85)
    # w, h 不需要加偏移（它们是在bottom_region里的尺寸）
    
    return [x, y, w, h]
    
    
def calculate_iou(box1: list, box2: list) -> float:
    """
    计算两个边界框的 IOU（交并比）
    
    :param box1: 边界框1 [x,y,w,h]
    :param box2: 边界框2 [x,y,w,h]
    :return: IOU 值（0-1）
    """
    x1 = max(box1[0], box2[0])
    y1 = max(box1[1], box2[1])
    x2 = min(box1[0] + box1[2], box2[0] + box2[2])
    y2 = min(box1[1] + box1[3], box2[1] + box2[3])
    
    if x2 < x1 or y2 < y1:
        return 0.0
    
    intersection = (x2 - x1) * (y2 - y1)
    area1 = box1[2] * box1[3]
    area2 = box2[2] * box2[3]
    union = area1 + area2 - intersection
    
    return intersection / union if union > 0 else 0.0


def check_subtitle_safe_area(image_path: str, subtitle_box: list = None, interactive: bool = False) -> dict:
    """
    检查字幕安全区域
    
    :param image_path: 图片路径
    :param subtitle_box: 字幕框 [x,y,w,h]（如果为 None，自动检测）
    :param interactive: 是否交互式标注
    :return: 质检报告字典
    """
    # 读取图片
    img = cv2.imread(image_path)
    if img is None:
        print(f"[ERROR] 无法读取图片：{image_path}")
        return {}
    
    height, width = img.shape[:2]
    print(f"[INFO] 检查图片：{image_path} ({width}x{height})")
    
    # 检测字幕框
    if subtitle_box is None:
        if interactive:
            print("[INFO] 交互式标注模式：请手动框选字幕区域")
            roi = cv2.selectROI("Select Subtitle Region", img, showCrosshair=True)
            cv2.destroyAllWindows()
            subtitle_box = [int(roi[0]), int(roi[1]), int(roi[2]), int(roi[3])]
        else:
            subtitle_box = detect_subtitle_box(img)
    
    print(f"[INFO] 字幕框：x={subtitle_box[0]}, y={subtitle_box[1]}, w={subtitle_box[2]}, h={subtitle_box[3]}")
    
    # 检测人脸
    print("[INFO] 检测人脸...")
    face_boxes = detect_faces(img)
    print(f"[INFO] 找到 {len(face_boxes)} 个人脸")
    
    # 检测身体
    print("[INFO] 检测身体...")
    body_boxes = detect_body(img)
    print(f"[INFO] 找到 {len(body_boxes)} 个身体区域")
    
    # 检查问题
    issues = []
    
    # 1. 检查是否遮挡脸部
    for i, face_box in enumerate(face_boxes):
        iou = calculate_iou(subtitle_box, face_box)
        if iou > 0.1:  # IOU > 10% 认为有遮挡
            severity = "high" if iou > 0.5 else "medium"
            issues.append({
                "type": "face_occlusion",
                "severity": severity,
                "bbox": face_box,
                "iou": round(iou, 2),
                "message": f"字幕遮挡脸部（IOU={iou:.2f}）"
            })
            print(f"[ISSUE] 字幕遮挡脸部（IOU={iou:.2f}）")
    
    # 2. 检查是否遮挡身体
    for i, body_box in enumerate(body_boxes):
        iou = calculate_iou(subtitle_box, body_box)
        if iou > 0.2:  # IOU > 20% 认为有遮挡
            severity = "high" if iou > 0.6 else "medium"
            issues.append({
                "type": "body_occlusion",
                "severity": severity,
                "bbox": body_box,
                "iou": round(iou, 2),
                "message": f"字幕遮挡身体（IOU={iou:.2f}）"
            })
            print(f"[ISSUE] 字幕遮挡身体（IOU={iou:.2f}）")
    
    # 3. 检查是否贴边
    edge_threshold = 20  # 像素
    if subtitle_box[0] < edge_threshold:
        issues.append({
            "type": "edge_too_close",
            "severity": "low",
            "distance": subtitle_box[0],
            "message": f"字幕距离左边缘太近（{subtitle_box[0]}px）"
        })
        print(f"[ISSUE] 字幕距离左边缘太近（{subtitle_box[0]}px）")
    
    if (subtitle_box[0] + subtitle_box[2]) > (width - edge_threshold):
        distance = (subtitle_box[0] + subtitle_box[2]) - width
        issues.append({
            "type": "edge_too_close",
            "severity": "low",
            "distance": abs(distance),
            "message": f"字幕距离右边缘太近（{abs(distance)}px）"
        })
        print(f"[ISSUE] 字幕距离右边缘太近（{abs(distance)}px）")
    
    # 4. 检查对比度（字幕是否看不清）
    # 简化方案：计算字幕区域和背景的平均亮度差
    subtitle_region = img[
        subtitle_box[1]:subtitle_box[1]+subtitle_box[3],
        subtitle_box[0]:subtitle_box[0]+subtitle_box[2]
    ]
    
    if subtitle_region.size > 0:
        # 计算字幕区域的平均亮度
        subtitle_gray = cv2.cvtColor(subtitle_region, cv2.COLOR_BGR2GRAY)
        subtitle_brightness = np.mean(subtitle_gray)
        
        # 计算背景区域的平均亮度（字幕框上方的区域）
        background_region = img[
            max(0, subtitle_box[1]-subtitle_box[3]):subtitle_box[1],
            subtitle_box[0]:subtitle_box[0]+subtitle_box[2]
        ]
        
        if background_region.size > 0:
            background_gray = cv2.cvtColor(background_region, cv2.COLOR_BGR2GRAY)
            background_brightness = np.mean(background_gray)
            
            # 计算对比度（亮度差）
            contrast = abs(subtitle_brightness - background_brightness)
            contrast_ratio = contrast / max(subtitle_brightness, background_brightness)
            
            if contrast_ratio < 0.3:  # 对比度 < 30% 认为看不清
                severity = "high" if contrast_ratio < 0.1 else "medium"
                issues.append({
                    "type": "low_contrast",
                    "severity": severity,
                    "contrast_ratio": round(contrast_ratio, 2),
                    "message": f"字幕对比度太低（{contrast_ratio:.2f}），可能看不清"
                })
                print(f"[ISSUE] 字幕对比度太低（{contrast_ratio:.2f}），可能看不清")
    
    # 生成报告
    report = {
        "image_path": image_path,
        "video_width": width,
        "video_height": height,
        "subtitle_box": subtitle_box,
        "face_boxes": face_boxes,
        "body_boxes": body_boxes,
        "issues": issues,
        "safe": len(issues) == 0,
        "issue_count": len(issues)
    }
    
    if report["safe"]:
        print(f"[OK] 字幕安全区域检查通过")
    else:
        print(f"[WARN] 发现 {len(issues)} 个问题")
    
    return report


def batch_check_video(video_path: str, ass_path: str, output_dir: str, sample_interval: float = 1.0):
    """
    批量检查视频帧
    
    :param video_path: 视频文件路径
    :param ass_path: ASS 字幕文件路径
    :param output_dir: 输出目录
    :param sample_interval: 采样间隔（秒）
    """
    # 打开视频
    cap = cv2.VideoCapture(video_path)
    if not cap.isOpened():
        print(f"[ERROR] 无法打开视频：{video_path}")
        return
    
    fps = cap.get(cv2.CAP_PROP_FPS)
    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    
    print(f"[INFO] 视频信息：{total_frames} 帧，{fps} FPS")
    
    # 创建输出目录
    os.makedirs(output_dir, exist_ok=True)
    frames_dir = os.path.join(output_dir, "problem-frames/")
    os.makedirs(frames_dir, exist_ok=True)
    
    # 读取 ASS 字幕文件（获取字幕时间点）
    # 简化方案：假设字幕在底部，直接检测
    # TODO: 解析 ASS 文件，获取准确的字幕时间点
    
    reports = []
    frame_count = 0
    
    while True:
        ret, frame = cap.read()
        if not ret:
            break
        
        timestamp = frame_count / fps
        
        # 按采样间隔检查
        if timestamp % sample_interval < (1.0 / fps):
            # 保存帧为临时文件
            temp_frame_path = os.path.join(output_dir, f"temp_frame_{frame_count:06d}.png")
            cv2.imwrite(temp_frame_path, frame)
            
            # 检查字幕安全区域
            report = check_subtitle_safe_area(temp_frame_path)
            
            if report and not report["safe"]:
                # 保存问题帧
                problem_frame_path = os.path.join(frames_dir, f"problem_{frame_count:06d}.png")
                cv2.imwrite(problem_frame_path, frame)
                
                report["timestamp"] = timestamp
                report["problem_frame"] = problem_frame_path
                reports.append(report)
                
                print(f"[WARN] 发现问题的帧：{timestamp:.2f}s")
            
            # 删除临时文件
            os.remove(temp_frame_path)
        
        frame_count += 1
    
    cap.release()
    
    # 保存报告
    report_path = os.path.join(output_dir, "qc-report.json")
    with open(report_path, "w", encoding="utf-8") as f:
        json.dump(reports, f, ensure_ascii=False, indent=2)
    
    print(f"\n[OK] 批量检查完成")
    print(f"[INFO] 共检查 {frame_count} 帧，发现 {len(reports)} 个有问题帧")
    print(f"[INFO] 报告已保存：{report_path}")


def main():
    parser = argparse.ArgumentParser(
        description="Subtitle Safe Area QC · 字幕安全区域质检器",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
示例：
  python subtitle-safe-area-qc.py --frame frame.png --subtitle-box 100,1750,880,120 --output qc-report.json
  python subtitle-safe-area-qc.py --video input.mp4 --subtitle-ass subtitles.ass --output qc-report.json
  python subtitle-safe-area-qc.py --frame frame.png --interactive
"""
    )
    parser.add_argument("--frame", help="单张帧图片路径")
    parser.add_argument("--video", help="视频文件路径（批量检查）")
    parser.add_argument("--subtitle-box", help="字幕框坐标 x,y,w,h（逗号分隔）")
    parser.add_argument("--subtitle-ass", help="ASS 字幕文件路径（用于批量检查）")
    parser.add_argument("--output", help="输出报告文件路径")
    parser.add_argument("--interactive", action="store_true", help="交互式标注模式")
    parser.add_argument("--generate-report", action="store_true", help="生成质检报告（包含问题帧截图）")
    parser.add_argument("--output-dir", help="输出目录（用于批量检查）")
    parser.add_argument("--sample-interval", type=float, default=1.0, help="采样间隔（秒，用于批量检查）")
    
    args = parser.parse_args()
    
    if args.frame:
        # 单张帧检查
        subtitle_box = None
        if args.subtitle_box:
            subtitle_box = [int(x) for x in args.subtitle_box.split(",")]
        
        report = check_subtitle_safe_area(args.frame, subtitle_box, args.interactive)
        
        if not report:
            sys.exit(1)
        
        if args.output:
            with open(args.output, "w", encoding="utf-8") as f:
                json.dump(report, f, ensure_ascii=False, indent=2)
            print(f"[OK] 质检报告已保存：{args.output}")
    
    elif args.video:
        # 批量检查
        if not args.subtitle_ass:
            print("[ERROR] 批量检查需要指定 --subtitle-ass")
            sys.exit(1)
        
        output_dir = args.output_dir or "./qc-report/"
        batch_check_video(args.video, args.subtitle_ass, output_dir, args.sample_interval)
    
    else:
        print("[ERROR] 请指定 --frame 或 --video")
        sys.exit(1)


if __name__ == "__main__":
    main()