# Copyright 2023 LiveKit, Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from __future__ import annotations

import asyncio
import contextlib
import logging
from collections import deque
from typing import List, Optional

import numpy as np
import torch
from livekit import agents, rtc


class VAD(agents.vad.VAD):
    def __init__(
        self, *, model_path: Optional[str] = None, use_onnx: bool = True
    ) -> None:
        if model_path:
            model = torch.jit.load(model_path)
            model.eval()
        else:
            model, _ = torch.hub.load(
                repo_or_dir="snakers4/silero-vad",
                model="silero_vad",
                onnx=use_onnx,
            )
        self._model = model

    def stream(
        self,
        *,
        min_speaking_duration: float = 0.5,
        min_silence_duration: float = 0.5,
        padding_duration: float = 0.1,
        sample_rate: int = 16000,
        max_buffered_speech: float = 45.0,
        threshold: float = 0.5,
    ) -> "VADStream":
        return VADStream(
            self._model,
            min_speaking_duration=min_speaking_duration,
            min_silence_duration=min_silence_duration,
            padding_duration=padding_duration,
            sample_rate=sample_rate,
            max_buffered_speech=max_buffered_speech,
            threshold=threshold,
        )


# Based on https://github.com/snakers4/silero-vad/blob/94504ece54c8caeebb808410b08ae55ee82dba82/utils_vad.py#L428
class VADStream(agents.vad.VADStream):
    def __init__(
        self,
        model,
        *,
        min_speaking_duration: float,
        min_silence_duration: float,
        padding_duration: float,
        sample_rate: int,
        max_buffered_speech: float,
        threshold: float,
    ) -> None:
        self._min_speaking_duration = min_speaking_duration
        self._min_silence_duration = min_silence_duration
        self._padding_duration = padding_duration
        self._sample_rate = sample_rate
        self._max_buffered_speech = max_buffered_speech
        self._threshold = threshold

        if sample_rate not in [8000, 16000]:
            raise ValueError("Silero VAD only supports 8KHz and 16KHz sample rates")

        self._queue = asyncio.Queue[rtc.AudioFrame | None]()
        self._event_queue = asyncio.Queue[agents.vad.VADEvent | None]()
        self._model = model

        self._closed = False
        self._speaking = False
        self._waiting_start = False
        self._waiting_end = False
        self._current_sample = 0
        self._min_speaking_samples = min_speaking_duration * sample_rate
        self._min_silence_samples = min_silence_duration * sample_rate
        self._padding_duration_samples = padding_duration * sample_rate
        self._max_buffered_samples = max_buffered_speech * sample_rate

        self._queued_frames: deque[rtc.AudioFrame] = deque()
        self._original_frames: deque[rtc.AudioFrame] = deque()
        self._buffered_frames: List[rtc.AudioFrame] = []
        self._main_task = asyncio.create_task(self._run())

        def log_exception(task: asyncio.Task) -> None:
            if not task.cancelled() and task.exception():
                logging.error(f"silero vad task failed: {task.exception()}")

        self._main_task.add_done_callback(log_exception)

    def push_frame(self, frame: rtc.AudioFrame) -> None:
        if self._closed:
            raise ValueError("cannot push frame to closed stream")

        self._queue.put_nowait(frame)

    async def aclose(self, *, wait: bool = True) -> None:
        self._closed = True
        if not wait:
            self._main_task.cancel()

        self._queue.put_nowait(None)
        with contextlib.suppress(asyncio.CancelledError):
            await self._main_task

    async def _run(self):
        try:
            while True:
                frame = await self._queue.get()
                if frame is None:
                    break  # None is sent inside aclose

                self._queue.task_done()

                # resample to silero's sample rate
                resampled_frame = frame.remix_and_resample(
                    self._sample_rate, 1
                )  # TODO: This is technically wrong, fix when we have a better resampler
                self._original_frames.append(frame)
                self._queued_frames.append(resampled_frame)

                # run inference by chunks of 40ms until we run out of data
                while True:
                    available_length = sum(
                        f.samples_per_channel for f in self._queued_frames
                    )

                    samples_40ms = self._sample_rate // 1000 * 40
                    if available_length < samples_40ms:
                        break

                    await asyncio.shield(self._run_inference())
        finally:
            self._event_queue.put_nowait(None)

    async def _run_inference(self) -> None:
        # merge the first 4 frames (we know each is 10ms)
        if len(self._queued_frames) < 4:
            return

        original_frames = [self._original_frames.popleft() for _ in range(4)]
        merged_frame = agents.utils.merge_frames(
            [self._queued_frames.popleft() for _ in range(4)]
        )

        # convert data_40ms to tensor & f32
        tensor = torch.from_numpy(np.frombuffer(merged_frame.data, dtype=np.int16))
        tensor = tensor.to(torch.float32) / 32768.0

        # run inference
        speech_prob = await asyncio.to_thread(
            lambda: self._model(tensor, self._sample_rate).item()
        )
        self._dispatch_event(speech_prob, original_frames)
        self._current_sample += merged_frame.samples_per_channel

    def _dispatch_event(self, speech_prob: int, original_frames: List[rtc.AudioFrame]):
        """
        Dispatches a VAD event based on the speech probability and the options
        Args:
            speech_prob: speech probability of the current frame
            original_frames: original frames of the current inference
        """

        samples_10ms = self._sample_rate / 100
        padding_count = int(
            self._padding_duration_samples // samples_10ms
        )  # number of frames to keep for the padding (one side)

        self._buffered_frames.extend(original_frames)
        if (
            not self._speaking
            and not self._waiting_start
            and len(self._buffered_frames) > padding_count
        ):
            self._buffered_frames = self._buffered_frames[
                len(self._buffered_frames) - padding_count :
            ]

        max_buffer_len = padding_count + max(
            int(self._max_buffered_samples // samples_10ms),
            int(self._min_speaking_samples // samples_10ms),
        )
        if len(self._buffered_frames) > max_buffer_len:
            # if unaware of this, may be hard to debug, so logging seems ok here
            logging.warning(
                f"VAD buffer overflow, dropping {len(self._buffered_frames) - max_buffer_len} frames"
            )
            self._buffered_frames = self._buffered_frames[
                len(self._buffered_frames) - max_buffer_len :
            ]

        if speech_prob >= self._threshold:
            # speaking, wait for min_speaking_duration to trigger START_OF_SPEECH
            self._waiting_end = False
            if not self._waiting_start and not self._speaking:
                self._waiting_start = True
                self._start_speech = self._current_sample

            if self._waiting_start and (
                self._current_sample - self._start_speech >= self._min_speaking_samples
            ):
                self._waiting_start = False
                self._speaking = True
                event = agents.vad.VADEvent(
                    type=agents.vad.VADEventType.START_OF_SPEECH,
                    samples_index=self._start_speech,
                )
                self._event_queue.put_nowait(event)

                # since we're waiting for the min_spaking_duration to trigger START_OF_SPEECH,
                # the SPEAKING data is missing the first few frames, trigger it here
                # TODO(theomonnom): Maybe it is better to put the data inside the START_OF_SPEECH event?
                event = agents.vad.VADEvent(
                    type=agents.vad.VADEventType.SPEAKING,
                    samples_index=self._start_speech,
                    speech=self._buffered_frames[padding_count:],
                )

                return

        if self._speaking:
            # we don't check the speech_prob here
            event = agents.vad.VADEvent(
                type=agents.vad.VADEventType.SPEAKING,
                samples_index=self._current_sample,
                speech=original_frames,
            )
            self._event_queue.put_nowait(event)

        if speech_prob < self._threshold:
            # stopped speaking, wait for min_silence_duration to trigger END_OF_SPEECH,
            self._waiting_start = False
            if not self._waiting_end and self._speaking:
                self._waiting_end = True
                self._end_speech = self._current_sample

            if self._waiting_end and (
                self._current_sample - self._end_speech
                >= max(self._min_silence_samples, self._padding_duration_samples)
            ):
                self._waiting_end = False
                self._speaking = False
                event = agents.vad.VADEvent(
                    type=agents.vad.VADEventType.END_OF_SPEECH,
                    samples_index=self._end_speech,
                    duration=(self._current_sample - self._start_speech)
                    / self._sample_rate,
                    speech=self._buffered_frames,
                )
                self._event_queue.put_nowait(event)

    async def __anext__(self) -> agents.vad.VADEvent:
        evt = await self._event_queue.get()
        if evt is None:
            raise StopAsyncIteration

        return evt
