# -*- coding: utf-8 -*-
"""
阶段 A:抽帧 + OCR(产出原始缓存 ocr_raw.jsonl)
================================================
策略:OCR 优先,每帧都 OCR。绝不使用亮度判空/dHash/IoU 过滤。
按 idx 断点续跑:读取已有 ocr_raw.jsonl 中最大 idx,从 max_idx+1 继续。
"""
import base64
import json
import os
import subprocess
import sys
import threading
from concurrent.futures import ThreadPoolExecutor, as_completed
from pathlib import Path

import requests

# 根治 Windows 控制台 GBK 崩溃:强制 stdout 用 UTF-8
sys.stdout.reconfigure(encoding="utf-8", errors="replace")

# ========================================================================
# 全局参数(放文件顶部做常量,方便后续调)
# ========================================================================
FPS = 1
CROP_BOTTOM_RATIO = 0.2
SIM_THRESHOLD = 0.85
OLLAMA_URL = "http://localhost:11434/api/generate"
MODEL = "deepseek-ocr"
OCR_PROMPT = "Free OCR."
# 并发 OCR 线程数(环境变量 OCR_NUM_WORKERS 可覆盖)。
# 需配合 Ollama 服务端 OLLAMA_NUM_PARALLEL >= 此值,否则服务端仍排队、提速无效。
NUM_WORKERS = int(os.environ.get("OCR_NUM_WORKERS", "4"))

# 工作目录(脚本所在目录)
WORK_DIR = Path(__file__).resolve().parent
SOURCE_VIDEO_DIR = WORK_DIR / "source"
FRAMES_DIR = WORK_DIR / "frames_v2"
OCR_RAW_PATH = WORK_DIR / "ocr_raw.jsonl"


# ========================================================================
# A1: 定位源视频
# ========================================================================
def find_source_video():
    """在工作目录 source/ 或根目录下定位 .mp4 或 .mkv 源视频"""
    for ext in (".mp4", ".mkv"):
        candidates = list(SOURCE_VIDEO_DIR.glob(f"*{ext}"))
        if candidates:
            return candidates[0]
    # fallback: 也检查工作目录根(ep002 等视频直接在根目录的情况)
    for ext in (".mp4", ".mkv"):
        candidates = list(WORK_DIR.glob(f"*{ext}"))
        if candidates:
            return candidates[0]
    raise FileNotFoundError(
        f"在 {SOURCE_VIDEO_DIR} 和 {WORK_DIR} 中均未找到 .mp4 或 .mkv 文件"
    )


def get_video_info(video_path):
    """用 ffprobe 获取时长(秒)和分辨率(宽x高)"""
    cmd = [
        "ffprobe", "-v", "error",
        "-select_streams", "v:0",
        "-show_entries", "stream=duration,width,height",
        "-of", "json",
        str(video_path),
    ]
    result = subprocess.run(cmd, capture_output=True, text=True, encoding="utf-8", errors="replace")
    if result.returncode != 0:
        raise RuntimeError(f"ffprobe 失败:\n{result.stderr}")
    info = json.loads(result.stdout)
    stream = info["streams"][0]
    duration = float(stream.get("duration", 0))
    width = stream["width"]
    height = stream["height"]
    return duration, width, height


# ========================================================================
# A2: ffmpeg 抽帧
# ========================================================================
def extract_frames(video_path):
    """
    用 ffmpeg 按 1fps 抽帧,裁切下方 20% 字幕区域。
    输出到 frames_v2/frame_%04d.png(从 0001 开始)。
    返回按文件名排序的帧路径列表。
    """
    FRAMES_DIR.mkdir(parents=True, exist_ok=True)

    # 清理旧帧避免残留
    for old in FRAMES_DIR.glob("frame_*.png"):
        old.unlink()

    crop_expr = (
        f"fps={FPS},"
        f"crop=iw:ih*{CROP_BOTTOM_RATIO}:0:ih*{1 - CROP_BOTTOM_RATIO}"
    )
    frame_pattern = str(FRAMES_DIR / "frame_%04d.png")

    cmd = [
        "ffmpeg",
        "-i", str(video_path),
        "-vf", crop_expr,
        "-q:v", "2",
        frame_pattern,
        "-y",
    ]

    result = subprocess.run(cmd, capture_output=True, text=True, encoding="utf-8", errors="replace")
    if result.returncode != 0:
        raise RuntimeError(f"ffmpeg 抽帧失败:\n{result.stderr}")

    frames = sorted(FRAMES_DIR.glob("frame_*.png"))
    return frames


# ========================================================================
# OCR 单帧
# ========================================================================
def ocr_single_frame(image_path):
    """
    调用本地 Ollama OCR 识别单帧。
    API: POST /api/generate,单轮,图片走 base64,stream=false。
    返回 strip 后的文本;异常向上抛出让调用方处理。
    """
    with open(image_path, "rb") as fh:
        img_bytes = fh.read()
    img_b64 = base64.b64encode(img_bytes).decode("utf-8")

    body = {
        "model": MODEL,
        "prompt": OCR_PROMPT,
        "images": [img_b64],
        "stream": False,
    }

    resp = requests.post(OLLAMA_URL, json=body, timeout=120)
    resp.raise_for_status()
    data = resp.json()
    text = data.get("response", "").strip()
    return text


