from abc import abstractmethod
import asyncio
import copy
import functools
import itertools
import logging
import queue
import threading
import time
from typing import AsyncGenerator
import numpy as np
import pyaudio
import torch
import torchaudio
from enum import Enum
import datetime

from typing import AsyncGenerator, AsyncIterable, BinaryIO, Iterable
from dailyai.queue_aggregators import LLMAssistantContextAggregator, LLMUserContextAggregator

from dailyai.queue_frame import (
    AudioQueueFrame,
    EndStreamQueueFrame,
    ImageQueueFrame,
    QueueFrame,
    SpriteQueueFrame,
    StartStreamQueueFrame,
    TranscriptionQueueFrame,
    TTSCompletedFrame,
    UserStartedSpeakingFrame,
    UserStoppedSpeakingFrame
)

torch.set_num_threads(1)

model, utils = torch.hub.load(repo_or_dir='snakers4/silero-vad',
                              model='silero_vad',
                              force_reload=False)

(get_speech_timestamps,
 save_audio,
 read_audio,
 VADIterator,
 collect_chunks) = utils

# Taken from utils_vad.py


def validate(model,
             inputs: torch.Tensor):
    with torch.no_grad():
        outs = model(inputs)
    return outs

# Provided by Alexander Veysov


def int2float(sound):
    abs_max = np.abs(sound).max()
    sound = sound.astype('float32')
    if abs_max > 0:
        sound *= 1/32768
    sound = sound.squeeze()  # depends on the use case
    return sound


FORMAT = pyaudio.paInt16
CHANNELS = 1
SAMPLE_RATE = 16000
CHUNK = int(SAMPLE_RATE / 10)

audio = pyaudio.PyAudio()


class VADState(Enum):
    QUIET = 1
    STARTING = 2
    SPEAKING = 3
    STOPPING = 4


