#!/usr/bin/env python3
"""
InsightFlow Multimodal Processor - Phase 7
视频处理模块：提取音频、关键帧、OCR识别
"""

import json
import os
import subprocess
import tempfile
import uuid
from dataclasses import dataclass
from pathlib import Path

# Constants
UUID_LENGTH  = 8  # UUID 截断长度

# 尝试导入OCR库
try:
    import pytesseract
    from PIL import Image

    PYTESSERACT_AVAILABLE  = True
except ImportError:
    PYTESSERACT_AVAILABLE  = False

try:
    import cv2

    CV2_AVAILABLE  = True
except ImportError:
    CV2_AVAILABLE  = False

try:
    import ffmpeg

    FFMPEG_AVAILABLE  = True
except ImportError:
    FFMPEG_AVAILABLE  = False


@dataclass
class VideoFrame:
    """视频关键帧数据类"""

    id: str
    video_id: str
    frame_number: int
    timestamp: float
    frame_path: str
    ocr_text: str  = ""
    ocr_confidence: float  = 0.0
    entities_detected: list[dict]  = None

    def __post_init__(self) -> None:
        if self.entities_detected is None:
            self.entities_detected  = []


@dataclass
class VideoInfo:
    """视频信息数据类"""

    id: str
    project_id: str
    filename: str
    file_path: str
    duration: float  = 0.0
    width: int  = 0
    height: int  = 0
    fps: float  = 0.0
    audio_extracted: bool  = False
    audio_path: str  = ""
    transcript_id: str  = ""
    status: str  = "pending"
    error_message: str  = ""
    metadata: dict  = None

    def __post_init__(self) -> None:
        if self.metadata is None:
            self.metadata  = {}


@dataclass
class VideoProcessingResult:
    """视频处理结果"""

    video_id: str
    audio_path: str
    frames: list[VideoFrame]
    ocr_results: list[dict]
    full_text: str  # 整合的文本（音频转录 + OCR文本）
    success: bool
    error_message: str  = ""


