"""faster-whisper ASR 节点。

单体版中作为进程内节点模块，由调度器直接调用。模型权重默认优先从本地
目录加载，避免从远端下载，仅在本地找不到模型时才回退到远端 large-v3。
CUDA 动态库通过 ctypes 在进程内预加载，替代分布式版的 LD_LIBRARY_PATH 注入。
"""

from __future__ import annotations

import ctypes
import os
import subprocess
import sysconfig
import time
import wave
from pathlib import Path

from nodes.ffmpeg import _ffmpeg_bin
from wov_app.logging import get_logger
from wov_sdk.models import InvokeRequest, InvokeResponse

# 转写进度日志：输出到主进程控制台，长音频分块时可见每块进度。
logger = get_logger("whisper")

# 暂停信号文件名：位于 run 根目录（<storage>/runs/<run_id>/paused.flag），
# 与 subtitle-ocr 节点约定一致；批量暂停时由暂停接口写入，分块间检查即中止。
PAUSE_FLAG = "paused.flag"

def _is_windows() -> bool:
    """判断当前是否为 Windows，供测试单独注入覆盖。"""
    return os.name == "nt"


def _load_cuda_libraries() -> None:
    """在进程内预加载 pip 安装的 NVIDIA 动态库。

    分布式版通过给节点子进程注入 LD_LIBRARY_PATH（Windows 为 PATH）解决；
    单体版没有进程边界，必须在导入 faster-whisper 前加载 nvidia 轮子自带
    的 .so/.dll，否则 ctranslate2 初始化 CUDA 时找不到 libcublas.so.12。
    """
    # site-packages 目录，nvidia 各包的动态库位于其下。
    site_packages = Path(sysconfig.get_paths()["purelib"])
    for vendor in ("cublas", "cudnn", "cuda_nvrtc"):
        # Windows 轮子把 dll 放在 bin/，Linux 放在 lib/。
        for subdir in ("bin", "lib"):
            lib_dir = site_packages / "nvidia" / vendor / subdir
            if not lib_dir.is_dir():
                continue
            if _is_windows():
                # Windows 通过 DLL 搜索目录注册，等价于进程内 PATH 注入。
                os.add_dll_directory(str(lib_dir))
            else:
                for so_file in sorted(lib_dir.glob("*.so*")):
                    try:
                        ctypes.CDLL(str(so_file))
                    except OSError:
                        # 个别依赖缺失（如 libcudart）时跳过，交由 ctranslate2 报错。
                        continue

def _local_model_candidates() -> list[Path]:
    """返回本地模型候选目录：单体根目录 model/ 优先，其次 nodes/ 同级 model/。

    单体根目录 model/ 对应仓库根下的 model/faster-whisper-large-v3，
    nodes/ 同级 model/ 允许部署时把权重随代码目录一起携带。
    """
    monolith_root = Path(__file__).resolve().parent.parent
    return [
        monolith_root / "model" / "faster-whisper-large-v3",
        monolith_root / "nodes" / "model" / "faster-whisper-large-v3",
    ]


def resolve_model_path(
    params: dict,
    env: dict | None = None,
    candidates: list[Path] | None = None,
) -> str:
    """按 参数 > 环境变量 > 本地候选目录 > 远端 large-v3 的顺序解析模型路径。

    本地优先是默认行为：只要候选目录存在且包含 model.bin 就使用本地权重，
    避免从 Hugging Face 下载；远端下载仅在全部本地候选缺失时作为兜底。
    参数/环境变量传入的是裸模型名（不含路径分隔符）时，会先在本地模型
    目录（model/）下按名解析，方便工作流直接引用下载好的模型。
    candidates 参数供测试注入临时目录，默认使用 _local_model_candidates()。
    """
    env = env if env is not None else os.environ
    candidates = candidates if candidates is not None else _local_model_candidates()
    explicit = params.get("model_path") or env.get("WHISPER_MODEL_PATH")
    if explicit:
        explicit_str = str(explicit)
        # 裸模型名按 <模型目录>/<名称> 在本地解析，例如
        # "whisper-large-v3-translate-zh-v0.1-lt-ct2"。
        if "/" not in explicit_str and "\\" not in explicit_str:
            named = candidates[0].parent / explicit_str
            if (named / "model.bin").is_file():
                return str(named)
        return explicit_str
    for candidate in candidates:
        # model.bin 是 CTranslate2 权重的必需文件，存在才认为模型完整。
        if candidate.is_dir() and (candidate / "model.bin").is_file():
            return str(candidate)
    return "large-v3"

def format_timestamp(seconds: float) -> str:
    """把秒数格式化为 SRT 时间戳，例如 01:00:00,500。"""
    # 先换算成毫秒再逐级拆分为时/分/秒/毫秒，避免浮点误差。
    total_ms = int(seconds * 1000)
    hours, remainder = divmod(total_ms, 3600000)
    minutes, remainder = divmod(remainder, 60000)
    secs, millis = divmod(remainder, 1000)
    return f"{hours:02d}:{minutes:02d}:{secs:02d},{millis:03d}"


