from functools import partial
from itertools import chain
from contextlib import contextmanager
from typing import Iterable, Union, Callable, Optional

from stream2py.stream_buffer import StreamBuffer
from stream2py.sources.audio import PyAudioSourceReader
from taped.util import DFLT_SR, DFLT_SAMPLE_WIDTH, DFLT_CHK_SIZE, \
    DFLT_STREAM_BUF_SIZE_S, bytes_to_waveform, ensure_source_input_device_index

from itertools import islice
from dataclasses import dataclass
from creek import Creek

# TODO: revise wrapping so it shows this module and name, not some <locals> stuff
@Creek.wrap
@dataclass
class BufferItems(StreamBuffer):
    """A generator of live chunks of audio bytes taken from a stream sourced from specified microphone.

    :param input_device_index: Index of Input Device to use. Unspecified (or None) uses default device.
    :param sr: Specifies the desired sample rate (in Hz)
    :param sample_bytes: Sample width in bytes (1, 2, 3, or 4)
    :param sample_width: Specifies the number of frames per buffer.
    :param stream_buffer_size_s: How many seconds of data to keep in the buffer (i.e. how far in the past you can see)
    """
    input_device_index: Optional[int] = None
    sr: int = DFLT_SR
    sample_width: int = DFLT_SAMPLE_WIDTH
    chk_size: int = DFLT_CHK_SIZE
    stream_buffer_size_s: Union[float, int] = DFLT_STREAM_BUF_SIZE_S

    def __post_init__(self):
        self.input_device_index = ensure_source_input_device_index(self.input_device_index)
        seconds_per_read = self.chk_size / self.sr

        self.maxlen = int(self.stream_buffer_size_s / seconds_per_read)
        self.source_reader = PyAudioSourceReader(rate=self.sr, width=self.sample_width,
                                                 unsigned=True,
                                                 input_device_index=self.input_device_index,
                                                 frames_per_buffer=self.chk_size)

        super().__init__(source_reader=self.source_reader, maxlen=self.maxlen)


class ByteChunks(BufferItems):
    def data_to_obj(self, data):
        return data[1]


# TODO: chain from ByteChunks instead of BufferItems
# TODO: use more stable and flexible bytes_to_waveform
class WfChunks(ByteChunks):
    sample_width = 2

    def data_to_obj(self, data):
        data = super().data_to_obj(data)
        return bytes_to_waveform(data, sample_width=self.sample_width)


# TODO: chain from ByteChunks instead of BufferItems
# TODO: use more stable and flexible bytes_to_waveform
class LiveWf(WfChunks):
    def post_iter(self, obj):
        return chain.from_iterable(obj)

    def __getitem__(self, item):
        if not isinstance(item, slice):
            item = slice(item, item + 1)  # to emulate usual list[i] interface
        return list(islice(self, item.start, item.stop, item.step))




# class LiveAudioChks(StreamBuffer):
#     """A generator of live chunks of audio bytes taken from a stream sourced from specified microphone.
#
#     :param input_device_index: Index of Input Device to use. Unspecified (or None) uses default device.
#     :param sr: Specifies the desired sample rate (in Hz)
#     :param sample_bytes: Sample width in bytes (1, 2, 3, or 4)
#     :param sample_width: Specifies the number of frames per buffer.
#     :param stream_buffer_size_s: How many seconds of data to keep in the buffer (i.e. how far in the past you can see)
#     """
#
#     def __init__(self, input_device_index = None,
#         sr = DFLT_SR,
#         sample_width = DFLT_SAMPLE_WIDTH,
#         chk_size = DFLT_CHK_SIZE,
#         stream_buffer_size_s = DFLT_STREAM_BUF_SIZE_S):
#             self.input_device_index = input_device_index
#             self.sr = sr
#             self.sample_width = sample_width
#             self.chk_size = chk_size
#             self.stream_buffer_size_s = stream_buffer_size_s
#
#             self.input_device_index = ensure_source_input_device_index(self.input_device_index)
#             seconds_per_read = self.chk_size / self.sr
#
#             self.maxlen = int(self.stream_buffer_size_s / seconds_per_read)
#             self.source_reader = PyAudioSourceReader(rate=self.sr, width=self.sample_width,
#                                                      unsigned=True,
#                                                      input_device_index=self.input_device_index,
#                                                      frames_per_buffer=self.chk_size)
#
#             super().__init__(source_reader=self.source_reader, maxlen=self.maxlen)
#
#     def __iter__(self):
#         yield from self.source_reader


