"""Contains the SpeechToKeys class."""

import logging
from queue import Empty, Queue
from threading import Thread
from time import sleep
from typing import Optional, Tuple

import numpy as np
import torch
import whisper
from pynput.keyboard import Key
from pynput.keyboard import Listener as KeyboardListener
from pynput.mouse import Listener as MouseListener

from .audio_capture import AudioCapture
from .essential_thread import essential_thread
from .rate_limited_keyboard import Controller as KeyboardController

CUDA_SUCKS_MESSAGE = (
    "NVidia driver errors occur after waking from suspend. You can try restarting the"
    " NVidia driver with the following commands:\n\n"
    "sudo rmmod nvidia_uvm\nsudo modprobe nvidia_uvm\n\n"
)


# pylint: disable=too-many-instance-attributes
class SpeechToKeys:
    """
    Class to convert speech to keyboard input.
    """

    # pylint: disable=too-many-arguments,too-many-positional-arguments
    def __init__(
        self,
        model_name: str = "turbo",
        device: Optional[str | int] = None,
        max_key_rate: float = 500.0,
        ambient_duration: float = 1.0,
        activity_ambient_multiplier: float = 1.5,
        phrase_timeout: float = 3.0,
    ):
        self._model_name = model_name
        self._phrase_timeout = phrase_timeout
        self._data_queue = Queue[Tuple[float, float, np.ndarray]]()
        self._float_samples: np.ndarray = np.array([], dtype=np.float32)
        self._last_block_end_time = None
        self._buffer = ""
        self._is_first_phrase = True
        self._dictation_active = False
        self._keyboard = KeyboardController(1.0 / max_key_rate)
        self._keyboard_listener = None
        self._mouse_listener = None
        self._is_programmatic_typing = False
        self._exit = False

        self._audio_capture = AudioCapture(
            on_audio=self._on_audio,
            device=device,
            ambient_duration=ambient_duration,
            activity_ambient_multiplier=activity_ambient_multiplier,
        )

        self._initialize_model()

        audio_thread = Thread(
            target=self._loop, daemon=True, name="TranscriptionThread"
        )
        audio_thread.start()
        logging.info("Processing thread started.")

    @staticmethod
    def query_devices():
        """
        Lists the available input audio devices using sounddevice.
        """
        print(AudioCapture.query_devices())

    def shutdown(self):
        """
        Shuts down the speech to keys.
        """
        self.enabled = False
        self._audio_capture.shutdown()
        self._exit = True

        while not self._data_queue.empty():
            try:
                self._data_queue.get_nowait()
            except Empty:
                break
        logging.info("SpeechToKeys shutdown complete.")

    def _initialize_model(self):
        logging.info("Loading Whisper model: %s", self._model_name)
        try:
            self._audio_model = whisper.load_model(self._model_name)
        except Exception as e:
            raise RuntimeError(
                "Error loading Whisper model. " + CUDA_SUCKS_MESSAGE
            ) from e

        logging.info("Pytorch device type: %s", self._audio_model.device.type)
        if self._audio_model.device.type == "cpu":
            logging.warning(
                "No CUDA accelerators found. Using CPU. This will be slow. %s",
                CUDA_SUCKS_MESSAGE,
            )
        logging.info("Whisper model loaded successfully.")

    def _reset(self):
        self._float_samples = np.array([], dtype=np.float32)
        self._last_block_end_time = None
        self._buffer = ""
        self._is_first_phrase = True
        while not self._data_queue.empty():
            try:
                self._data_queue.get_nowait()
            except Empty:
                break

    def _on_audio(
        self, block_start_time: float, block_end_time: float, audio: np.ndarray
    ) -> None:
        """
        Threaded callback function to receive audio data when recordings finish.
        audio: An AudioData containing the recorded bytes.
        """
        if self._dictation_active:
            logging.debug("Recording callback received and dictation is active.")
            self._data_queue.put((block_start_time, block_end_time, audio))
        else:
            logging.warning("Recording callback received but dictation is not active.")

    def _on_press(self, key):
        """Callback for keyboard listener."""
        if self._is_programmatic_typing:
            return  # Ignore key presses generated by the program itself

        # We don't need to know which key, just that a press occurred.
        logging.debug("User key press detected. %s", key)
        self._reset()

    def _on_click(self, x, y, button, _pressed):
        """Callback for mouse listener."""
        # We don't need to know which button, just that a click occurred.
        logging.debug("User mouse click detected. x=%s, y=%s, button=%s", x, y, button)
        self._reset()

    def _start_keyboard_listener(self):
        if self._keyboard_listener is None:
            logging.info("Starting keyboard listener.")
            self._keyboard_listener = KeyboardListener(on_press=self._on_press)
            self._keyboard_listener.start()
            logging.info("Keyboard listener started.")
        else:
            logging.info("Keyboard listener already running.")

    def _stop_keyboard_listener(self):
        if self._keyboard_listener is not None:
            logging.info("Stopping keyboard listener.")
            self._keyboard_listener.stop()
            self._keyboard_listener = None
            logging.info("Keyboard listener stopped.")
        else:
            logging.info("Keyboard listener not running or already stopped.")

    def _start_mouse_listener(self):
        if self._mouse_listener is None:
            logging.info("Starting mouse listener.")
            self._mouse_listener = MouseListener(on_click=self._on_click)
            self._mouse_listener.start()
            logging.info("Mouse listener started.")
        else:
            logging.info("Mouse listener already running.")

    def _stop_mouse_listener(self):
        if self._mouse_listener is not None:
            logging.info("Stopping mouse listener.")
            self._mouse_listener.stop()
            self._mouse_listener = None
            logging.info("Mouse listener stopped.")
        else:
            logging.info("Mouse listener not running or already stopped.")

    def _process_blocks(self, blocks: list[np.ndarray], phrase_break: bool):
        if blocks:
            new_samples = np.concatenate(blocks).astype(np.float32) / 32768.0
            new_samples = new_samples.flatten()
            self._float_samples = np.concatenate((self._float_samples, new_samples))
            self._transcribe(phrase_break)
            if phrase_break:
                self._float_samples = np.array([], dtype=np.float32)

    def _loop(self):
        """Processes audio from the queue and performs transcription."""
        with essential_thread():
            while not self._exit:
                if not self._dictation_active:
                    sleep(0.1)
                    continue

                new_blocks = []

                while not self._data_queue.empty():
                    try:
                        block_start_time, block_end_time, audio = (
                            self._data_queue.get_nowait()
                        )
                    except Empty:
                        break

                    if (
                        self._last_block_end_time is not None
                        and block_start_time
                        > self._last_block_end_time + self._phrase_timeout
                    ):
                        self._process_blocks(new_blocks, True)

                    self._last_block_end_time = block_end_time
                    new_blocks.append(audio)

                self._process_blocks(new_blocks, False)
                sleep(0.1)

    def _transcribe(self, phrase_break: bool):
        try:
            self._is_programmatic_typing = True
            if phrase_break:
                # Sometimes the transciption misses ending punctuation if it had
                # thought more words would come, but did not.
                if self._buffer and self._buffer.rstrip()[-1] not in [".", "!", "?"]:
                    self._keyboard.type(".")
                self._buffer = ""
                self._is_first_phrase = False

            if self._float_samples.size == 0:
                return

            logging.debug("Transcribing audio.")
            try:
                result = self._audio_model.transcribe(
                    self._float_samples, fp16=torch.cuda.is_available()
                )
            except RuntimeError:
                logging.error("Error transcribing audio. Reloading model.")
                self._initialize_model()
                result = self._audio_model.transcribe(
                    self._float_samples, fp16=torch.cuda.is_available()
                )
            text = result["text"]
            if self._is_first_phrase:
                # There was no previous phrase, so remove any leading whitespace that
                # the transcription might have added.
                text = text.lstrip()
            logging.info("Transcribed text: '%s'", text)

            if self._buffer:
                # find the first index where the text and buffer differ
                index = next(
                    (i for i, (t, b) in enumerate(zip(text, self._buffer)) if t != b),
                    len(self._buffer),
                )

                for _ in range(index, len(self._buffer)):
                    self._keyboard.tap(Key.backspace)
            else:
                index = 0

            for i in range(index, len(text)):
                self._keyboard.type(text[i])
            self._buffer = text
        finally:
            self._is_programmatic_typing = False

    @property
    def enabled(self) -> bool:
        """
        Returns the enabled state of the speech to keys.
        """
        return self._dictation_active

    @enabled.setter
    def enabled(self, value: bool):
        if value == self._dictation_active:
            logging.info(
                "SpeechToKeys.enabled changing from %s to %s (no change)",
                self._dictation_active,
                value,
            )
            return

        logging.info(
            "SpeechToKeys.enabled changing from %s to %s",
            self._dictation_active,
            value,
        )

        if value:
            self._start_keyboard_listener()
            self._start_mouse_listener()

            self._reset()
            self._dictation_active = True

            try:
                self._audio_capture.start()
                logging.info("Background audio listener started.")
            except (OSError, AttributeError, RuntimeError) as e:
                logging.error(
                    "Error starting background audio listener: %s", e, exc_info=True
                )
                self._stop_keyboard_listener()
                self._stop_mouse_listener()
                return

        else:
            self._dictation_active = False  # Set this first
            self._stop_keyboard_listener()  # Stop keyboard listener
            self._stop_mouse_listener()  # Stop mouse listener
            self._audio_capture.stop()

            self._reset()