def _split_audio(
    audio_path: Path,
    output_dir: Path,
    chunk_seconds: int,
    ffmpeg_bin: str | None,
) -> list[Path]:
    """用 ffmpeg 把音频切成 chunk_seconds 秒一块的 wav，返回块路径列表。

    分块是应用层工程策略（内存有界、失败粒度小），与模型 30s 窗口无关；
    whisper 训练与推理都按 30s 窗口解码，任意块大小都适用。以下情况回退
    为整段单次转写：chunk_seconds <= 0、找不到 ffmpeg、切块失败、音频本身
    不足一块（ffmpeg 产出单块）。
    """
    if chunk_seconds <= 0 or not ffmpeg_bin:
        return [audio_path]
    chunk_dir = output_dir / "chunks"
    chunk_dir.mkdir(parents=True, exist_ok=True)
    pattern = str(chunk_dir / "chunk_%03d.wav")
    # 音频已是 16kHz 单声道 WAV，流拷贝切块即可，无需重编码。
    result = subprocess.run(
        [
            ffmpeg_bin,
            "-y",
            "-i",
            str(audio_path),
            "-f",
            "segment",
            "-segment_time",
            str(chunk_seconds),
            "-c",
            "copy",
            pattern,
        ],
        capture_output=True,
        text=True,
    )
    if result.returncode != 0:
        # 切块失败（输入损坏等）回退整段，不让转写流程中断。
        return [audio_path]
    chunks = sorted(chunk_dir.glob("chunk_*.wav"))
    return chunks or [audio_path]


def _wav_duration_seconds(path: Path, fallback: float) -> float:
    """从 WAV 头精确读取时长；文件非法/非 WAV 时用 fallback 兜底。

    分块偏移必须用每块的实际时长累积，而不是 块序号×块长 的假设值——
    ffmpeg 切出的块实际时长并不精确等于块长（如 60.05s），假设值会随
    块数累积漂移，造成字幕时间轴逐渐错位。
    """
    try:
        with wave.open(str(path), "rb") as wav:
            rate = wav.getframerate()
            return wav.getnframes() / rate if rate else fallback
    except (wave.Error, EOFError, OSError):
        # 文件损坏/非 WAV（如 mp4 直传）时用 fallback 兜底。
        return fallback


def _append_srt_lines(lines: list[str], segments, offset: float, start_index: int) -> int:
    """把一段转写结果按 SRT 格式追加到 lines，时间加上 offset 偏移。

    分块合并时每块 offset 为前面所有块的实际时长累积；单次调用 offset=0。
    每写出一条字幕就打印其编号与完整视频角度的时间范围，便于对照对齐。
    返回本段新增的条数，用于全局序号递增。
    """
    count = 0
    for segment in segments:
        # 完整视频角度的时间 = 模型预测时间 + 累积偏移。
        start_time = segment.start + offset
        end_time = segment.end + offset
        lines.extend(
            [
                str(start_index + count),
                f"{format_timestamp(start_time)} --> {format_timestamp(end_time)}",
                segment.text.strip(),
                "",
            ]
        )
        logger.info(
            "分段 #%d: %s --> %s",
            start_index + count,
            format_timestamp(start_time),
            format_timestamp(end_time),
        )
        count += 1
    return count

