from . import command, audio_fn
from vosk import Model, KaldiRecognizer, SetLogLevel
import json
import sys
import logging
import io
from sh import ffmpeg as ffmpeg_unbaked

ffmpeg = ffmpeg_unbaked.bake(_tty_out=False)

model = Model('speechrec/model-nl-smaller')
sample_rate=48000

logger = logging.getLogger('speechrec')

def voice_or_audio_fn(update):
    if not update.message:
        raise ValueError('no message?')

    audio_message = update.message.reply_to_message or update.message
    if (not audio_message or (not audio_message.voice and not audio_message.audio)):
        raise ValueError('no audio found in message or reply-message.')

    if (audio_message.voice):
        return audio_message.voice
    if (audio_message.audio):
        return audio_message.audio
    raise ValueError('no audio found in message or reply-message.')

@command('stt', additional_match_fn=voice_or_audio_fn)
def stt(bea, bot, update, audio):
    rec = KaldiRecognizer(model, sample_rate)
    fl = audio.get_file()
    fl.download(custom_path='speechrec/input.audio')
    buf = io.BytesIO()
    ffmpeg('-loglevel', 'quiet', '-i', 'speechrec/input.audio', '-ar', f'{sample_rate}', '-ac', '1', '-f', 's16le', '-', _out=buf)
    buf.seek(0)
    if (rec.AcceptWaveform(buf.getvalue())):
        res_interm = rec.Result()
    res = rec.FinalResult()
    data = json.loads(res)
    data_interm = json.loads(res_interm)
    result = data.get('text') or data_interm.get('text')
    if result:
        bea.reply(result)
    else:
        bea.reply('Kan het resultaat van de spraakherkenner niet parsen, zie logs.')
        logger.warning(f'Unable to parse: {data!r} {data_interm!r}')

def init(bea, _config):
    return [stt]
