# -*- coding: utf-8 -*-
"""video.ipynb

Automatically generated by Colab.

Original file is located at
    https://colab.research.google.com/drive/1d-uwY0B5q7gOItN6fyA-1RoBcDwkOxb6
"""
# ✅ 全局加载模型一次（推荐）
from faster_whisper import WhisperModel
whisper_model = WhisperModel("tiny", device="cpu", compute_type="int8")

# 将整段video.py封装为analyze_video函数
def analyze_video(file, lang=None):
    output = ""  # 用于收集结果文本
    
    import os
    import zipfile
    from dotenv import load_dotenv
    from ultralytics import YOLO
    import cv2
    import requests
    import base64
    import json
    
    # === 配置 ===
    VIDEO_PATH = file
    YOLO_MODEL_PATH = "trained_dataset/runs/detect/train/weights/best.pt"
    # ✅ 自动解压模型文件（如果还没解压）
    if not os.path.exists("trained_dataset"):
        print("📦 正在解压 YOLO 模型文件...")
        with zipfile.ZipFile("trained_dataset.zip", "r") as zip_ref:
            zip_ref.extractall(".")
        print("✅ 解压完成")
    
    # ✅ 加载环境变量
    from dotenv import load_dotenv
    load_dotenv()
    API_KEY = os.getenv("GOOGLE_VISION_API_KEY")
    VISION_API_URL = f"https://vision.googleapis.com/v1/images:annotate?key={API_KEY}"
    LIKELIHOOD_MAPPING = {"UNKNOWN": 0, "VERY_UNLIKELY": 1, "UNLIKELY": 2, "POSSIBLE": 3, "LIKELY": 4, "VERY_LIKELY": 5}
    
    # === 初始化 ===      
    model = YOLO(YOLO_MODEL_PATH)
    cap = cv2.VideoCapture(VIDEO_PATH)
    
    safe_search_results = {"adult": 0, "spoof": 0, "medical": 0, "violence": 0, "racy": 0}
    detected_texts = set()
    detected_labels = set()
    total_frames = 0
    
    while cap.isOpened():
        ret, frame = cap.read()
        if not ret:
            break
        total_frames += 1
    
        # YOLO 检测
        results = model.predict(source=frame, conf=0.2, save=False, verbose=False)
        if len(results[0].boxes) > 0:
            for c in results[0].boxes.cls:
                class_id = int(c.item())
                class_name = results[0].names[class_id]
                detected_labels.add(class_name)
    
        # 每 10 帧 Google 内容分析
        if total_frames % 10 == 0:
            _, buffer = cv2.imencode('.jpg', frame)
            img_base64 = base64.b64encode(buffer).decode()
            payload = {
                "requests": [{
                    "image": {"content": img_base64},
                    "features": [
                        {"type": "SAFE_SEARCH_DETECTION"},
                        {"type": "TEXT_DETECTION"}
                    ]
                }]
            }
            response = requests.post(VISION_API_URL, json=payload)
            result = response.json()
    
            if "responses" in result and len(result["responses"]) > 0:
                safe = result["responses"][0].get("safeSearchAnnotation", {})
                for key in safe_search_results.keys():
                    if LIKELIHOOD_MAPPING.get(safe.get(key, "UNKNOWN"), 0) >= 4:
                        safe_search_results[key] += 1
    
                texts = result["responses"][0].get("textAnnotations", [])
                for text in texts:
                    detected_texts.add(text["description"])
    
    cap.release()
    
    # === Whisper 音频转文字 ===
    output += "\n🎙 正在转录音频文本...:\n"
    
    segments, info = whisper_model.transcribe(VIDEO_PATH, language=lang if lang != "auto" else None)
    transcribed_text = "".join([seg.text for seg in segments])
    
    # === 输出结果 ===
    output += "\n🎯 YOLO检测到的图案类别:\n"
    if detected_labels:
        output += "✅ 检测到图案：" + "、".join(detected_labels) + "\n"
    else:
        output += "❌ 未检测到任何图案\n"
    
    def risk_level(count, total):
        if count > total * 0.05:
            return "⛔️ 高风险"
        elif count > 5:
            return "⚠️ 中等风险"
        else:
            return "✅ 低风险"
    
    output += "\n🔎 内容安全风险分析（每类满足可能性 ≥ LIKELY 的帧计数）:\n"
    for k, v in safe_search_results.items():
        output += f"{k.capitalize():<10}: {risk_level(v, total_frames)}（{v} 帧）\n"
    
    output += "\n📝 视觉文字识别（OCR）:\n"
    output += " ".join(detected_texts) + "\n" if detected_texts else "无可识别文字\n"
    
    output += "\n🔊 Whisper语音识别结果:\n"
    output += transcribed_text + "\n" if transcribed_text.strip() else "无有效语音\n"
    
    return output