import time
import wave
from contextlib import contextmanager

import pyaudiowpatch as pyaudio
from TheCodeLabs_BaseUtils.DefaultLogger import DefaultLogger

LOG_FORMAT = '[%(levelname)-7s] - %(asctime)s - %(message)s'
LOGGER = DefaultLogger().create_logger_if_not_exists('AudioRecorder', logFormat=LOG_FORMAT)


class AudioRecorder:
    _CHUNK_SIZE = 512

    def __init__(self, deviceName: str, destinationFilePath: str):
        self._deviceName = deviceName
        self._destinationFilePath = destinationFilePath

    def __get_device_by_name(self, pyaudioInstance, deviceName: str):
        deviceNames = []
        for loopback in pyaudioInstance.get_loopback_device_info_generator():
            deviceNames.append(loopback['name'])
            if deviceName in loopback['name']:
                return loopback
        else:
            message = f'Device with name {self._deviceName} not found.'
            message += f'\nAvailable devices:'
            for deviceName in deviceNames:
                message += f'\n\t{deviceName}'
            raise RuntimeError(message)

    @contextmanager
    def record(self):
        with pyaudio.PyAudio() as pyaudioInstance:
            device = self.__get_device_by_name(pyaudioInstance, self._deviceName)
            LOGGER.info(f'Recording from: ({device["index"]}) {device["name"]}')

            waveFile = wave.open(self._destinationFilePath, 'wb')
            waveFile.setnchannels(device['maxInputChannels'])
            waveFile.setsampwidth(pyaudio.get_sample_size(pyaudio.paInt16))
            waveFile.setframerate(int(device['defaultSampleRate']))

            def callback(in_data, frame_count, time_info, status):
                waveFile.writeframes(in_data)
                return in_data, pyaudio.paContinue

            with pyaudioInstance.open(format=pyaudio.paInt16,
                                      channels=device['maxInputChannels'],
                                      rate=int(device['defaultSampleRate']),
                                      frames_per_buffer=self._CHUNK_SIZE,
                                      input=True,
                                      input_device_index=device['index'],
                                      stream_callback=callback
                                      ) as stream:
                yield stream

            waveFile.close()
            LOGGER.info(f'Recording stopped. File: "{self._destinationFilePath}"')


if __name__ == '__main__':
    recorder = AudioRecorder(
        deviceName='3/4 - Musik (2- GIGAPort HD Audio driver) [Loopback]',
        destinationFilePath='C:/Users/RobertG/Desktop/output.wav'
    )

    with recorder.record() as stream:
        time.sleep(5)