#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
SHOT-QC-AUTOMATION
镜头QC自动化 — 每个镜头自动拆帧，检查竖屏、字幕、换脸、牌匾、遮挡、现代物品。

功能:
1. 输入视频文件
2. 自动拆帧
3. 检查:
   - 竖屏 (9:16)
   - 字幕存在性和位置
   - 换脸 (与参考图对比)
   - 牌匾文字正确性
   - 遮挡 (人物被遮挡)
   - 现代物品 (手机、汽车等)
4. 输出 QC 报告 JSON

依赖:
  pip install opencv-python numpy  # 基础
  pip install pytesseract  # OCR (需要系统安装 tesseract)
  # YOLO 可选: pip install ultralytics

用法:
  python shot-qc-automation.py --video input.mp4 --character CHAR-003-SuBai
  python shot-qc-automation.py --batch video_list.json
  python shot-qc-automation.py --video input.mp4 --output qc_report.json
"""

import os
import sys
import json
import argparse
from pathlib import Path
from datetime import datetime
import numpy as np

try:
    import cv2
    CV2_AVAILABLE = True
except ImportError:
    CV2_AVAILABLE = False
    print("⚠️  OpenCV (cv2) 未安装，将使用简化模式")

try:
    import pytesseract
    TESSERACT_AVAILABLE = True
except ImportError:
    TESSERACT_AVAILABLE = False
    print("⚠️  pytesseract 未安装，OCR 功能不可用")

try:
    from PIL import Image
    PIL_AVAILABLE = True
except ImportError:
    PIL_AVAILABLE = False


PROJECT_ROOT = Path(__file__).parent.parent
sys.path.insert(0, str(PROJECT_ROOT / "engines"))


class ShotQCAutomation:
    """镜头QC自动化"""

    def __init__(self, character_id=None, reference_images=None):
        self.character_id = character_id
        self.reference_images = reference_images or []
        self.qc_items = [
            "vertical_screen",  # 竖屏
            "subtitle",         # 字幕
            "face_swap",        # 换脸
            "plaque_text",      # 牌匾文字
            "occlusion",        # 遮挡
            "modern_items"       # 现代物品
        ]

        # 加载参考图
        self.reference_images_data = []
        if character_id:
            self._load_reference_images()

    def _load_reference_images(self):
        """加载角色参考图"""
        if not CV2_AVAILABLE:
            return

        char_dir = PROJECT_ROOT / "assets" / "characters" / self.character_id / "approved"
        if not char_dir.exists():
            print(f"⚠️  角色目录不存在: {char_dir}")
            return

        for img_file in char_dir.glob("*.png"):
            img = cv2.imread(str(img_file))
            if img is not None:
                self.reference_images_data.append({
                    "path": str(img_file),
                    "data": img,
                    "name": img_file.name
                })
                print(f"  ✓ 已加载参考图: {img_file.name}")

        print(f"  共加载 {len(self.reference_images_data)} 张参考图")

    def qc_video(self, video_path, output_path=None):
        """
        QC 单个视频

        返回:
          {
            "video_path": str,
            "passed": bool,
            "score": float,  # 0-10
            "checks": {
              "vertical_screen": {"passed": bool, "detail": str},
              "subtitle": {...},
              ...
            },
            "frames_checked": int,
            "issues": list
          }
        """
        print(f"\n🔍 QC 视频: {Path(video_path).name}")
        print("=" * 60)

        video_path = Path(video_path)
        if not video_path.exists():
            return {"passed": False, "error": f"视频不存在: {video_path}"}

        if not CV2_AVAILABLE:
            print("⚠️  OpenCV 不可用，跳过 QC")
            return {
                "passed": None,
                "warning": "OpenCV 不可用，QC 未执行",
                "qc_skipped": True
            }

        # 打开视频
        cap = cv2.VideoCapture(str(video_path))
        if not cap.isOpened():
            return {"passed": False, "error": f"无法打开视频: {video_path}"}

        # 视频信息
        fps = cap.get(cv2.CAP_PROP_FPS)
        frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
        width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
        height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))

        print(f"  分辨率: {width}x{height}")
        print(f"  FPS: {fps}")
        print(f"  帧数: {frame_count}")

        # 检查项
        results = {
            "video_path": str(video_path),
            "resolution": f"{width}x{height}",
            "fps": fps,
            "frame_count": frame_count,
            "passed": True,
            "score": 10.0,
            "checks": {},
            "frames_checked": 0,
            "issues": []
        }

        # 1. 竖屏检查
        print(f"\n  [1/6] 竖屏检查...")
        vertical_result = self._check_vertical_screen(width, height)
        results["checks"]["vertical_screen"] = vertical_result
        if not vertical_result["passed"]:
            results["passed"] = False
            results["score"] -= 2.0
            results["issues"].append("竖屏比例错误")

        # 2. 字幕检查 (抽帧)
        print(f"  [2/6] 字幕检查...")
        subtitle_result = self._check_subtitle(cap, frame_count, fps)
        results["checks"]["subtitle"] = subtitle_result
        if not subtitle_result["passed"]:
            results["passed"] = False
            results["score"] -= 1.5
            results["issues"].append("字幕检查失败")

        # 3. 换脸检查 (与参考图对比)
        print(f"  [3/6] 换脸检查...")
        face_swap_result = self._check_face_swap(cap, frame_count, fps)
        results["checks"]["face_swap"] = face_swap_result
        if not face_swap_result["passed"]:
            results["passed"] = False
            results["score"] -= 2.0
            results["issues"].append("疑似换脸")

        # 4. 牌匾文字检查
        print(f"  [4/6] 牌匾文字检查...")
        plaque_result = self._check_plaque_text(cap, frame_count, fps)
        results["checks"]["plaque_text"] = plaque_result
        if not plaque_result["passed"]:
            results["passed"] = False
            results["score"] -= 1.5
            results["issues"].append("牌匾文字错误")

        # 5. 遮挡检查
        print(f"  [5/6] 遮挡检查...")
        occlusion_result = self._check_occlusion(cap, frame_count, fps)
        results["checks"]["occlusion"] = occlusion_result
        if not occlusion_result["passed"]:
            results["passed"] = False
            results["score"] -= 1.0
            results["issues"].append("人物被遮挡")

        # 6. 现代物品检查
        print(f"  [6/6] 现代物品检查...")
        modern_result = self._check_modern_items(cap, frame_count, fps)
        results["checks"]["modern_items"] = modern_result
        if not modern_result["passed"]:
            results["passed"] = False
            results["score"] -= 1.0
            results["issues"].append("检测到现代物品")

        # 确保分数在 0-10 之间
        results["score"] = max(0, min(10, results["score"]))

        # 重置视频到开头
        cap.set(cv2.CAP_PROP_POS_FRAMES, 0)
        cap.release()

        # 打印总结
        print(f"\n📊 QC 总结")
        print(f"  通过: {'✅' if results['passed'] else '❌'}")
        print(f"  分数: {results['score']:.1f}/10")
        print(f"  问题数: {len(results['issues'])}")
        for issue in results["issues"]:
            print(f"    - {issue}")

        # 保存报告
        if output_path is None:
            output_path = PROJECT_ROOT / "outputs" / "qc_reports" / f"{video_path.stem}_qc.json"
        else:
            output_path = Path(output_path)

        output_path.parent.mkdir(parents=True, exist_ok=True)
        with open(output_path, "w", encoding="utf-8") as f:
            json.dump(results, f, ensure_ascii=False, indent=2)

        print(f"\n  报告已保存: {output_path}")

        return results

    def _check_vertical_screen(self, width, height):
        """检查竖屏 (9:16)"""
        # 竖屏: 宽度 < 高度，比例接近 9:16
        if width >= height:
            return {
                "passed": False,
                "detail": f"横屏 {width}x{height}，应为竖屏 9:16",
                "aspect_ratio": width / height
            }

        # 检查比例是否接近 9:16
        ratio = width / height
        target_ratio = 9 / 16  # ≈ 0.5625

        if abs(ratio - target_ratio) < 0.05:
            return {
                "passed": True,
                "detail": f"竖屏比例正确 {width}x{height} (ratio={ratio:.3f})",
                "aspect_ratio": ratio
            }
        else:
            return {
                "passed": False,
                "detail": f"竖屏比例不正确 {width}x{height} (ratio={ratio:.3f}, target={target_ratio:.3f})",
                "aspect_ratio": ratio
            }

    def _check_subtitle(self, cap, frame_count, fps):
        """检查字幕 (抽帧 + OCR)"""
        if not TESSERACT_AVAILABLE:
            return {
                "passed": True,  # 无法检查，默认通过
                "detail": "Tesseract OCR 不可用，跳过字幕检查",
                "skipped": True
            }

        # 抽帧: 每秒抽1帧
        sample_interval = int(fps)
        if sample_interval < 1:
            sample_interval = 1

        frames_to_check = []
        for i in range(0, frame_count, sample_interval):
            frames_to_check.append(i)

        # 限制最多检查 30 帧
        if len(frames_to_check) > 30:
            step = len(frames_to_check) // 30
            frames_to_check = frames_to_check[::step][:30]

        subtitle_found = False
        subtitle_positions = []

        for frame_idx in frames_to_check:
            cap.set(cv2.CAP_PROP_POS_FRAMES, frame_idx)
            ret, frame = cap.read()
            if not ret:
                continue

            # OCR 检测字幕 (通常在画面底部 1/4 区域)
            height, width = frame.shape[:2]
            subtitle_region = frame[int(height * 0.75):, :]  # 底部 1/4

            try:
                text = pytesseract.image_to_string(subtitle_region, config='--psm 6')
                if text.strip():
                    subtitle_found = True
                    subtitle_positions.append(frame_idx / fps)  # 秒数
            except Exception as e:
                pass

        if subtitle_found:
            return {
                "passed": True,
                "detail": f"检测到字幕，出现位置: {len(subtitle_positions)} 处",
                "subtitle_positions": subtitle_positions[:10]  # 前10个位置
            }
        else:
            return {
                "passed": False,
                "detail": "未检测到字幕",
                "subtitle_positions": []
            }

    def _check_face_swap(self, cap, frame_count, fps):
        """检查换脸 (与参考图对比)"""
        if len(self.reference_images_data) == 0:
            return {
                "passed": True,  # 无参考图，无法检查
                "detail": "无参考图，跳过换脸检查",
                "skipped": True
            }

        # 抽帧: 每分钟抽1帧
        sample_interval = int(fps * 60)
        if sample_interval < 1:
            sample_interval = 1

        frames_to_check = []
        for i in range(0, frame_count, sample_interval):
            frames_to_check.append(i)

        # 限制最多检查 10 帧
        if len(frames_to_check) > 10:
            step = len(frames_to_check) // 10
            frames_to_check = frames_to_check[::step][:10]

        face_swap_detected = False
        suspicious_frames = []

        for frame_idx in frames_to_check:
            cap.set(cv2.CAP_PROP_POS_FRAMES, frame_idx)
            ret, frame = cap.read()
            if not ret:
                continue

            # 简化方法: 比较直方图
            frame_hist = cv2.calcHist([frame], [0, 1, 2], None, [8, 8, 8], [0, 256, 0, 256, 0, 256])
            cv2.normalize(frame_hist, frame_hist)

            for ref in self.reference_images_data:
                ref_hist = cv2.calcHist([ref["data"]], [0, 1, 2], None, [8, 8, 8], [0, 256, 0, 256, 0, 256])
                cv2.normalize(ref_hist, ref_hist)

                # 比较直方图相关性
                similarity = cv2.compareHist(frame_hist, ref_hist, cv2.HISTCMP_CORREL)

                if similarity < 0.3:  # 低相似度 = 可能换脸
                    face_swap_detected = True
                    suspicious_frames.append({
                        "frame": frame_idx,
                        "time": frame_idx / fps,
                        "similarity": float(similarity)
                    })

        if not face_swap_detected:
            return {
                "passed": True,
                "detail": f"未检测到换脸 (检查了 {len(frames_to_check)} 帧)",
                "frames_checked": len(frames_to_check)
            }
        else:
            return {
                "passed": False,
                "detail": f"疑似换脸 (检测到 {len(suspicious_frames)} 处异常)",
                "suspicious_frames": suspicious_frames[:5]
            }

    def _check_plaque_text(self, cap, frame_count, fps):
        """检查牌匾文字 (OCR)"""
        if not TESSERACT_AVAILABLE:
            return {
                "passed": True,
                "detail": "Tesseract OCR 不可用，跳过牌匾文字检查",
                "skipped": True
            }

        # 抽帧: 牌匾通常静止，抽 5 帧即可
        frames_to_check = [0, int(frame_count * 0.25), int(frame_count * 0.5), int(frame_count * 0.75), frame_count - 1]
        frames_to_check = [f for f in frames_to_check if f < frame_count]

        plaque_text_detected = []
        correct_text = "天道宗"  # 期望的牌匾文字

        for frame_idx in frames_to_check:
            cap.set(cv2.CAP_PROP_POS_FRAMES, frame_idx)
            ret, frame = cap.read()
            if not ret:
                continue

            # OCR 整帧
            try:
                text = pytesseract.image_to_string(frame, config='--psm 6')
                if correct_text in text:
                    plaque_text_detected.append({
                        "frame": frame_idx,
                        "time": frame_idx / fps,
                        "text": text.strip()[:50]
                    })
            except Exception as e:
                pass

        if len(plaque_text_detected) > 0:
            return {
                "passed": True,
                "detail": f"牌匾文字正确 '{correct_text}' (在 {len(plaque_text_detected)} 帧中检测到)",
                "detected": plaque_text_detected
            }
        else:
            return {
                "passed": False,
                "detail": f"未检测到牌匾文字 '{correct_text}'",
                "warning": "可能牌匾文字错误或未出现在画面中"
            }

    def _check_occlusion(self, cap, frame_count, fps):
        """检查遮挡 (人物被遮挡)"""
        # 简化方法: 检测画面中是否突然出现大块纯色区域 (可能是水印或遮挡)

        sample_interval = int(fps * 10)  # 每10秒抽1帧
        if sample_interval < 1:
            sample_interval = 1

        frames_to_check = []
        for i in range(0, frame_count, sample_interval):
            frames_to_check.append(i)

        occlusion_detected = False

        for frame_idx in frames_to_check:
            cap.set(cv2.CAP_PROP_POS_FRAMES, frame_idx)
            ret, frame = cap.read()
            if not ret:
                continue

            # 转灰度
            gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)

            # 计算灰度直方图
            hist = cv2.calcHist([gray], [0], None, [256], [0, 256])
            hist = hist.flatten()

            # 如果某个灰度值占比过高 = 可能有遮挡/水印
            max_ratio = np.max(hist) / (frame.shape[0] * frame.shape[1])
            if max_ratio > 0.3:  # 30% 以上像素是同一颜色
                occlusion_detected = True
                break

        if not occlusion_detected:
            return {
                "passed": True,
                "detail": f"未检测到明显遮挡 (检查了 {len(frames_to_check)} 帧)"
            }
        else:
            return {
                "passed": False,
                "detail": "检测到可能的遮挡 (画面中有大块纯色区域)"
            }

    def _check_modern_items(self, cap, frame_count, fps):
        """检查现代物品 (手机、汽车等)"""
        # 简化方法: 检测画面中是否有现代物品的特征颜色/形状

        # TODO: 使用 YOLO 检测现代物品
        # 暂时跳过，返回通过

        return {
            "passed": True,
            "detail": "现代物品检查 (TODO: 需要 YOLO 模型)",
            "skipped": True,
            "todo": "Implement YOLO-based modern item detection"
        }

    def batch_qc(self, video_list_config):
        """
        批量 QC

        video_list_config 格式:
          {
            "videos": [
              {"path": "ep01-shot01.mp4", "character": "CHAR-003-SuBai"},
              ...
            ]
          }
        """
        if isinstance(video_list_config, str):
            config_file = Path(video_list_config)
            with open(config_file, "r", encoding="utf-8") as f:
                config = json.load(f)
            videos = config.get("videos", [])
        elif isinstance(video_list_config, list):
            videos = video_list_config
        else:
            videos = []

        print(f"\n📦 批量 QC: {len(videos)} 个视频")

        results = []
        for i, video_config in enumerate(videos):
            print(f"\n  进度: [{i+1}/{len(videos)}]")

            video_path = video_config.get("path")
            character_id = video_config.get("character", self.character_id)

            qc = ShotQCAutomation(character_id=character_id)
            result = qc.qc_video(video_path)
            results.append(result)

        # 统计
        passed_count = sum(1 for r in results if r.get("passed"))
        print(f"\n✅ 批量 QC 完成: {passed_count}/{len(results)} 通过")

        # 保存批量报告
        report_path = PROJECT_ROOT / "outputs" / "qc_reports" / f"batch_qc_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json"
        report_path.parent.mkdir(parents=True, exist_ok=True)

        with open(report_path, "w", encoding="utf-8") as f:
            json.dump({
                "total": len(results),
                "passed": passed_count,
                "results": results,
                "generated_at": datetime.now().isoformat()
            }, f, ensure_ascii=False, indent=2)

        print(f"  报告已保存: {report_path}")

        return results


def main():
    parser = argparse.ArgumentParser(description="SHOT-QC-AUTOMATION")
    parser.add_argument("--video", type=str, help="输入视频路径")
    parser.add_argument("--character", type=str, help="角色ID (用于换脸检查)")
    parser.add_argument("--output", type=str, help="输出 QC 报告路径")
    parser.add_argument("--batch", type=str, help="批量 QC 配置文件 (JSON)")

    args = parser.parse_args()

    if args.batch:
        # 批量模式
        qc = ShotQCAutomation(character_id=args.character)
        results = qc.batch_qc(args.batch)
        sys.exit(0 if all(r.get("passed") for r in results) else 1)

    if not args.video:
        parser.print_help()
        sys.exit(1)

    # 单文件模式
    qc = ShotQCAutomation(character_id=args.character)
    result = qc.qc_video(args.video, output_path=args.output)

    if result.get("passed"):
        print(f"\n✅ QC 通过")
        sys.exit(0)
    elif result.get("passed") is None and result.get("qc_skipped"):
        print(f"\n⚠️  QC 跳过 (依赖不可用)")
        sys.exit(0)
    else:
        if result.get("error"):
            print(f"\n❌ QC 失败: {result['error']}")
        else:
            issues = result.get("issues") or []
            issue_text = "；".join(issues) if issues else "未通过阈值"
            print(f"\n❌ QC 未通过: {issue_text}")
        sys.exit(1)


if __name__ == "__main__":
    main()