# ========================================================================
# A3: 单帧验证
# ========================================================================
def test_single_frame():
    """
    A3: 验证单帧 OCR。
    对 frame_0076.png 调一次 OCR,打印原始返回。
    若对不上预期则中止,不进入批量。
    """
    test_frame = FRAMES_DIR / "frame_0076.png"
    if not test_frame.exists():
        print("[A3] [WARN] test frame not found:", test_frame)
        print("[A3] skip single-frame check")
        return

    print("[A3] test frame:", test_frame.name)
    try:
        raw_text = ocr_single_frame(test_frame)
    except Exception as exc:
        print("[A3] [FAIL] OCR call failed:", exc)
        print("[A3] check Ollama / deepseek-ocr, abort.")
        sys.exit(1)

    print("[A3] raw return:", repr(raw_text))

    # 通用冒烟检查:OCR 返回非空中文即视为链路正常(不绑定某期具体台词);
    # 真正的"Ollama 挂了"由上面 ocr_single_frame 抛异常 + sys.exit(1) 兜底。
    if any("一" <= ch <= "鿿" for ch in raw_text):
        print("[A3] [OK] OCR 链路正常(返回非空中文)。")
    else:
        print("[A3] [WARN] OCR 返回无中文,链路可能异常,但不阻断,继续批量:", repr(raw_text))


# ========================================================================
# A4: 批量 OCR(按 idx 断点续跑 + 单帧容错)
# ========================================================================
def batch_ocr(frames):
    """
    并发 OCR(线程池 NUM_WORKERS 路) + 按 idx 断点续跑 + 单帧容错。
    - 读 ocr_raw.jsonl 已写入的 idx 集合(并发中断可能留空洞,用集合而非 max),只补未写帧
    - NUM_WORKERS 个线程并发调 Ollama;每帧独立,输出与串行逐帧完全一致
    - 写盘按完成顺序、加锁逐条 append + flush(stage_b 读后会按 idx 重排,顺序无所谓)
    - 单帧 OCR 报错写 error 字段,不崩整个进程
    """
    total = len(frames)

    # --- 读已写入的 idx 集合(含 error 记录,避免重跑出重复 idx 行) ---
    done_idx = set()
    if OCR_RAW_PATH.exists():
        with open(OCR_RAW_PATH, "r", encoding="utf-8") as fh:
            for line in fh:
                line = line.strip()
                if not line:
                    continue
                try:
                    rec = json.loads(line)
                    if "idx" in rec:
                        done_idx.add(rec["idx"])
                except json.JSONDecodeError:
                    pass

    # --- 构建 idx -> frame_path 映射 ---
    idx_to_path = {}
    for fp in frames:
        idx = int(fp.stem.split("_")[1])
        idx_to_path[idx] = fp

    pending = [idx for idx in sorted(idx_to_path) if idx not in done_idx]
    if not pending:
        print(f"[A4] all {total} frames already done, skip.")
        return

    print(
        f"[A4] 已完成 {len(done_idx)}/{total},待处理 {len(pending)} 帧,"
        f"并发 {NUM_WORKERS} 路"
    )

    def _ocr_one(idx):
        t_sec = idx - 1
        frame_path = idx_to_path.get(idx)
        if frame_path is None:
            return {"idx": idx, "t_sec": t_sec, "text": "", "error": "frame file missing"}
        try:
            text = ocr_single_frame(frame_path)
            return {"idx": idx, "t_sec": t_sec, "text": text}
        except Exception as exc:
            return {"idx": idx, "t_sec": t_sec, "text": "", "error": str(exc)}

    write_lock = threading.Lock()
    done_count = len(done_idx)

    with open(OCR_RAW_PATH, "a", encoding="utf-8") as fh:
        with ThreadPoolExecutor(max_workers=NUM_WORKERS) as pool:
            futures = {pool.submit(_ocr_one, idx): idx for idx in pending}
            for fut in as_completed(futures):
                record = fut.result()
                with write_lock:
                    fh.write(json.dumps(record, ensure_ascii=False) + "\n")
                    fh.flush()
                    done_count += 1
                    if done_count % 50 == 0:
                        print(f"{done_count}/{total}", flush=True)

    print(f"{total}/{total}", flush=True)
    print(f"[A4] OCR done, {total} records in {OCR_RAW_PATH}")


# ========================================================================
# 主流程
# ========================================================================
def main():
    # A1
    video_path = find_source_video()
    duration, width, height = get_video_info(video_path)
    print(f"[A1] video: {video_path}")
    print(f"[A1] duration: {duration:.1f}s  resolution: {width}x{height}")

    # A2: skip re-extract if frames already exist
    FRAMES_DIR.mkdir(parents=True, exist_ok=True)
    existing = sorted(FRAMES_DIR.glob("frame_*.png"))
    if existing:
        print(f"[A2] {len(existing)} frames already exist, skip re-extract")
        frames = existing
    else:
        print(f"[A2] ffmpeg extract (fps={FPS}, crop_bottom_ratio={CROP_BOTTOM_RATIO})...")
        frames = extract_frames(video_path)
        print(f"[A2] extracted {len(frames)} frames -> {FRAMES_DIR}")

    # A3
    test_single_frame()

    # A4
    batch_ocr(frames)

    print("\nstage A done.")


if __name__ == "__main__":
    main()