#!/usr/local/bin/python3.12

from vosk import KaldiRecognizer, Model
import os.path
import json
import os


FRAME = 4_000

SAMPLE_RATE = 16_000

MODELS = list(
  map(
    lambda model: os.path.join(os.getcwd(), 'third-party/vosk-models', model),
    [
      'vosk-model-en-us-0.22-lgraph', # English (United States)
      'vosk-model-small-ru-0.22', # Russian
      'vosk-model-small-uk-v3-small', # Ukrainian
    ]
  )
)

DATASET = os.path.join(os.getcwd(), 'files/dataset.json')

AUDIOS = os.path.join(os.getcwd(), 'files/audios')


def score_transcription(result: dict) -> float:
  words = result.get("result", [])
  values = list(
    map(lambda word: word.get("conf", 0.0), words)
  )
  try:
    return sum(values) / len(values)
  except ZeroDivisionError:
    return 0.0


def main() -> None:
  recognizers = []
  transcriptions = []

  for model in MODELS:
    model = Model(model)
    recognizer = KaldiRecognizer(model, SAMPLE_RATE)
    recognizer.SetWords(True)
    recognizers.append(recognizer)

  for child in os.listdir(AUDIOS):
    if child == ".gitkeep": continue

    file = os.path.join(AUDIOS, child)
    if not os.path.isfile(file): continue

    with open(file, 'rb') as audio:
      while True:
        frame = audio.read(FRAME)
        if len(frame) == 0:
          break
        for recognizer in recognizers:
          recognizer.AcceptWaveform(frame)

      results = []
      for i, recognizer in enumerate(recognizers):
        result = json.loads(recognizer.FinalResult())
        if result.get("text", "") == "": continue

        score = score_transcription(result)
        if score == 0.0: continue

        results.append({'model': MODELS[i], 'score': score, 'result': result})

      if len(results) == 0:
        continue
      results = list(
        sorted(results, key=lambda result: result['score'], reverse=True)
      )

      transcription = results[0]
      transcription.update({'video': child})

      transcriptions.append(transcription)

  dataset = []
  for transcription in transcriptions:
    words = transcription['result']['result']
    message = []
    for i, word in enumerate(words):
      message.append(f'[{i + 1}]{word['word']}')
    dataset.append({'user': ' '.join(message), 'assistant': None})

  with open(DATASET, 'w') as file:
    json.dump(dataset, file, ensure_ascii=False, indent=2)


if __name__ == "__main__":
  main()