import asyncio
import json
import logging
from typing import Optional

from vosk import Model, KaldiRecognizer, SetLogLevel

from py_speech_service import speech_service_pb2
from py_speech_service.grammar_parser import GrammarParser

import queue
import sounddevice as sd

class SpeechRecognition:
    model: Model
    grammar_parser: GrammarParser
    required_confidence: float = 80
    stop_speech_recognition_event: Optional[asyncio.Event] = None
    continue_speech_recognition: bool = True
    grpc_response_queue: Optional[asyncio.Queue] = None
    shutdown_event = asyncio.Event()
    recognition_queue = asyncio.Queue()

    def __init__(self):
        SetLogLevel(-1)
        self.stop_speech_recognition_event = None

    def set_vosk_details(self, vosk_model: str, grammar_file: str, required_confidence: float = 80):
        self.grammar_parser = GrammarParser()
        self.grammar_parser.set_grammar_file(grammar_file)

        if vosk_model.count("/") > 0 or vosk_model.count("\\") > 0:
            self.model = Model(vosk_model)
        else:
            self.model = Model(None, vosk_model)

        self.required_confidence = required_confidence

    async def start_speech_recognition(self, context):
        asyncio.create_task(self.process_recognition_queue())
        await asyncio.to_thread(self.listen)

    async def process_recognition_queue(self):
        logging.info("Started processing speech recognition queue")
        print("Started processing speech recognition queue")
        while not self.shutdown_event.is_set():  # Keep running unless shutdown is triggered
            try:
                item = await asyncio.wait_for(self.recognition_queue.get(), timeout=1.0)  # Wait for 1 second
                await self.process_speech(item)
            except asyncio.TimeoutError:
                continue  # Retry checking
        self.recognition_queue.task_done()
        logging.info("Stopped processing speech recognition queue")
        print("Stopped processing speech recognition queue")

    def listen(self):
        if self.stop_speech_recognition_event:
            self.stop_speech_recognition_event.set()

        stop_speech_recognition_event = asyncio.Event()
        self.stop_speech_recognition_event = stop_speech_recognition_event

        device_info = sd.query_devices(sd.default.device[0], 'input')
        samplerate = int(device_info['default_samplerate'])
        speech_queue = queue.Queue()

        def record_callback(indata, frames, time, status):
            if status:
                logging.error(str(status))
            speech_queue.put(bytes(indata))

        words_json = json.dumps(self.grammar_parser.all_words)
        recognizer = KaldiRecognizer(self.model, samplerate, words_json)
        recognizer.SetWords(False)
        recognizer.SetGrammar(words_json)

        try:
            logging.info("Started listening to voice via VOSK")
            print("Started listening to voice via VOSK")
            with sd.RawInputStream(dtype='int16',
                                   channels=1,
                                   callback=record_callback):
                while not stop_speech_recognition_event.is_set() and not self.shutdown_event.is_set():
                    data = speech_queue.get()
                    if recognizer.AcceptWaveform(data):
                        recognized_text = recognizer.Result()
                        if recognized_text:
                            asyncio.run(self.recognition_queue.put(recognized_text))
            logging.info("Stopped listening to voice via VOSK")
            print("Stopped listening to voice via VOSK")
        except KeyboardInterrupt:
            print('Finished recording due to keyboard interrupt')
            logging.error("Finished recording due to keyboard interrupt")
        except Exception as e:
            print("Error from VOSK: " + str(e))
            logging.error("Error from VOSK")
            logging.error(e)

    async def process_speech(self, recognizer_result: str):
        try:
            result_dict = json.loads(recognizer_result)
            text_recognized = result_dict.get("text", "")
            if not text_recognized == "":
                match = self.grammar_parser.find_match(text_recognized, self.required_confidence)

                if match is not None:
                    if self.grpc_response_queue:
                        response = speech_service_pb2.SpeechServiceResponse()
                        response.speech_recognized.heard_text = text_recognized
                        response.speech_recognized.recognized_text = match.matched_text
                        response.speech_recognized.recognized_rule = match.rule
                        response.speech_recognized.confidence = round(match.confidence, 2)
                        response.speech_recognized.semantics.update(match.values)
                        await self.grpc_response_queue.put(response)
                    else:
                        print("I heard: '" + text_recognized + "', but I am " + str(round(match.confidence, 2)) + "% sure you said '" + match.matched_text + "'")

        except Exception as e:
            print("Error processing speech: " + str(e))
            logging.error("Error processing speech")
            logging.error(e)

    def set_grpc_response_queue(self, queue: asyncio.Queue):
        self.grpc_response_queue = queue

    def stop_speech_recognition(self):
        if self.stop_speech_recognition_event is not None:
            self.stop_speech_recognition_event.set()

    def shutdown(self):
        self.shutdown_event.set()