def invoke(request: InvokeRequest) -> InvokeResponse:
    """转写音频并生成 SRT 字幕，产物为 transcript.srt。"""
    audio_uri = request.inputs.get("audio_uri")
    if not audio_uri:
        return InvokeResponse(status="failed", error="audio_uri is required")

    # 文件不存在时提前失败，避免进入耗时的模型加载流程。
    audio_path = Path(audio_uri)
    if not audio_path.is_file():
        return InvokeResponse(status="failed", error="audio file not found")

    try:
        # 延迟导入 faster-whisper，保证节点注册与调度等轻量路径不依赖重型依赖；
        # 导入前先预加载 NVIDIA 动态库，否则 ctranslate2 找不到 libcublas。
        _load_cuda_libraries()
        from faster_whisper import WhisperModel
        # 模型路径默认本地优先：参数 > 环境变量 > 工作区本地目录 > 远端兜底。
        model_path = resolve_model_path(request.params)
        device = str(request.params.get("device") or os.getenv("WHISPER_DEVICE", "auto"))
        # auto 让 faster-whisper 根据硬件自动选择 float16/int8 等计算类型。
        compute_type = str(request.params.get("compute_type") or "float16")
        model = WhisperModel(
            model_path,
            device=device,
            compute_type=compute_type,
        )
        # language 默认日语；vad_filter 默认开启（用户 2026-08 决定）：过滤静音
        # 段以提速并减少无语音处幻觉；长静音时 VAD 压缩时间轴可能轻微错位，
        # 如需极致对齐可在工作流参数中显式关闭。
        # task 默认 transcribe，中文直出模型可传 translate 直接翻译为目标语言。
        # condition_on_previous_text 默认 False：长音频下开启会导致重复/漂移，
        # 关闭后每个 30s 窗口独立解码，是 faster-whisper 官方建议的长音频方案。
        output_dir = Path(request.output_dir)
        output_dir.mkdir(parents=True, exist_ok=True)
        # 分块转写：默认每 1 分钟一块（chunk_seconds=60），切块失败自动回退整段。
        chunk_seconds = int(request.params.get("chunk_seconds", 60))
        chunks = _split_audio(audio_path, output_dir, chunk_seconds, _ffmpeg_bin())
        # 每视频自适应 VAD：若开启 vad_filter 且未显式传 vad_parameters，则根据本音频
        # 信号分析自动确定 VAD 参数（BGM 覆盖/静音比例/长停顿），改善碎片化与漏识别。
        # 门控 WOV_AUTO_VAD=0 可关闭；显式传入 vad_parameters 时跳过。
        vad_parameters = request.params.get("vad_parameters")
        vad_filter = bool(request.params.get("vad_filter", True))
        if vad_filter and not vad_parameters and os.getenv("WOV_AUTO_VAD", "1") == "1":
            try:
                from nodes.vad_profiler import vad_parameters_for_audio

                vad_parameters = vad_parameters_for_audio(
                    audio_path,
                    sample_rate=int(request.params.get("sample_rate", 16000)),
                    whisper_invoke=None,  # 片段网格需模型；生产默认仅信号分析（快速）
                    run_dir=Path(request.output_dir),
                    chunk_seconds=chunk_seconds,
                )
                logger.info("自动 VAD 参数: %s", vad_parameters)
            except Exception as exc:  # noqa: BLE001
                # 分析失败不影响转写：回退默认 bentenVAD，仅记录。
                logger.warning("自动 VAD 分析失败，回退默认: %s", exc)
                vad_parameters = None
        # 逐块转写并合并：offset 用每块实际时长累积（WAV 头精确），SRT 序号连续。
        logger.info("转写开始: %d 个分块", len(chunks))
        lines: list[str] = []
        offset = 0.0
        # SRT 序号从 1 开始，跨块连续递增。
        srt_number = 1
        transcribe_started = time.monotonic()
        for chunk_index, chunk in enumerate(chunks, start=1):
            # 暂停检查：批量暂停时在 run 根目录写 paused.flag，whisper 在分块
            # 之间检查该信号（默认 60s 一块，暂停粒度不超过一块），检测到即抛
            # 异常，由 invoke 转 failed、调度器保持任务 PAUSED，继续时整个节点
            # 重新转写（whisper 没有节点级断点存档，产物只在结束时一次性写出）。
            if (Path(request.output_dir).parent.parent / PAUSE_FLAG).exists():
                raise RuntimeError(f"whisper 被暂停（run {request.run_id}）")
            chunk_started = time.monotonic()
            segments, _info = model.transcribe(
                str(chunk),
                language=str(request.params.get("language", "ja")),
                task=str(request.params.get("task", "transcribe")),
                beam_size=int(request.params.get("beam_size", 1)),
                vad_filter=vad_filter,
                vad_parameters=vad_parameters,
                condition_on_previous_text=bool(
                    request.params.get("condition_on_previous_text", False)
                ),
            )
            # 进度日志：块序号/总数、单块耗时、实时倍率（块音频时长/墙钟耗时）
            # 与转写累计耗时，直观反映数据处理速度。
            chunk_elapsed = time.monotonic() - chunk_started
            srt_number += _append_srt_lines(lines, segments, offset, srt_number)
            # 偏移按本块实际时长推进，避免假设块长导致的累积漂移。
            offset += _wav_duration_seconds(chunk, chunk_seconds)
            logger.info(
                "分块 %d/%d 完成 offset=%.2fs 耗时 %.1fs (%.2fx 实时, 累计 %.1fs)",
                chunk_index, len(chunks), offset, chunk_elapsed,
                chunk_seconds / chunk_elapsed if chunk_elapsed > 0 else 0.0,
                time.monotonic() - transcribe_started,
            )
        output_path = output_dir / "transcript.srt"
        output_path.write_text("\n".join(lines), encoding="utf-8")
        return InvokeResponse(status="completed", outputs={"srt_uri": str(output_path)})
    except Exception as exc:  # noqa: BLE001
        # 模型加载或转写异常统一转换为 failed 响应，不让调度线程崩溃。
        return InvokeResponse(status="failed", error=str(exc))