"""字幕 OCR 汇总节点。

读取 frame-extract 产出的 frames.json，逐帧调用 vlm-ocr 节点识别字幕文字；
过滤无文字帧的垃圾输出（glm-ocr 在空帧上会输出无用文字），折叠模型重复
循环输出，合并连续相同的字幕（记录最后可见帧时间），最终组装为带时间轴的
SRT 基准数据：每条字幕消失时间 = 最后可见帧时间 + 采样间隔（间隔从帧清单
时间轴推导），与视频烧录时间对齐。
"""

from __future__ import annotations

import json
import threading
from pathlib import Path

from nodes.adaptive_pool import AdaptiveThreadPool
from nodes.whisper import format_timestamp
from wov_app import registry
from wov_app.logging import get_logger
from wov_sdk.models import InvokeRequest, InvokeResponse

logger = get_logger("subtitle-ocr")

# 节点级断点存档的写锁：多线程 OCR 并发完成时串行化追加写，避免行交错。
_partial_lock = threading.Lock()

# 断点存档文件名：每行 {"frame": 帧序号(0-based), "text": 该帧OCR文本}。
# OCR 每帧完成后立即追加一行；进程重启后从存档恢复已处理帧，只对未处理
# 帧重新调用 vlm-ocr——2 小时视频级任务中断后不浪费已完成的帧。
_PARTIAL_NAME = "ocr_partial.jsonl"

# 暂停信号文件名：位于 run 根目录（<storage>/runs/<run_id>/paused.flag），
# 由暂停接口写入、继续/重试/删除时清除；节点逐帧检查，存在即中止。
_PAUSE_FLAG = "paused.flag"


class PauseRequested(Exception):
    """节点内暂停信号：OCR 检测到任务被暂停后抛出，由调度器保持 PAUSED。

    不把暂停误报为 FAILED：调度器捕获异常时若任务状态已是 PAUSED，
    则保持暂停等待用户 resume，从断点存档继续未处理帧。
    """


def _load_partial(output_dir: Path) -> dict[int, str]:
    """读取节点级断点存档，返回 {帧序号: OCR 文本}；无存档时返回空字典。"""
    path = output_dir / _PARTIAL_NAME
    if not path.is_file():
        return {}
    result: dict[int, str] = {}
    for line in path.read_text(encoding="utf-8").splitlines():
        if not line.strip():
            continue
        try:
            item = json.loads(line)
        except json.JSONDecodeError:
            # 进程被杀时可能残留半行写入：跳过该行，对应帧视为未处理。
            continue
        # 存档区分"成功空帧/skipped（超长跳过）"与失败——失败不写成功存档。
        # 无 status 的历史空串无法区分超时与无文字，需重新识别；非空结果可复用。
        text = str(item["text"])
        if item.get("status") in ("completed", "skipped") or ("status" not in item and text):
            result[int(item["frame"])] = text
    return result


def _format_eta(seconds: float) -> str:
    """把剩余秒数格式化为可读的预计完成时间（如 34分13秒 / 2小时05分）。"""
    total = max(0, int(seconds))
    hours, remainder = divmod(total, 3600)
    minutes, secs = divmod(remainder, 60)
    if hours:
        return f"{hours}小时{minutes:02d}分"
    if minutes:
        return f"{minutes}分{secs:02d}秒"
    return f"{secs}秒"


def _eta_suffix(done: int, total: int, rate: float) -> str:
    """根据当前处理速度计算剩余时间后缀（供进度日志追加）。

    剩余时间 = 剩余帧数 / 当前速度；速度为 0（刚开始或耗时不可测）时
    返回空串，进度提示不显示 ETA。
    """
    if rate <= 0:
        return ""
    return f", 预计剩余 {_format_eta((total - done) / rate)}"
