"""transcribe: faster-whisper with word timestamps."""

from __future__ import annotations

import os
import threading
from typing import Any

from ..errors import BadRequest, Cancelled
from ..schemas import TranscribeRequest
from .common import Ctx, clamp01, open_input
from .devices import ctranslate2_device


class WhisperTranscriber:
    def __init__(self, model_dir: str, settings: Any) -> None:
        from faster_whisper import WhisperModel

        device = ctranslate2_device(settings)
        self.model = WhisperModel(
            model_dir,
            device=device,
            compute_type="float16" if device == "cuda" else "int8",
        )

    def transcribe(
        self, path: str, language: str | None, cancel: threading.Event
    ) -> tuple[str | None, list[dict[str, Any]]]:
        segments, info = self.model.transcribe(path, language=language, word_timestamps=True, vad_filter=True)
        words: list[dict[str, Any]] = []
        for segment in segments:
            if cancel.is_set():
                raise Cancelled()
            for word in segment.words or []:
                start = int(round(word.start * 1000))
                words.append(
                    {
                        "startMs": start,
                        "endMs": max(int(round(word.end * 1000)), start),
                        "text": word.word.strip(),
                        "confidence": clamp01(word.probability),
                    }
                )
        return info.language, words


def load_transcriber(files: dict[str, str], settings: Any) -> WhisperTranscriber:
    return WhisperTranscriber(os.path.dirname(files["model.bin"]), settings)


def run(ctx: Ctx, body: TranscribeRequest, cancel: threading.Event) -> dict[str, Any]:
    real, info = open_input(ctx, body.path, cancel, need_video=False)
    if not info.has_audio:
        raise BadRequest("input has no audio stream")
    transcriber = ctx.registry.runtime("transcribe")
    language, words = transcriber.transcribe(real, body.language, cancel)
    return {"language": language, "words": [w for w in words if w["text"]]}