class BaseTransportService():

    def __init__(
        self,
        **kwargs,
    ) -> None:
        self._mic_enabled = kwargs.get("mic_enabled") or False
        self._mic_sample_rate = kwargs.get("mic_sample_rate") or 16000
        self._camera_enabled = kwargs.get("camera_enabled") or False
        self._camera_width = kwargs.get("camera_width") or 1024
        self._camera_height = kwargs.get("camera_height") or 768
        self._speaker_enabled = kwargs.get("speaker_enabled") or False
        self._speaker_sample_rate = kwargs.get("speaker_sample_rate") or 16000
        self._fps = kwargs.get("fps") or 8
        self._vad_start_s = kwargs.get("vad_start_s") or 0.2
        self._vad_stop_s = kwargs.get("vad_stop_s") or 1.2
        self._context = kwargs.get("context") or []

        self._vad_samples = 1536
        vad_frame_s = self._vad_samples / SAMPLE_RATE
        self._vad_start_frames = round(self._vad_start_s / vad_frame_s)
        self._vad_stop_frames = round(self._vad_stop_s / vad_frame_s)
        self._vad_starting_count = 0
        self._vad_stopping_count = 0
        self._vad_state = VADState.QUIET

        duration_minutes = kwargs.get("duration_minutes") or 10
        self._expiration = time.time() + duration_minutes * 60

        self.send_queue = asyncio.Queue()
        self.receive_queue = asyncio.Queue()

        self._threadsafe_send_queue = queue.Queue()

        self._images = None
        self._user_is_speaking = False
        self._current_phrase = ""

        try:
            self._loop: asyncio.AbstractEventLoop | None = asyncio.get_running_loop()
        except RuntimeError:
            self._loop = None

        self._stop_threads = threading.Event()
        self._is_interrupted = threading.Event()

        self._logger: logging.Logger = logging.getLogger()

    def update_messages(self, new_context: list[dict[str, str]], task: asyncio.Task | None):
        if task:
            if not task.cancelled():
                self._current_phrase = ""
                self._context = new_context

    def append_to_context(self, role, text):
        last_context_item = self._context[-1]
        if last_context_item and last_context_item['role'] == role:
            last_context_item['content'] += f" {text}"
        else:
            self._context.append({"role": role, "content": text})

    async def run_pipeline(self, frame, context):
        print(f"starting to speak_after_delay, {frame}")
        print(f"past asyncio sleep, context is {context}")
        # TODO-CB: This exception for missing class gets eaten!
        tma_in = LLMUserContextAggregator(
            context, self._my_participant_id, complete_sentences=False
        )
        tma_out = LLMAssistantContextAggregator(
            context, self._my_participant_id
        )
        print(f"about to call the runner, tma_in is {tma_in}")
        await self._runner(frame, tma_in, tma_out)

    async def run_conversation(self, runner: Iterable[QueueFrame]
                               | AsyncIterable[QueueFrame]
                               | asyncio.Queue[QueueFrame],
                               ) -> AsyncGenerator[QueueFrame, None]:
        current_response_task = None
        self._runner = runner

        async for frame in self.get_receive_frames():
            print(f"got frame of type: {type(frame)}")
            if isinstance(frame, EndStreamQueueFrame):
                break
            # elif not isinstance(frame, TranscriptionQueueFrame):
                # continue

            if hasattr(frame, 'participantId') and frame.participantId == self._my_participant_id:
                continue

            if current_response_task:
                # TODO-CB: Maybe not always interrupt? Are there frame types we can pass through?
                current_response_task.cancel()
                self.interrupt()

            # self._current_phrase += " " + frame.text
            current_llm_context = copy.deepcopy(self._context)
            current_response_task = asyncio.create_task(
                self.run_pipeline(
                    frame, current_llm_context)
            )
            current_response_task.add_done_callback(
                functools.partial(self.update_messages, current_llm_context)
            )

    async def run(self):
        self._prerun()

        async_output_queue_marshal_task = asyncio.create_task(
            self._marshal_frames())

        self._camera_thread = threading.Thread(
            target=self._run_camera, daemon=True)
        self._camera_thread.start()

        self._frame_consumer_thread = threading.Thread(
            target=self._frame_consumer, daemon=True)
        self._frame_consumer_thread.start()

        if self._speaker_enabled:
            # TODO-CB: This is interesting
            # self._receive_audio_thread = threading.Thread(
            #     target=self._receive_audio, daemon=True)
            # self._receive_audio_thread.start()

            self._vad_thread = threading.Thread(target=self._vad, daemon=True)
            self._vad_thread.start()

        try:
            while (
                time.time() < self._expiration
                and not self._stop_threads.is_set()
            ):
                await asyncio.sleep(1)
        except Exception as e:
            self._logger.error(f"Exception {e}")
            raise e
        finally:
            # Do anything that must be done to clean up
            self._post_run()

        self._stop_threads.set()

        await self.send_queue.put(EndStreamQueueFrame())
        await async_output_queue_marshal_task
        await self.send_queue.join()
        self._frame_consumer_thread.join()

        if self._speaker_enabled:
            self._receive_audio_thread.join()

    def _post_run(self):
        # Note that this function must be idempotent! It can be called multiple times
        # if, for example, a keyboard interrupt occurs.
        pass

    def stop(self):
        self._stop_threads.set()

    async def stop_when_done(self):
        await self._wait_for_send_queue_to_empty()
        self.stop()

    async def _wait_for_send_queue_to_empty(self):
        await self.send_queue.join()
        self._threadsafe_send_queue.join()

    @abstractmethod
    def write_frame_to_camera(self, frame: bytes):
        pass

    @abstractmethod
    def write_frame_to_mic(self, frame: bytes):
        pass

    @abstractmethod
    def read_audio_frames(self, desired_frame_count):
        return bytes()

    @abstractmethod
    def _prerun(self):
        pass

    def _vad(self):
        # CB: Starting silero VAD stuff
        # TODO-CB: Probably need to force virtual speaker creation if we're
        # going to build this in?
        # TODO-CB: pyaudio installation
        while not self._stop_threads.is_set():
            audio_chunk = self.read_audio_frames(self._vad_samples)
            audio_int16 = np.frombuffer(audio_chunk, np.int16)
            audio_float32 = int2float(audio_int16)
            new_confidence = model(
                torch.from_numpy(audio_float32), 16000).item()
            speaking = new_confidence > 0.5

            if speaking:
                match self._vad_state:
                    case VADState.QUIET:
                        self._vad_state = VADState.STARTING
                        self._vad_starting_count = 1
                    case VADState.STARTING:
                        self._vad_starting_count += 1
                    case VADState.STOPPING:
                        self._vad_state = VADState.SPEAKING
                        self._vad_stopping_count = 0
            else:
                match self._vad_state:
                    case VADState.STARTING:
                        self._vad_state = VADState.QUIET
                        self._vad_starting_count = 0
                    case VADState.SPEAKING:
                        self._vad_state = VADState.STOPPING
                        self._vad_stopping_count = 1
                    case VADState.STOPPING:
                        self._vad_stopping_count += 1

            if self._vad_state == VADState.STARTING and self._vad_starting_count >= self._vad_start_frames:
                print(
                    f'!!! {datetime.datetime.utcnow().isoformat()} queueing start frame')
                asyncio.run_coroutine_threadsafe(
                    self.receive_queue.put(
                        UserStartedSpeakingFrame()), self._loop
                )
                print(f"!!! VAD started, calling interrupt")
                self.interrupt()
                self._vad_state = VADState.SPEAKING
                self._vad_starting_count = 0
            if self._vad_state == VADState.STOPPING and self._vad_stopping_count >= self._vad_stop_frames:
                print(
                    f'!!! {datetime.datetime.utcnow().isoformat()} queueing stop frame')
                asyncio.run_coroutine_threadsafe(
                    self.receive_queue.put(
                        UserStoppedSpeakingFrame()), self._loop
                )
                self._vad_state = VADState.QUIET
                self._vad_stopping_count = 0

    async def _marshal_frames(self):
        while True:
            frame: QueueFrame | list = await self.send_queue.get()
            self._threadsafe_send_queue.put(frame)
            self.send_queue.task_done()
            if isinstance(frame, EndStreamQueueFrame):
                break

    def interrupt(self):
        print(f"!!! setting interrupt")
        self._is_interrupted.set()

    async def get_receive_frames(self) -> AsyncGenerator[QueueFrame, None]:
        while True:
            frame = await self.receive_queue.get()
            yield frame
            if isinstance(frame, EndStreamQueueFrame):
                break

    def _receive_audio(self):
        if not self._loop:
            self._logger.error("No loop available for audio thread")
            return

        seconds = 1
        desired_frame_count = self._speaker_sample_rate * seconds
        while not self._stop_threads.is_set():
            buffer = self.read_audio_frames(desired_frame_count)
            if len(buffer) > 0:
                frame = AudioQueueFrame(buffer)
                asyncio.run_coroutine_threadsafe(
                    self.receive_queue.put(frame), self._loop
                )

        asyncio.run_coroutine_threadsafe(
            self.receive_queue.put(EndStreamQueueFrame()), self._loop
        )

    def _set_image(self, image: bytes):
        self._images = itertools.cycle([image])

    def _set_images(self, images: list[bytes], start_frame=0):
        self._images = itertools.cycle(images)

    def _run_camera(self):
        try:
            while not self._stop_threads.is_set():
                if self._images:
                    this_frame = next(self._images)
                    self.write_frame_to_camera(this_frame)

                time.sleep(1.0 / self._fps)
        except Exception as e:
            self._logger.error(f"Exception {e} in camera thread.")
            raise e

    def _frame_consumer(self):
        self._logger.info("🎬 Starting frame consumer thread")
        b = bytearray()
        smallest_write_size = 3200
        all_audio_frames = bytearray()
        while True:
            try:
                frames_or_frame: QueueFrame | list[QueueFrame] = (
                    self._threadsafe_send_queue.get()
                )
                if isinstance(frames_or_frame, QueueFrame):
                    frames: list[QueueFrame] = [frames_or_frame]
                elif isinstance(frames_or_frame, list):
                    frames: list[QueueFrame] = frames_or_frame
                else:
                    raise Exception("Unknown type in output queue")

                for frame in frames:
                    if isinstance(frame, EndStreamQueueFrame):
                        self._logger.info("Stopping frame consumer thread")
                        self._threadsafe_send_queue.task_done()
                        return

                    # if interrupted, we just pull frames off the queue and discard them
                    if not self._is_interrupted.is_set():
                        print(
                            f"~~~ not interrupted so popping frame of type {type(frame)}")
                        if frame:
                            if isinstance(frame, AudioQueueFrame):
                                chunk = frame.data
                                print(
                                    f"~~~ length of this chunk: {len(chunk)}")
                                all_audio_frames.extend(chunk)

                                b.extend(chunk)
                                truncated_length: int = len(b) - (
                                    len(b) % smallest_write_size
                                )
                                if truncated_length:
                                    self.write_frame_to_mic(
                                        bytes(b[:truncated_length]))
                                    b = b[truncated_length:]
                            elif isinstance(frame, ImageQueueFrame):
                                self._set_image(frame.image)
                            elif isinstance(frame, SpriteQueueFrame):
                                self._set_images(frame.images)
                            elif isinstance(frame, TTSCompletedFrame):
                                self.append_to_context(
                                    "assistant", frame.text)
                        elif len(b):
                            self.write_frame_to_mic(bytes(b))
                            b = bytearray()
                    else:
                        # if there are leftover audio bytes, write them now; failing to do so
                        # can cause static in the audio stream.
                        print(f"!!! interrupted, flushing audio")
                        if len(b):
                            truncated_length = len(b) - (len(b) % 160)
                            self.write_frame_to_mic(
                                bytes(b[:truncated_length]))
                            b = bytearray()

                        if isinstance(frame, StartStreamQueueFrame):
                            self._is_interrupted.clear()

                self._threadsafe_send_queue.task_done()
            except queue.Empty:
                if len(b):
                    self.write_frame_to_mic(bytes(b))

                b = bytearray()
            except Exception as e:
                self._logger.error(
                    f"Exception in frame_consumer: {e}, {len(b)}")
                raise e