import os
from threading import Lock

import torch
import whisper

# TODO use pydantic config
model_name = os.getenv("ASR_MODEL", "base")
if torch.cuda.is_available():
    model = whisper.load_model(model_name).cuda()
else:
    model = whisper.load_model(model_name)
model_lock = Lock()

# TODO move transcribe to a modeling worker
def transcribe(audio):
    # options_dict = {"task" : task}
    # if language:
    #     options_dict["language"] = language
    # if initial_prompt:
    #     options_dict["initial_prompt"] = initial_prompt
    with model_lock:
        # result = model.transcribe(audio, **options_dict)
        result = model.transcribe(audio)

    return result
