from logging import getLogger

import numpy as np
from funasr import AutoModel

import config

logger = getLogger(__name__)


class FunASR:
    def __init__(self, source_lange: str = 'en', warmup=True) -> None:
        self.source_lange = source_lange

        self.model = AutoModel(
            model="paraformer-zh", vad_model="fsmn-vad", punc_model="ct-punc"
        )
        if warmup:
            self.warmup()

    def warmup(self, warmup_steps=1):
        warmup_soundfile = f"{config.ASSERT_DIR}/jfk.flac"
        for _ in range(warmup_steps):
            self.model.generate(input=warmup_soundfile, disable_pbar=True, hotword=config.hotwords_file.as_posix())

    def transcribe(self, audio_buffer: bytes, language):
        audio_frames = np.frombuffer(audio_buffer, dtype=np.float32)
        # sf.write(f'{config.ASSERT_DIR}/{time.time()}.wav', audio_frames, samplerate=16000)
        try:
            output = self.model.generate(input=audio_frames, disable_pbar=True, hotword=config.hotwords_file.as_posix())
            return output
        except Exception as e:
            logger.error(e)
            return []