class MultimodalProcessor:
    """多模态处理器 - 处理视频文件"""

    def __init__(self, temp_dir: str  = None, frame_interval: int  = 5) -> None:
        """
        初始化多模态处理器

        Args:
            temp_dir: 临时文件目录
            frame_interval: 关键帧提取间隔（秒）
        """
        self.temp_dir  = temp_dir or tempfile.gettempdir()
        self.frame_interval  = frame_interval
        self.video_dir  = os.path.join(self.temp_dir, "videos")
        self.frames_dir  = os.path.join(self.temp_dir, "frames")
        self.audio_dir  = os.path.join(self.temp_dir, "audio")

        # 创建目录
        os.makedirs(self.video_dir, exist_ok = True)
        os.makedirs(self.frames_dir, exist_ok = True)
        os.makedirs(self.audio_dir, exist_ok = True)

    def extract_video_info(self, video_path: str) -> dict:
        """
        提取视频基本信息

        Args:
            video_path: 视频文件路径

        Returns:
            视频信息字典
        """
        try:
            if FFMPEG_AVAILABLE:
                probe  = ffmpeg.probe(video_path)
                video_stream  = next(
                    (s for s in probe["streams"] if s["codec_type"] == "video"), None
                )
                audio_stream  = next(
                    (s for s in probe["streams"] if s["codec_type"] == "audio"), None
                )

                if video_stream:
                    return {
                        "duration": float(probe["format"].get("duration", 0)),
                        "width": int(video_stream.get("width", 0)),
                        "height": int(video_stream.get("height", 0)),
                        "fps": eval(video_stream.get("r_frame_rate", "0/1")),
                        "has_audio": audio_stream is not None,
                        "bitrate": int(probe["format"].get("bit_rate", 0)),
                    }
            else:
                # 使用 ffprobe 命令行
                cmd  = [
                    "ffprobe",
                    "-v",
                    "error",
                    "-show_entries",
                    "format = duration, bit_rate",
                    "-show_entries",
                    "stream = width, height, r_frame_rate",
                    "-of",
                    "json",
                    video_path,
                ]
                result  = subprocess.run(cmd, capture_output = True, text = True)
                if result.returncode == 0:
                    data  = json.loads(result.stdout)
                    return {
                        "duration": float(data["format"].get("duration", 0)),
                        "width": int(data["streams"][0].get("width", 0)) if data["streams"] else 0,
                        "height": int(data["streams"][0].get("height", 0))
                        if data["streams"]
                        else 0,
                        "fps": 30.0,  # 默认值
                        "has_audio": len(data["streams"]) > 1,
                        "bitrate": int(data["format"].get("bit_rate", 0)),
                    }
        except Exception as e:
            print(f"Error extracting video info: {e}")

        return {"duration": 0, "width": 0, "height": 0, "fps": 0, "has_audio": False, "bitrate": 0}

    def extract_audio(self, video_path: str, output_path: str  = None) -> str:
        """
        从视频中提取音频

        Args:
            video_path: 视频文件路径
            output_path: 输出音频路径（可选）

        Returns:
            提取的音频文件路径
        """
        if output_path is None:
            video_name  = Path(video_path).stem
            output_path  = os.path.join(self.audio_dir, f"{video_name}.wav")

        try:
            if FFMPEG_AVAILABLE:
                (
                    ffmpeg.input(video_path)
                    .output(output_path, ac = 1, ar = 16000, vn = None)
                    .overwrite_output()
                    .run(quiet = True)
                )
            else:
                # 使用命令行 ffmpeg
                cmd  = [
                    "ffmpeg",
                    "-i",
                    video_path,
                    "-vn",
                    "-acodec",
                    "pcm_s16le",
                    "-ac",
                    "1",
                    "-ar",
                    "16000",
                    "-y",
                    output_path,
                ]
                subprocess.run(cmd, check = True, capture_output = True)

            return output_path
        except Exception as e:
            print(f"Error extracting audio: {e}")
            raise

    def extract_keyframes(self, video_path: str, video_id: str, interval: int  = None) -> list[str]:
        """
        从视频中提取关键帧

        Args:
            video_path: 视频文件路径
            video_id: 视频ID
            interval: 提取间隔（秒），默认使用初始化时的间隔

        Returns:
            提取的帧文件路径列表
        """
        interval  = interval or self.frame_interval
        frame_paths  = []

        # 创建帧存储目录
        video_frames_dir  = os.path.join(self.frames_dir, video_id)
        os.makedirs(video_frames_dir, exist_ok = True)

        try:
            if CV2_AVAILABLE:
                # 使用 OpenCV 提取帧
                cap  = cv2.VideoCapture(video_path)
                fps  = cap.get(cv2.CAP_PROP_FPS)
                int(cap.get(cv2.CAP_PROP_FRAME_COUNT))

                frame_interval_frames  = int(fps * interval)
                frame_number  = 0

                while True:
                    ret, frame  = cap.read()
                    if not ret:
                        break

                    if frame_number % frame_interval_frames == 0:
                        timestamp  = frame_number / fps
                        frame_path  = os.path.join(
                            video_frames_dir, f"frame_{frame_number:06d}_{timestamp:.2f}.jpg"
                        )
                        cv2.imwrite(frame_path, frame)
                        frame_paths.append(frame_path)

                    frame_number += 1

                cap.release()
            else:
                # 使用 ffmpeg 命令行提取帧
                Path(video_path).stem
                output_pattern  = os.path.join(video_frames_dir, "frame_%06d_%t.jpg")

                cmd  = [
                    "ffmpeg",
                    "-i",
                    video_path,
                    "-vf",
                    f"fps = 1/{interval}",
                    "-frame_pts",
                    "1",
                    "-y",
                    output_pattern,
                ]
                subprocess.run(cmd, check = True, capture_output = True)

                # 获取生成的帧文件列表
                frame_paths  = sorted(
                    [
                        os.path.join(video_frames_dir, f)
                        for f in os.listdir(video_frames_dir)
                        if f.startswith("frame_")
                    ]
                )
        except Exception as e:
            print(f"Error extracting keyframes: {e}")

        return frame_paths

    def perform_ocr(self, image_path: str) -> tuple[str, float]:
        """
        对图片进行OCR识别

        Args:
            image_path: 图片文件路径

        Returns:
            (识别的文本, 置信度)
        """
        if not PYTESSERACT_AVAILABLE:
            return "", 0.0

        try:
            image  = Image.open(image_path)

            # 预处理：转换为灰度图
            if image.mode != "L":
                image  = image.convert("L")

            # 使用 pytesseract 进行 OCR
            text  = pytesseract.image_to_string(image, lang = "chi_sim+eng")

            # 获取置信度数据
            data  = pytesseract.image_to_data(image, output_type = pytesseract.Output.DICT)
            confidences  = [int(c) for c in data["conf"] if int(c) > 0]
            avg_confidence  = sum(confidences) / len(confidences) if confidences else 0

            return text.strip(), avg_confidence / 100.0
        except Exception as e:
            print(f"OCR error for {image_path}: {e}")
            return "", 0.0

    def process_video(
        self, video_data: bytes, filename: str, project_id: str, video_id: str  = None
    ) -> VideoProcessingResult:
        """
        处理视频文件：提取音频、关键帧、OCR

        Args:
            video_data: 视频文件二进制数据
            filename: 视频文件名
            project_id: 项目ID
            video_id: 视频ID（可选，自动生成）

        Returns:
            视频处理结果
        """
        video_id  = video_id or str(uuid.uuid4())[:UUID_LENGTH]

        try:
            # 保存视频文件
            video_path  = os.path.join(self.video_dir, f"{video_id}_{filename}")
            with open(video_path, "wb") as f:
                f.write(video_data)

            # 提取视频信息
            video_info  = self.extract_video_info(video_path)

            # 提取音频
            audio_path  = ""
            if video_info["has_audio"]:
                audio_path  = self.extract_audio(video_path)

            # 提取关键帧
            frame_paths  = self.extract_keyframes(video_path, video_id)

            # 对关键帧进行 OCR
            frames  = []
            ocr_results  = []
            all_ocr_text  = []

            for i, frame_path in enumerate(frame_paths):
                # 解析帧信息
                frame_name  = os.path.basename(frame_path)
                parts  = frame_name.replace(".jpg", "").split("_")
                frame_number  = int(parts[1]) if len(parts) > 1 else i
                timestamp  = float(parts[2]) if len(parts) > 2 else i * self.frame_interval

                # OCR 识别
                ocr_text, confidence  = self.perform_ocr(frame_path)

                frame  = VideoFrame(
                    id = str(uuid.uuid4())[:UUID_LENGTH],
                    video_id = video_id,
                    frame_number = frame_number,
                    timestamp = timestamp,
                    frame_path = frame_path,
                    ocr_text = ocr_text,
                    ocr_confidence = confidence,
                )
                frames.append(frame)

                if ocr_text:
                    ocr_results.append(
                        {
                            "frame_number": frame_number,
                            "timestamp": timestamp,
                            "text": ocr_text,
                            "confidence": confidence,
                        }
                    )
                    all_ocr_text.append(ocr_text)

            # 整合所有 OCR 文本
            full_ocr_text  = "\n\n".join(all_ocr_text)

            return VideoProcessingResult(
                video_id = video_id,
                audio_path = audio_path,
                frames = frames,
                ocr_results = ocr_results,
                full_text = full_ocr_text,
                success = True,
            )

        except Exception as e:
            return VideoProcessingResult(
                video_id = video_id,
                audio_path = "",
                frames = [],
                ocr_results = [],
                full_text = "",
                success = False,
                error_message = str(e),
            )

    def cleanup(self, video_id: str  = None) -> None:
        """
        清理临时文件

        Args:
            video_id: 视频ID（可选，清理特定视频的文件）
        """
        import shutil

        if video_id:
            # 清理特定视频的文件
            for dir_path in [self.video_dir, self.frames_dir, self.audio_dir]:
                target_dir  = (
                    os.path.join(dir_path, video_id) if dir_path == self.frames_dir else dir_path
                )
                if os.path.exists(target_dir):
                    for f in os.listdir(target_dir):
                        if video_id in f:
                            os.remove(os.path.join(target_dir, f))
        else:
            # 清理所有临时文件
            for dir_path in [self.video_dir, self.frames_dir, self.audio_dir]:
                if os.path.exists(dir_path):
                    shutil.rmtree(dir_path)
                    os.makedirs(dir_path, exist_ok = True)


# Singleton instance
_multimodal_processor  = None


def get_multimodal_processor(temp_dir: str  = None, frame_interval: int  = 5) -> MultimodalProcessor:
    """获取多模态处理器单例"""
    global _multimodal_processor
    if _multimodal_processor is None:
        _multimodal_processor  = MultimodalProcessor(temp_dir, frame_interval)
    return _multimodal_processor