# 默认垃圾词：无文字帧的模型输出可能反复出现这些词。
def _sampling_interval(manifest: list[dict], default: float) -> float:
    """从帧清单时间轴推导采样间隔（相邻帧时间差的中位数）。

    帧时间由 frame-extract 按固定间隔生成，取相邻差的中位数即可还原真实
    采样间隔，避免结束时间与抽取参数不一致。清单不足两帧时回退默认值。
    """
    diffs = [
        float(manifest[i + 1]["time"]) - float(manifest[i]["time"])
        for i in range(len(manifest) - 1)
        if float(manifest[i + 1]["time"]) > float(manifest[i]["time"])
    ]
    if not diffs:
        return default
    diffs.sort()
    # 取 3 位小数：与 frame-extract 的 round(秒,3) 时间戳精度一致，避免浮点漂移。
    return round(diffs[len(diffs) // 2], 3)


def _assemble_srt(
    kept: list[tuple[float, float, str]],
    interval_seconds: float,
) -> list[str]:
    """把 (起始帧时间, 最后可见帧时间, 文本) 序列组装为 SRT 行列表。

    每条字幕的结束时间 = 最后可见帧时间 + 采样间隔：字幕在最后一个被识别
    到的帧之后的一个采样间隔内消失，与视频烧录时间对齐；两段字幕之间的
    空白段（无字幕帧）不再被并入前一条字幕。
    """
    lines: list[str] = []
    for index, (start, last_seen, text) in enumerate(kept):
        end = last_seen + interval_seconds
        lines.extend(
            [
                str(index + 1),
                f"{format_timestamp(start)} --> {format_timestamp(end)}",
                text,
                "",
            ]
        )
    return lines


def _merge_kept(
    manifest: list[dict], texts: list[str]
) -> list[tuple[float, float, str]]:
    """按 manifest 时间轴把逐帧 OCR 文本合并为字幕条目。

    kept 元素为 (起始帧时间, 最后可见帧时间, 文本)：空文本（无文字帧）跳过；
    连续帧相同字幕合并为一条（字幕停留多帧属正常现象），仅更新最后可见帧
    时间，起始时间保持首次出现。要求 texts 与 manifest 按帧顺序一一对应
    （调用方保证），重组装旧数据时也复用此逻辑保证行为一致。
    """
    kept: list[tuple[float, float, str]] = []
    for index, text in enumerate(texts):
        if not text:
            continue
        time = float(manifest[index]["time"])
        if kept and index > 0 and texts[index - 1] == text:
            kept[-1] = (kept[-1][0], time, text)
            continue
        kept.append((time, time, text))
    return kept


def invoke(request: InvokeRequest) -> InvokeResponse:
    """逐帧 OCR 并汇总字幕，产物为 subtitle.srt。"""
    manifest_uri = request.inputs.get("frames_manifest")
    if not manifest_uri:
        return InvokeResponse(status="failed", error="frames_manifest is required")
    manifest_path = Path(manifest_uri)
    if not manifest_path.is_file():
        return InvokeResponse(status="failed", error="frames manifest not found")
    manifest = json.loads(manifest_path.read_text(encoding="utf-8"))

    # 结果长度上限：超过即视为模型异常（重复循环等），该帧直接报错跳过。
    max_result_chars = int(request.params.get("max_result_chars", 200))
    # 采样间隔优先从帧清单时间轴推导（与 frame-extract 实际抽取间隔一致），
    # 参数仅作清单退化时的兜底。
    interval = _sampling_interval(
        manifest, float(request.params.get("interval_seconds", 2.0))
    )
    # 透传给 vlm-ocr 的参数（仅传已提供的，避免覆盖其默认值）。
    vlm_params = {
        key: request.params.get(key)
        for key in (
            "model", "ollama_host", "prompt", "timeout_seconds", "keep_alive",
            "temperature", "repeat_penalty", "num_predict",
        )
        if request.params.get(key) is not None
    }

    # 节点级断点：读取已处理帧存档，只对未处理帧调用 vlm-ocr。
    output_dir = Path(request.output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)
    partial_path = output_dir / _PARTIAL_NAME
    partial_texts = _load_partial(output_dir)
    pending = [(i, item) for i, item in enumerate(manifest) if i not in partial_texts]
    if partial_texts:
        logger.info(
            "检测到节点级断点：%d/%d 帧已处理，本次只处理剩余 %d 帧",
            len(partial_texts), len(manifest), len(pending),
        )

    # 单帧 OCR：成功返回文字或空串；临时失败抛异常，不写成功存档。
    # 超长输出按既有规则确定跳过，以 skipped 存档区别于成功无文字。
    def ocr_frame(payload) -> str:
        index, item = payload
        # 暂停检查：调度器置 PAUSED 并向 run 根写入 paused.flag 后，工作线程
        # 立即中止（不 OCR、不写存档，该帧恢复时重跑），让 map 快速结束。
        if (Path(request.output_dir).parent.parent / _PAUSE_FLAG).exists():
            # 抑制后续进度回调：队列中剩余大量帧会快速失败退出，避免逐项
            # 打印"OCR 进度"导致日志井喷。
            pool.cancel()
            raise PauseRequested(f"OCR 被暂停（run {request.run_id}）")
        text = ""
        try:
            response = registry.invoke(
                "vlm-ocr",
                InvokeRequest(
                    run_id=request.run_id,
                    node_instance_id="",
                    inputs={"image_uri": str(item["image_uri"])},
                    params=vlm_params,
                    output_dir=str(Path(request.output_dir) / "ocr_frames" / f"{index:04d}"),
                ),
            )
            if response.status != "completed":
                raise RuntimeError(f"帧 {index} OCR 失败: {response.error}")
            if not isinstance(response.outputs.get("text"), str):
                raise ValueError(f"帧 {index} OCR 缺少有效 text")
        except Exception:
            pool.report_failure()
            raise
        status = "completed"
        logger.info("帧 %d/%d OCR 完成: %r", index + 1, len(manifest), response.outputs.get("text"))
        text = response.outputs["text"].strip()
        if len(text) > max_result_chars:
            # 超长输出视为模型异常（重复循环等），直接报错并跳过该帧。
            logger.warning(
                "帧 %d OCR 输出超长(%d > %d)，跳过: %r",
                index, len(text), max_result_chars, text[:60],
            )
            text = ""
            status = "skipped"
        # 只存成功或确定跳过的结果，网络故障保持未处理，恢复时重新识别。
        with _partial_lock:
            with partial_path.open("a", encoding="utf-8") as fh:
                fh.write(json.dumps({"frame": index, "text": text, "status": status}, ensure_ascii=False) + "\n")
        return text

    if pending:
        # 进度日志：打印已识别帧数、总数、平均处理速度（帧/s）、最近窗口平均
        # 单帧耗时、当前线程数与**预计剩余完成时间**（剩余帧/当前速度）——
        # 便于判断多线程是否因单帧处理过慢而未启用（窗口平均响应 ≥
        # fast_threshold 时自适应池不会扩容）以及整体还需要多久。
        def log_progress(done: int, total: int, rate: float, avg_time: float, workers: int) -> None:
            logger.info(
                "OCR 进度 %d/%d 帧 (%.1f 帧/s, 平均 %.2fs/帧, 线程 %d/%d%s)",
                done, total, rate, avg_time, workers, pool.max_workers,
                _eta_suffix(done, total, rate),
            )

        # 自适应并发调用 vlm-ocr：按窗口内平均响应快慢增/减额度（上限
        # pool_max_workers、下限 1），避免盲目并发压垮本地 Ollama。
        pool = AdaptiveThreadPool(
            worker=ocr_frame,
            on_progress=log_progress,
            min_workers=int(request.params.get("pool_min_workers", 1)),
            max_workers=int(request.params.get("pool_max_workers", 16)),
            window_seconds=float(request.params.get("pool_window_seconds", 10.0)),
            fast_threshold=float(request.params.get("pool_fast_threshold", 0.3)),
            slow_threshold=float(request.params.get("pool_slow_threshold", 1.0)),
        )
        pending_texts = pool.map(pending)
        # 任一工作线程检测到暂停信号即整体中止：已写盘的断点存档保留，
        # resume 后从剩余帧续跑；以 failed 返回让调度器保持 PAUSED（不误报失败）。
        if any(isinstance(text, PauseRequested) for text in pending_texts):
            return InvokeResponse(status="failed", error=f"OCR 被暂停（run {request.run_id}）")
        # 失败帧在收紧后的并发额度下重试一轮，成功帧（含空帧）不重复调用。
        failed = [payload for payload, text in zip(pending, pending_texts) if isinstance(text, Exception)]
        new_by_index = {i: t for (i, _item), t in zip(pending, pending_texts) if not isinstance(t, Exception)}
        if failed:
            logger.warning("OCR %d 帧失败，重试一次", len(failed))
            retried = pool.map(failed)
            for (index, _item), text in zip(failed, retried):
                if isinstance(text, Exception):
                    # 不发布残缺字幕；已成功写盘的其他帧下次直接复用。
                    return InvokeResponse(status="failed", error=f"OCR 帧 {index} 重试失败: {text}")
                new_by_index[index] = text
    else:
        new_by_index = {}

    # 完整帧序文本：优先取存档，其次取本次新增（未处理帧必有其一条目）。
    texts = [
        partial_texts.get(i, new_by_index.get(i, ""))
        for i in range(len(manifest))
    ]
    # 按帧顺序合并连续相同字幕（与重组装旧数据共用 _merge_kept）。
    kept = _merge_kept(manifest, texts)
    # output_dir 已在断点初始化时创建（mkdir 幂等），此处直接使用。
    output_path = output_dir / "subtitle.srt"
    output_path.write_text("\n".join(_assemble_srt(kept, interval)), encoding="utf-8")
    logger.info("字幕汇总完成: %d 条", len(kept))
    return InvokeResponse(
        status="completed",
        outputs={"srt_uri": str(output_path), "count": len(kept)},
    )