# -*- coding: utf-8 -*-
"""
视频双路拆分 - P1 核心模块
=================================================
功能:
  A 路:视频帧 → pHash 变化检测 → OCR → B 稿 txt
  B 路:视频 → 16kHz/单声道/16bit WAV

不引入 ffmpeg-python 等 wrapper,只用 subprocess 调系统 ffmpeg。
"""

import hashlib
import json
import os
import shutil
import subprocess
import tempfile
from pathlib import Path
from typing import Dict, List, Tuple, Optional

from PIL import Image
import imagehash


# ========================================================================
# 凭证(从环境变量读取,供 OCR 调用 DeepSeek Vision)
# ========================================================================

DEEPSEEK_API_KEY = os.environ.get("DEEPSEEK_API_KEY", "").strip()


# ========================================================================
# FFmpeg 封装
# ========================================================================


def check_ffmpeg():
    """检查 ffmpeg 是否在 PATH 中"""
    result = shutil.which("ffmpeg")
    if result is None:
        raise RuntimeError(
            "ffmpeg 未找到,请先安装 ffmpeg 并加入 PATH。"
            "下载地址: https://ffmpeg.org/download.html"
        )
    return result


def extract_frames(
    video_path: Path,
    output_dir: Path,
    fps: int = 1,
) -> List[Tuple[int, int, Path]]:
    """
    按固定 fps 抽帧
    返回: [(frame_index, timestamp_ms, image_path), ...]
    """
    check_ffmpeg()

    frames_dir = output_dir / "frames"
    frames_dir.mkdir(parents=True, exist_ok=True)

    # ffmpeg 抽帧,格式 frame_%04d.png
    frame_pattern = str(frames_dir / "frame_%04d.png")
    cmd = [
        "ffmpeg",
        "-i", str(video_path),
        "-vf", f"fps={fps}",
        "-q:v", "2",  # JPEG 质量
        frame_pattern,
        "-y",  # 覆盖
    ]

    result = subprocess.run(
        cmd,
        capture_output=True,
        text=True,
    )
    if result.returncode != 0:
        raise RuntimeError(f"ffmpeg 抽帧失败: {result.stderr}")

    # 收集抽出的帧
    frames = []
    for i, f in enumerate(sorted(frames_dir.glob("frame_*.png"))):
        # 时间戳:第 i 帧就是 i 秒
        timestamp_ms = i * 1000
        frames.append((i, timestamp_ms, f))

    return frames


def extract_audio(
    video_path: Path,
    output_path: Path,
) -> Path:
    """
    用 ffmpeg 提取音频,转为 16kHz/单声道/16bit WAV
    """
    check_ffmpeg()

    cmd = [
        "ffmpeg",
        "-i", str(video_path),
        "-ac", "1",           # 单声道
        "-ar", "16000",       # 16kHz
        "-sample_fmt", "s16", # 16bit
        str(output_path),
        "-y",
    ]

    result = subprocess.run(
        cmd,
        capture_output=True,
        text=True,
    )
    if result.returncode != 0:
        raise RuntimeError(f"ffmpeg 音频提取失败: {result.stderr}")

    return output_path


# ========================================================================
# pHash 变化检测
# ========================================================================


def compute_phash(image_path: Path) -> str:
    """计算图片的 pHash,返回 hex 字符串"""
    img = Image.open(image_path)
    ph = imagehash.phash(img)
    return str(ph)


def find_keyframes(
    frames: List[Tuple[int, int, Path]],
    threshold: int = 8,
) -> List[Dict]:
    """
    基于 pHash 海明距离找出字幕变化的关键帧

    算法:
    - 第一帧总是关键帧
    - 后续帧:如果与上一个关键帧的 pHash 海明距离 > threshold,则是新关键帧

    threshold: 海明距离阈值,默认 8
    """
    if not frames:
        return []

    keyframes = []
    last_keyframe_phash = None

    for frame_index, timestamp_ms, image_path in frames:
        phash = compute_phash(image_path)

        is_keyframe = False
        if last_keyframe_phash is None:
            # 第一帧总是关键帧
            is_keyframe = True
        else:
            # 计算海明距离
            hamming = hamming_distance(last_keyframe_phash, phash)
            if hamming > threshold:
                is_keyframe = True

        if is_keyframe:
            keyframes.append({
                "frame_index": frame_index,
                "timestamp_ms": timestamp_ms,
                "frame_image_path": str(image_path),
                "phash": phash,
                "ocr_text": "",  # P2 调用 DeepSeek Vision 填充
            })
            last_keyframe_phash = phash

    return keyframes


def hamming_distance(s1: str, s2: str) -> int:
    """计算两个 hex pHash 字符串的海明距离"""
    if len(s1) != len(s2):
        # pHash 长度不一致,取较长字符串的长度作为海明距离上限
        return max(len(s1), len(s2))
    return sum(c1 != c2 for c1, c2 in zip(s1, s2))


# ========================================================================
# OCR 接口(P2 实现,目前返回占位)
# ========================================================================


