# -*- coding: utf-8 -*-

from __future__ import annotations

import logging

from ..sinks import AudioSink

log = logging.getLogger(__name__)

__all__ = [
    'SpeechRecognitionSink',
]

try:
    import speech_recognition as sr
except ImportError:

    def __getattr__(name: str):
        if name in __all__:
            raise RuntimeError('The SpeechRecognition module is required to use this sink.')

else:
    import time
    import array
    import asyncio
    import audioop

    from collections import defaultdict

    from ..rtp import SilencePacket

    from typing import TYPE_CHECKING, TypedDict

    if TYPE_CHECKING:
        from concurrent.futures import Future as CFuture
        from typing import Literal, Callable, Optional, Any, Final, Protocol, Awaitable, TypeVar

        from discord import Member

        from ..opus import VoiceData
        from ..types import MemberOrUser as User

        T = TypeVar('T')

        # [r.split('_', 1)[1] for r in dir(sr.Recognizer()) if r.startswith("recognize")]
        SRRecognizerMethod = Literal[
            'amazon',
            'api',
            'assemblyai',
            'azure',
            'bing',
            'faster_whisper',
            'google',
            'google_cloud',
            'groq',
            'houndify',
            'ibm',
            'lex',
            'openai',
            'sphinx',
            'tensorflow',
            'vosk',
            'whisper',
            'wit',
        ]

        class SRStopper(Protocol):
            def __call__(self, wait: bool = True, /) -> None: ...

        SRProcessDataCB = Callable[[sr.Recognizer, sr.AudioData, User], Optional[str]]
        SRTextCB = Callable[[User, str], Any]

    class _StreamData(TypedDict):
        stopper: Optional[SRStopper]
        recognizer: sr.Recognizer
        buffer: array.array[int]

    class SpeechRecognitionSink(AudioSink):
        def __init__(
            self,
            *,
            process_cb: Optional[SRProcessDataCB] = None,
            text_cb: Optional[SRTextCB] = None,
            default_recognizer: SRRecognizerMethod = 'google',
            phrase_time_limit: int = 10,
            ignore_silence_packets: bool = True,
        ):
            super().__init__(None)
            self.process_cb: Optional[SRProcessDataCB] = process_cb
            self.text_cb: Optional[SRTextCB] = text_cb
            self.phrase_time_limmit: int = phrase_time_limit
            self.ignore_silence_packets: bool = ignore_silence_packets

            self.default_recognizer: SRRecognizerMethod = default_recognizer
            self._stream_data: defaultdict[int, _StreamData] = defaultdict(
                lambda: _StreamData(stopper=None, recognizer=sr.Recognizer(), buffer=array.array('B'))
            )

        def _await(self, coro: Awaitable[T]) -> CFuture[T]:
            assert self.client is not None
            return asyncio.run_coroutine_threadsafe(coro, self.client.loop)  # type: ignore

        def wants_opus(self) -> bool:
            return False

        def write(self, user: Optional[User], data: VoiceData) -> None:
            if self.ignore_silence_packets and isinstance(data.packet, SilencePacket):
                return

            if user is None:
                return

            sdata = self._stream_data[user.id]
            sdata['buffer'].extend(data.pcm)

            if not sdata['stopper']:
                sdata['stopper'] = sdata['recognizer'].listen_in_background(
                    DiscordSRAudioSource(sdata['buffer']), self.background_listener(user), self.phrase_time_limmit
                )

        def background_listener(self, user: User):
            process_cb = self.process_cb or self.get_default_process_callback()
            text_cb = self.text_cb or self.get_default_text_callback()

            def callback(_recognizer: sr.Recognizer, _audio: sr.AudioData):
                output = process_cb(_recognizer, _audio, user)
                if output is not None:
                    text_cb(user, output)

            return callback

        def get_default_process_callback(self) -> SRProcessDataCB:
            def cb(recognizer: sr.Recognizer, audio: sr.AudioData, user: Optional[User]) -> Optional[str]:
                log.debug("Got %s, %s, %s", audio, audio.sample_rate, audio.sample_width)
                text: Optional[str] = None
                try:
                    # they changed recognize_google to be optionally assigned at runtime...
                    func = getattr(recognizer, 'recognize_' + self.default_recognizer, recognizer.recognize_google)  # type: ignore
                    text = func(audio)
                except sr.UnknownValueError:
                    log.debug("Bad speech chunk")
                    # self._debug_audio_chunk(audio)

                return text

            return cb

        def get_default_text_callback(self) -> SRTextCB:
            def cb(user: Optional[User], text: Optional[str]) -> Any:
                log.info("%s said: %s", user.display_name if user else 'Someone', text)

            return cb

        @AudioSink.listener()
        def on_voice_member_disconnect(self, member: Member, ssrc: Optional[int]) -> None:
            if member is not None:
                self._drop(member.id)

        def cleanup(self) -> None:
            for user_id in tuple(self._stream_data.keys()):
                self._drop(user_id)

        def _drop(self, user_id: int) -> None:
            data = self._stream_data.pop(user_id, None)
            if data is None:
                log.debug("Cannot drop user id: %s, no data", user_id)
                return

            stopper = data.get('stopper')
            if stopper:
                stopper()

            buffer = data.get('buffer')
            if buffer:
                # arrays don't have a clear function
                del buffer[:]

        def _debug_audio_chunk(self, audio: sr.AudioData, filename: str = 'sound.wav') -> None:
            import io, wave, discord

            with io.BytesIO() as b:
                with wave.open(b, 'wb') as writer:
                    writer.setframerate(48000)
                    writer.setsampwidth(2)
                    writer.setnchannels(2)
                    writer.writeframes(audio.get_wav_data())

                b.seek(0)
                f = discord.File(b, filename)
                self._await(self.voice_client.channel.send(file=f))  # type: ignore

    class DiscordSRAudioSource(sr.AudioSource):
        little_endian: Final[bool] = True
        SAMPLE_RATE: Final[int] = 48_000
        SAMPLE_WIDTH: Final[int] = 2
        CHANNELS: Final[int] = 2
        CHUNK: Final[int] = 960

        def __init__(self, buffer: array.array[int]):
            self.buffer = buffer
            self._entered: bool = False

        @property
        def stream(self):
            return self

        def __enter__(self):
            if self._entered:
                log.warning('Already entered sr audio source')
            self._entered = True
            return self

        def __exit__(self, *exc) -> None:
            self._entered = False
            if any(exc):
                log.exception('Error closing sr audio source')

        def read(self, size: int) -> bytes:
            # TODO: make this timeout configurable
            for _ in range(10):
                if len(self.buffer) < size * self.CHANNELS:
                    time.sleep(0.1)
                else:
                    break
            else:
                if len(self.buffer) == 0:
                    return b''

            chunksize = size * self.CHANNELS
            audiochunk = self.buffer[:chunksize].tobytes()
            del self.buffer[: min(chunksize, len(audiochunk))]
            audiochunk = audioop.tomono(audiochunk, 2, 1, 1)
            return audiochunk

        def close(self) -> None:
            self.buffer.clear()
