"""faster-whisper ASR 节点入口。

使用 faster-whisper 加载 Whisper 模型，将音频转写为带时间轴的 SRT 文件。
模型、设备与计算类型均可通过参数或环境变量配置。
"""

from __future__ import annotations

import json
import os
from pathlib import Path

from wov_sdk.models import InvokeRequest, InvokeResponse, NodeManifest
from wov_sdk.server import run_node


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 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，保证健康检查等轻量路径不依赖重型依赖。
        from faster_whisper import WhisperModel

        # 参数优先于环境变量；模型路径缺省使用 faster-whisper 的 large-v3。
        model_path = str(
            request.params.get("model_path")
            or os.getenv("WHISPER_MODEL_PATH", "large-v3")
        )
        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 "auto")
        model = WhisperModel(
            model_path,
            device=device,
            compute_type=compute_type,
        )
        # language 默认日语，vad_filter 过滤静音段以提升转写质量。
        segments, _info = model.transcribe(
            str(audio_path),
            language=str(request.params.get("language", "ja")),
            beam_size=int(request.params.get("beam_size", 1)),
            vad_filter=bool(request.params.get("vad_filter", True)),
        )

        output_dir = Path(request.output_dir)
        output_dir.mkdir(parents=True, exist_ok=True)
        output_path = output_dir / "transcript.srt"
        # 按 SRT 标准输出：序号、时间轴、文本和空行交替。
        lines: list[str] = []
        for index, segment in enumerate(segments, start=1):
            lines.extend(
                [
                    str(index),
                    f"{format_timestamp(segment.start)} --> {format_timestamp(segment.end)}",
                    segment.text.strip(),
                    "",
                ]
            )
        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))


def main() -> None:
    """加载节点清单并以本模块的 invoke 处理器启动服务。"""
    manifest_path = Path(__file__).resolve().parent.parent / "node.manifest.json"
    with open(manifest_path, "r", encoding="utf-8") as f:
        manifest = NodeManifest.from_dict(json.load(f))
    run_node(manifest, invoke)


if __name__ == "__main__":
    main()