def ocr_frame(image_path: Path) -> str:
    """
    识别帧内文字,返回纯文本

    P1: 返回占位文本
    P2: 调用 DeepSeek Vision API
    """
    if not DEEPSEEK_API_KEY:
        # 无 API Key,返回占位
        return f"[OCR待填充 frame={image_path.name}]"

    # P2 实现:调用 DeepSeek Vision
    # TODO: P2 实现
    return f"[OCR待填充 frame={image_path.name}]"


def ocr_keyframes(keyframes: List[Dict]) -> List[Dict]:
    """对关键帧列表逐一调用 OCR"""
    result = []
    for kf in keyframes:
        image_path = Path(kf["frame_image_path"])
        ocr_text = ocr_frame(image_path)
        kf_copy = kf.copy()
        kf_copy["ocr_text"] = ocr_text
        result.append(kf_copy)
    return result


# ========================================================================
# B 稿格式化
# ========================================================================


def format_timestamp(ms: int) -> str:
    """毫秒转 [Nm Ns] 格式"""
    total_sec = ms // 1000
    return f"{total_sec // 60}m{total_sec % 60}s"


def build_b_manuscript(keyframes: List[Dict]) -> List[str]:
    """
    将关键帧 OCR 结果合并为 B 稿
    合并相邻同文本的关键帧
    """
    lines = []
    last_text = None

    for kf in keyframes:
        text = kf["ocr_text"].strip()
        if not text:
            continue

        # 跳过占位文本
        if text.startswith("[OCR待填充"):
            text = ""

        if text and text != last_text:
            ts = format_timestamp(kf["timestamp_ms"])
            lines.append(f"[{ts}] {text}")
            last_text = text

    return lines


def write_b_manuscript(lines: List[str], output_path: Path) -> Path:
    """写入 B 稿 txt"""
    output_path.parent.mkdir(parents=True, exist_ok=True)
    with open(output_path, "w", encoding="utf-8") as f:
        f.write("\n".join(lines))
    return output_path


# ========================================================================
# 主流程
# ========================================================================


def split_video(
    video_path: Path,
    output_dir: Path,
    episode_id: str,
    phash_threshold: int = 8,
    fps: int = 1,
) -> Dict[str, any]:
    """
    视频双路拆分主流程

    参数:
        video_path: 输入视频路径
        output_dir: 输出目录(work/ 路径)
        episode_id: 节目 ID
        phash_threshold: pHash 海明距离阈值,默认 8
        fps: 抽帧帧率,默认 1(每秒一帧)

    返回:
        {
            "b_manuscript_path": Path,
            "audio_path": Path,
            "keyframes_path": Path,
            "keyframe_count": int,
        }
    """
    video_path = Path(video_path)
    output_dir = Path(output_dir)

    if not video_path.exists():
        raise FileNotFoundError(f"视频文件不存在: {video_path}")

    # 创建输出目录
    output_dir.mkdir(parents=True, exist_ok=True)
    frames_dir = output_dir / "frames"
    frames_dir.mkdir(parents=True, exist_ok=True)

    print(f"[video_split] 开始处理: {video_path.name}")
    print(f"[video_split] 抽帧 fps={fps}, pHash threshold={phash_threshold}")

    # ---- A 路:抽帧 + pHash 检测 + OCR ----
    print("[video_split] A路:抽帧...")
    frames = extract_frames(video_path, output_dir, fps=fps)
    print(f"[video_split] 抽帧完成,共 {len(frames)} 帧")

    print("[video_split] pHash 变化检测...")
    keyframes = find_keyframes(frames, threshold=phash_threshold)
    print(f"[video_split] 检测到 {len(keyframes)} 个关键帧")

    print("[video_split] OCR 关键帧...")
    keyframes = ocr_keyframes(keyframes)
    print(f"[video_split] OCR 完成")

    # ---- B 路:音频提取 ----
    print("[video_split] B路:提取音频...")
    audio_path = output_dir / "audio_16k.wav"
    extract_audio(video_path, audio_path)
    print(f"[video_split] 音频提取完成: {audio_path}")

    # ---- 输出产物 ----
    # B 稿
    b_lines = build_b_manuscript(keyframes)
    b_manuscript_path = output_dir / "b_manuscript.txt"
    write_b_manuscript(b_lines, b_manuscript_path)
    print(f"[video_split] B稿写入: {b_manuscript_path} ({len(b_lines)} 行)")

    # 关键帧索引 JSON
    keyframes_data = {
        "video_path": str(video_path),
        "fps_sampled": fps,
        "phash_threshold": phash_threshold,
        "keyframes": keyframes,
    }
    keyframes_path = output_dir / "keyframes.json"
    with open(keyframes_path, "w", encoding="utf-8") as f:
        json.dump(keyframes_data, f, ensure_ascii=False, indent=2)
    print(f"[video_split] 关键帧索引写入: {keyframes_path}")

    # 清理临时帧文件(可选,保留供调试)
    # shutil.rmtree(frames_dir)

    return {
        "b_manuscript_path": str(b_manuscript_path),
        "audio_path": str(audio_path),
        "keyframes_path": str(keyframes_path),
        "keyframe_count": len(keyframes),
    }