import base64, tempfile, os, torch
from transformers import AutoModelForSpeechSeq2Seq, AutoProcessor, pipeline
from functions.utils import getAudioDuration

MODEL_NAME = "openai/whisper-large-v3"
device = "cuda" if torch.cuda.is_available() else "cpu"

model = AutoModelForSpeechSeq2Seq.from_pretrained(MODEL_NAME, torch_dtype=torch.float16).to(device)
processor = AutoProcessor.from_pretrained(MODEL_NAME)
pipe = pipeline("automatic-speech-recognition", model=model, tokenizer=processor.tokenizer,
    feature_extractor=processor.feature_extractor, torch_dtype=torch.float16, device=device)


def arSTT(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, generate_kwargs={"language": "arabic"})
        text = result["text"]
        duration = getAudioDuration(tempAudioPath)
    finally:
        os.remove(tempAudioPath)

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