"""字幕 OCR 汇总节点。

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

from __future__ import annotations

import json
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")

# 默认垃圾词：无文字帧的模型输出可能反复出现这些词。
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 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
    }

    # 单帧 OCR：并行池的工作函数，返回该帧识别文本（失败/空/超长均返回空串）。
    def ocr_frame(payload) -> str:
        index, item = payload
        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":
            # 单帧失败不中断整体，跳过该帧继续汇总。
            logger.warning("帧 %d OCR 失败，跳过: %s", index, response.error)
            return ""
        logger.info("帧 %d/%d OCR 完成: %r", index + 1, len(manifest), response.outputs.get("text"))
        text = str(response.outputs.get("text", "")).strip()
        if not text:
            return ""
        # 超长输出视为模型异常（重复循环等），直接报错并跳过该帧。
        if len(text) > max_result_chars:
            logger.warning(
                "帧 %d OCR 输出超长(%d > %d)，跳过: %r",
                index, len(text), max_result_chars, text[:60],
            )
            return ""
        return text

    # 进度日志：打印已识别帧数、总数与平均处理速度（帧/s）。
    def log_progress(done: int, total: int, rate: float) -> None:
        logger.info("OCR 进度 %d/%d 帧 (%.1f 帧/s)", done, total, rate)

    # 自适应并发调用 vlm-ocr：10s 窗口内平均响应 < 0.3s 则加 1 线程（上限
    # pool_max_workers），> pool_slow_threshold 则减 1 线程（下限 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)),
    )
    texts = pool.map(list(enumerate(manifest)))

    # kept 元素为 (起始帧时间, 最后可见帧时间, 文本)；按帧顺序合并连续相同字幕。
    kept: list[tuple[float, float, str]] = []
    for index, text in enumerate(texts):
        if not text:
            continue
        time = float(manifest[index]["time"])
        # 连续帧相同字幕合并为一条（字幕停留多帧属正常现象）：
        # 仅更新最后可见帧时间，起始时间保持首次出现。
        if kept and kept[-1][2] == text:
            kept[-1] = (kept[-1][0], time, text)
            continue
        kept.append((time, time, text))

    output_dir = Path(request.output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)
    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)},
    )