import os

from speechbrain.pretrained import EncoderDecoderASR
from sziszapangma.integration.service_core.asr.asr_base_processor import AsrBaseProcessor
from sziszapangma.integration.service_core.asr.asr_result import AsrResult


class SpeechbrainAsrProcessor(AsrBaseProcessor):
    asr_model: EncoderDecoderASR

    def __init__(self):
        super().__init__()
        self.asr_model = EncoderDecoderASR.from_hparams(
            source="speechbrain/asr-transformer-transformerlm-librispeech"
        )

    def process_asr(self, audio_file_path: str) -> AsrResult:
        transcription = self.asr_model.transcribe_file(audio_file_path)
        os.remove(audio_file_path)
        words = [it.lower() for it in transcription.split(' ')]
        final_transcription = transcription.lower()
        return AsrResult(words=words, full_text=final_transcription, words_time_alignment=None)


if __name__ == '__main__':
    SpeechbrainAsrProcessor().start_processor()