class LiveWfOld(BufferItems):
    def __enter__(self):
        super().__enter__()
        self._bytes_to_waveform = partial(bytes_to_waveform,
                                          sr=self.sr,
                                          sample_width=self.sample_width)
        self.live_wf = chain.from_iterable(map(lambda x: self._bytes_to_waveform(x[1]), self))
        return self

    # TODO: Protect from using this, or make it work
    # def __iter__(self):
    #     raise NotImplementedError("")
    #     # yield from self.live_wf

    def __getitem__(self, item):
        if not isinstance(item, slice):
            item = slice(item, item + 1)  # to emulate usual list[i] interface
        return list(islice(self.live_wf, item.start, item.stop, item.step))

    def __exit__(self, *args, **kwargs):
        return super().__exit__(*args, **kwargs)


# from dataclasses import dataclass
#
#
# @dataclass
# class LiveWf:
#     input_device_index = None
#     sr = DFLT_SR
#     sample_width = DFLT_SAMPLE_WIDTH
#     n_channels = DFLT_N_CHANNELS
#     chk_size = DFLT_CHK_SIZE
#     stream_buffer_size_s = DFLT_STREAM_BUF_SIZE_S
#
#     # def __post_init__(self):
#     #     self._kwargs = self.__dict__
#
#     def __enter__(self):
#         # with live_audio_chks(**self._kwargs) as live_audio_chunks:
#         self.live_audio_chunks = live_audio_chks(
#             input_device_index=self.input_device_index,
#             sr=self.sr,
#             sample_width=self.sample_width,
#             n_channels=self.n_channels,
#             chk_size=self.chk_size,
#             stream_buffer_size_s=self.stream_buffer_size_s)
#         _bytes_to_waveform = partial(audio_pokes_version_of_bytes_to_waveform,
#                                      sr=self.sr,
#                                      n_channels=self.n_channels,
#                                      sample_width=self.sample_width)
#         self.live_wf = chain.from_iterable(map(lambda x: _bytes_to_waveform(x[1]), self.live_audio_chunks))
#         return self
#
#     def __exit__(self, *args, **kwargs):
#         self.live_audio_chunks.__exit__(*args, **kwargs)

#
# def live_wf(input_device_index=None, sr=DFLT_SR, sample_width=DFLT_SAMPLE_WIDTH,
#             chk_size=DFLT_CHK_SIZE, stream_buffer_size_s=DFLT_STREAM_BUF_SIZE_S):
#     """A generator of live waveform sample values taken from a stream sourced from specified microphone.
#
#     :param input_device_index: Index of Input Device to use. Unspecified (or None) uses default device.
#     :param sr: Specifies the desired sample rate (in Hz)
#     :param sample_width: Sample width in bytes (1, 2, 3, or 4)
#     :param n_channels: The desired number of input channels. Ignored if input_device is not specified (or None).
#     :param stream_buffer_size_s: How many seconds of data to keep in the buffer (i.e. how far in the past you can see)
#
#     >>> from time import sleep
#     >>> from itertools import islice
#     >>> # enter the id of your microphone and get a live waveform source!
#     >>> # (if None, will try to figure it out)
#     >>> wf_gen = live_wf(input_device_index=None)
#     >>>
#     >>> # Now wait a bit, say some silly things, then ask for a few samples...
#     >>> sleep(1.2)
#     >>> wf = list(islice(wf_gen, 0, 44100 * 1))
#     >>> # and now listen to that wf and be embarrassed...
#     >>> # ... or just look at the size (less fun though)
#     >>> len(wf)
#     44100
#
#     Don't forget to close! (or use live_wf_ctx context manager).
#     >>> wf_gen.close()
#
#     After wf_gen is closed, you can still ask it for data.
#     It just won't give you any.
#     >>> wf = list(islice(wf_gen, 0, 44100 * 1))
#     >>> len(wf)
#     0
#
#     Here wf_gen is a generator, so closing means: https://docs.python.org/2.5/whatsnew/pep-342.html
#     """
#     # TODO: Find a way to copy from containing function's signature and calling LiveAudioChunks with that
#     with live_wf_ctx(input_device_index=input_device_index,
#                      sr=sr, sample_width=sample_width,
#                      chk_size=chk_size,
#                      stream_buffer_size_s=stream_buffer_size_s) as live_wf:
#         yield from live_wf


