# from logging import Logger
import time
import torch
import gc
from transformers import WhisperForConditionalGeneration, WhisperProcessor

from transcription.engines.BatchSTT import BatchSTT

class WhisperEngine(BatchSTT):
    TARGET_SAMPLING_RATE = 16000
    
    def loadModel(self) -> None:
        self.processor = WhisperProcessor.from_pretrained(self.modelName)
        self.model = WhisperForConditionalGeneration.from_pretrained(
            self.modelName,
            torch_dtype = self.dType # check twice
        ).to(self.device) # ??? recheck
    
    def unloadModel(self) -> None:
        self.model = None
        self.processor = None

        # TODO: MPS?
        if self.device == "cuda": 
            torch.cuda.empty_cache()
    
    def transcribeBatch(
        self,
        batch,
    ) -> str:
        assert self.processor is not None
        assert self.model is not None

        inputs = self.processor(
            batch, 
            sampling_rate=self.TARGET_SAMPLING_RATE,
            return_tensors="pt",
            padding=True,
        )
        
        input_features = inputs.input_features.to(self.device).to(self.dType)

        with torch.no_grad():
            predicted_ids = self.model.generate(
                input_features,
                language=self.language,
                task="transcribe",
                temperature=0.0,
            )
        
        batchText = self.processor.batch_decode(
            predicted_ids, 
            skip_special_tokens=True,
        )

        inputs = None
        input_features = None
        predicted_ids = None
        gc.collect()

        # maybe do here something with MPS?
        if self.device.startswith("cuda"):
            torch.cuda.empty_cache()
            torch.cuda.ipc_collect()

        return " ".join(batchText)