import gc
import logging
import os

import torch
import whisperx


class Transcriber:
    """
    A class to transcribe audio files using whisperx.

    Attributes:
    ----------
        language_code : str (default: "en")
            Language code for whisperx transcription. Default is english.
        device : str (default: "cuda")
            Device to use for transcription. Default is "cuda" if available, else "cpu".
        compute_type : str (default: "float16")
            Precision type for whisperx transcription.

    Methods
    -------
        transcribe(wav_filename: str, batch_size: int = 16) -> dict:
            Transcribe the wav file.
    """

    def __init__(
        self,
        language_code: str = "en",
        device: str = "cuda",
        compute_type: str = "float16",
        asr_options=None,
    ):
        if asr_options is None:
            asr_options = {
                "repetition_penalty": 1,
                "prompt_reset_on_temperature": 0.5,
                "no_repeat_ngram_size": 0,
            }
        self.logger = logging.getLogger(__name__)

        self.device = device if torch.cuda.is_available() else "cpu"
        self.logger.info("Using device: %s", self.device)

        try:
            self.transcription_model = whisperx.load_model(
                "large-v2",
                device=self.device,
                compute_type=compute_type,
                language=language_code,
                asr_options=asr_options,
            )
            self.alignment_model, self.metadata = whisperx.load_align_model(
                language_code=language_code, device=self.device
            )
        except Exception as e:
            self.logger.exception("Failed to initialize models: %s", e)
            raise
        self.transcript_output = None
        self.hf_token = os.environ.get("HF_TOKEN")

    def transcribe(self, wav_filename: str, batch_size: int = 16) -> dict:
        """
        Transcribe the wav file.

        Parameters:
        ----------
        wav_filename : str
            Path to the wav file.
        batch_size : int (default: 16)
            Batch size for whisperx transcription.

        Returns:
        -------
        dict
            Transcription result.
        """
        try:
            audio = whisperx.load_audio(wav_filename)

            result = self.transcription_model.transcribe(audio, batch_size=batch_size)

            result = whisperx.align(
                result["segments"],
                self.alignment_model,
                self.metadata,
                audio,
                self.device,
                return_char_alignments=False,
            )

            self._clear_memory()

            diarize_model = whisperx.DiarizationPipeline(
                use_auth_token=self.hf_token, device=self.device
            )
            diarize_segments = diarize_model(audio)
            result = whisperx.assign_word_speakers(diarize_segments, result)

            self._clear_memory()
            self.transcript_output = result
            return result
        except Exception as e:
            self.logger.exception("Failed to transcribe: %s", e)
            raise

    def _clear_memory(self):
        """
        Clear GPU memory.
        """
        try:
            gc.collect()
            torch.cuda.empty_cache()
        except Exception as e:
            self.logger.exception("Failed to clear GPU memory: %s", e)
            raise