#!/usr/bin/env python3
import os
import torch
import rospkg
from pathlib import Path

WHISPER_CACHE = os.path.join(str(Path.home()), '.cache', 'whisper')
os.makedirs(WHISPER_CACHE, exist_ok=True)
os.environ["TIKTOKEN_CACHE_DIR"] = WHISPER_CACHE

import sys

if len(sys.argv) < 3:
    print('Usage:')
    print('rosrun lasr_speech_recognition transcribe_microphone by-index <device_index>')
    print('rosrun lasr_speech_recognition transcribe_microphone by-name <substring>')
    exit(1)
else:
    matcher = sys.argv[1]
    device_index = None
    if matcher == 'by-index':
        device_index = int(sys.argv[2])
    elif matcher == 'by-name':
        import speech_recognition as sr
        microphones = enumerate(sr.Microphone.list_microphone_names())

        target_name = sys.argv[2]
        for index, name in microphones:
            if target_name in name:
                device_index = index
                break
        
        if device_index is None:
            print('Could not find device!')
            exit(1)
    else:
        print('Invalid matcher')
        exit(1)

import rospy
from std_srvs.srv import Empty, EmptyResponse
rospy.init_node('transcribe_mic', anonymous=True)

from lasr_speech_recognition_whisper import SpeechRecognitionToTopic, MicrophonePhraseCollector, load_model

collector = MicrophonePhraseCollector(device_index=device_index)
collector.adjust_for_noise()

#model = load_model("base.en")

model = load_model("medium.en")

# try to run inference on the example file
r = rospkg.RosPack()
EXAMPLE_FILE = r.get_path('lasr_speech_recognition_whisper') + "/test.m4a"
rospy.loginfo("Running transcription on example file to ensure model is loaded...")
rospy.loginfo(model.transcribe(EXAMPLE_FILE, fp16=torch.cuda.is_available()))

worker = SpeechRecognitionToTopic(collector, model, "transcription", infer_partial = False)

def adjust_for_noise(_):
    collector.adjust_for_noise()
    return EmptyResponse()

def start_listening(_):
    worker.start()
    return EmptyResponse()

def stop_listening(_):
    worker.stop()
    return EmptyResponse()

rospy.Service('/whisper/adjust_for_noise', Empty, adjust_for_noise)
rospy.Service('/whisper/start_listening', Empty, start_listening)
rospy.Service('/whisper/stop_listening', Empty, stop_listening)

rospy.loginfo("Starting the Whisper worker!")
rospy.spin()
