import logging
import os
from typing import List, Optional, Tuple

from faster_whisper import WhisperModel
from numpy.typing import NDArray

from verbatim.logging_utils import status_enabled
from verbatim_audio.audio import samples_to_seconds

from ...transcript.words import Word
from .transcribe import Transcriber, WhisperConfig

LOG = logging.getLogger(__name__)


class FasterWhisperTranscriber(Transcriber):
    def __init__(
        self,
        *,
        model_size_or_path: str,
        device: str,
        whisper_beam_size: int = 3,
        whisper_best_of: int = 3,
        whisper_patience: float = 1.0,
        whisper_temperatures: Optional[List[float]] = None,
    ):
        if whisper_temperatures is None:
            whisper_temperatures = [0.0, 0.2, 0.4, 0.6, 0.8, 1.0]

        # Configure cache root and offline mode from environment
        cache_root = os.getenv("VERBATIM_MODEL_CACHE")
        download_root = os.path.join(cache_root, "faster-whisper") if cache_root else None
        offline_env = os.getenv("VERBATIM_OFFLINE", "0").lower() in ("1", "true", "yes")

        if device == "cpu":
            try:
                self.whisper_model = WhisperModel(
                    model_size_or_path=model_size_or_path,
                    device=device,
                    compute_type="default",
                    cpu_threads=16,
                    download_root=download_root,
                    local_files_only=offline_env,
                )
            except TypeError:
                # Older faster-whisper without these kwargs
                self.whisper_model = WhisperModel(
                    model_size_or_path=model_size_or_path,
                    device=device,
                    compute_type="default",
                    cpu_threads=16,
                )
        else:
            try:
                self.whisper_model = WhisperModel(
                    model_size_or_path=model_size_or_path,
                    device=device,
                    compute_type="float16",
                    download_root=download_root,
                    local_files_only=offline_env,
                )
            except TypeError:
                self.whisper_model = WhisperModel(
                    model_size_or_path=model_size_or_path,
                    device=device,
                    compute_type="float16",
                )

        self.whisper_beam_size = whisper_beam_size
        self.whisper_best_of = whisper_best_of
        self.whisper_patience = whisper_patience
        self.whisper_temperatures = whisper_temperatures

    def guess_language(self, audio: NDArray, lang: List[str]) -> Tuple[str, float]:
        language, language_probability, all_language_probs = self.whisper_model.detect_language(audio=audio)
        if language in lang:
            LOG.debug("detected '%s' with probability %s", language, language_probability)
            return language, language_probability

        guess_lang = lang[0]
        guess_prob = 0
        for lang_iter in lang:
            for t in all_language_probs:
                if t[0] == lang_iter:
                    if t[1] > guess_prob:
                        guess_lang = t[0]
                        guess_prob = t[1]
                        LOG.debug("detected '%s' with probability %s", lang_iter, guess_prob)
        return guess_lang, guess_prob

    def transcribe(
        self,
        *,
        audio: NDArray,
        lang: str,
        prompt: str,
        prefix: str,
        window_ts: int,
        audio_ts: int,
        whisper_beam_size: int = 3,
        whisper_best_of: int = 3,
        whisper_patience: float = 1.0,
        whisper_temperatures: Optional[List[float]] = None,
    ) -> List[Word]:
        LOG.debug("Transcription Prefix: %s", prefix)
        if whisper_temperatures is None:
            whisper_temperatures = [0.0, 0.2, 0.4, 0.6, 0.8, 1.0]

        show_progress = status_enabled()
        whisper_config: WhisperConfig = WhisperConfig()
        segment_iter, info = self.whisper_model.transcribe(
            audio=audio,
            language=lang,
            task=whisper_config.task,
            log_progress=show_progress,
            beam_size=whisper_beam_size,
            best_of=whisper_best_of,
            patience=whisper_patience,
            length_penalty=whisper_config.length_penalty,
            repetition_penalty=1.0,
            no_repeat_ngram_size=0,
            temperature=whisper_temperatures,
            compression_ratio_threshold=whisper_config.compression_ratio_threshold,
            log_prob_threshold=whisper_config.logprob_threshold,
            no_speech_threshold=whisper_config.no_speech_threshold,
            condition_on_previous_text=False,
            prompt_reset_on_temperature=0.0,
            initial_prompt=prompt,
            prefix=prefix,
            suppress_blank=whisper_config.suppress_blank,
            suppress_tokens=whisper_config.suppress_tokens,
            without_timestamps=False,
            max_initial_timestamp=1.0,
            word_timestamps=True,
            prepend_punctuations=whisper_config.prepend_punctuations,
            append_punctuations=whisper_config.append_punctuations,
            multilingual=False,
            vad_filter=False,
            vad_parameters=None,
            max_new_tokens=None,
            chunk_length=None,
            clip_timestamps="0.0",
            hallucination_silence_threshold=None,
            hotwords=None,
            language_detection_threshold=0.5,
            language_detection_segments=1,
        )

        LOG.debug(f"info={info}")

        transcript_words: List[Word] = []
        for segment in segment_iter:
            # Guard against missing word timestamps in segment
            words = getattr(segment, "words", None)
            if not words:
                continue
            for w in words:
                word = Word.from_word(word=w, lang=lang, ts_offset=window_ts)
                if word.end_ts > audio_ts:
                    continue

                LOG.debug(f"[{word.start_ts} ({samples_to_seconds(word.start_ts)})]: {w.word}")
                transcript_words.append(word)

        return transcript_words