# live_wf.list_device_info = PyAudioSourceReader.list_device_info


def simple_chunker(a: Iterable,
                   chk_size: int):
    """Generate fixed sized non-overlapping chunks of an iterable ``a``.

    >>> list(simple_chunker(range(7), 3))
    [(0, 1, 2), (3, 4, 5)]

    Most of the time, you'll want to fix the parameters of the chunker like this:

    >>> from functools import partial
    >>> chunker = partial(simple_chunker, chk_size=3)
    >>> list(chunker(range(7)))
    [(0, 1, 2), (3, 4, 5)]

    Note, the type of the chunks is always tuples, but you can easily change that using ``map``.
    For example, to change the type to be list:

    >>> list(map(list, chunker(range(7))))
    [[0, 1, 2], [3, 4, 5]]

    >>> a = range(6)
    >>> list(simple_chunker(a, 3))
    [(0, 1, 2), (3, 4, 5)]
    >>> list(simple_chunker(a, 2))
    [(0, 1), (2, 3), (4, 5)]
    >>> list(simple_chunker(a, 1))
    [(0,), (1,), (2,), (3,), (4,), (5,)]

    """
    return zip(*([iter(a)] * chk_size))


def rechunker(chks: Iterable[Iterable],
              chunker: Union[Callable, int]):
    """Generate fixed sized non-overlapping chunks of an iterable of chunks.
    That is, the rechunker applies a chunker to an unraveled stream of chunks,
    or more generally of iterables since they can be of varied sizes and types.

    >>> from functools import partial
    >>> chunker = partial(simple_chunker, chk_size=3)
    >>> chks = [[0], (1, 2, 3), [4, 5], iter((6, 7))]  # iterable of (different types of) iterables
    >>> list(rechunker(chks, chunker))
    [(0, 1, 2), (3, 4, 5)]

    """
    if isinstance(chunker, int):  # if chunker is an int, take it to be a the chk_size of a simple_chunker
        chk_size = chunker
        chunker = partial(simple_chunker, chk_size)
    yield from chunker(chain.from_iterable(chks))


def record_some_sound(save_to_file,
                      input_device_index=None, sr=DFLT_SR, sample_width=DFLT_SAMPLE_WIDTH,
                      chk_size=DFLT_CHK_SIZE, stream_buffer_size_s=DFLT_STREAM_BUF_SIZE_S, verbose=True,
                      ):
    def get_write_file_stream():
        if isinstance(save_to_file, str):
            return open(save_to_file, 'wb')
        else:
            return save_to_file  # assume it's already a stream

    def clog(*args, **kwargs):
        if verbose:
            print(*args, **kwargs)

    seconds_per_read = chk_size / sr
    maxlen = int(stream_buffer_size_s / seconds_per_read)

    source_reader = PyAudioSourceReader(rate=sr, width=sample_width, unsigned=True,
                                        input_device_index=input_device_index,
                                        frames_per_buffer=chk_size)

    with StreamBuffer(source_reader=source_reader, maxlen=maxlen) as stream_buffer:
        """keep open and save to file until stop event"""
        clog("starting the recording...")
        with get_write_file_stream() as write_stream:
            while True:
                try:
                    chk = source_reader.read()
                    print(type(chk), len(chk))
                except KeyboardInterrupt:
                    clog("stopping the recording...")
                    break

    clog('Done.')
