#!/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

# 尝试导入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):
        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):
        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())[:8]

        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())[:8],
                    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
