import base64, tempfile, os
from transformers import pipeline
from functions.utils import CUDA_AVAILABLE, TORCH_DTYPE, getAudioDuration

MODEL_NAME = "facebook/mms-1b-all"
pipe = pipeline(
    "automatic-speech-recognition",
    model=MODEL_NAME,
    model_kwargs={"target_lang": "mos", "torch_dtype": TORCH_DTYPE},
    device=0 if CUDA_AVAILABLE else -1,
)

# MODEL_NAME = "burkimbia/BIA-WHISPER-LARGE-SACHI_V3"
# pipe = pipeline("automatic-speech-recognition", model=MODEL_NAME)


def mooreSTT(audioBase64: str) -> dict:
    audioBytes = base64.b64decode(audioBase64)

    with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tempFile:
        tempFile.write(audioBytes)
        tempAudioPath = tempFile.name

    try:
        result = pipe(tempAudioPath)
        text = result["text"]
        duration = getAudioDuration(tempAudioPath)
    finally:
        os.remove(tempAudioPath)

    return {'text': text, 'language': 'mos', 'duration': duration}