import pyaudio
import wave, logging, threading, time, queue, signal, argparse
from os import path, makedirs

logging.basicConfig(format="%(asctime)s: %(message)s", level=logging.INFO,datefmt="%H:%M:%S")

class DiscordRecorder:
    def __init__(self, DEVICE_ID, CHUNK = 1024, FORMAT = pyaudio.paInt16, CHANNELS = 2, RATE = 48000, FILENAME = "./recs/radio.wav"):
        self.pa_instance = pyaudio.PyAudio()

        self.DEVICE_ID = DEVICE_ID
        self.CHUNK = CHUNK
        self.FORMAT = FORMAT
        self.CHANNELS = CHANNELS
        self.RATE = RATE

        self.FILENAME = FILENAME
        self._check_file_path_exists()

        self.queued_frames = queue.Queue()

        self.stop_threads = threading.Event()

        self.recording_thread = None
        self.saving_thread = None

        self.running_stream = None

    # Wrapper to check if the given filepath (not file itself) exists
    def _check_file_path_exists(self):
        if not path.exists(path.dirname(self.FILENAME)):            
            makedirs(path.dirname(self.FILENAME), exist_ok=True)

    # Wrapper for the recorder thread; Adds new data to the queue
    def _recorder(self):
        logging.info("* Recording Thread Starting")
        while True:
            data = self.running_stream.read(self.CHUNK)            
            self.queued_frames.put(data)

            # check for stop
            if self.stop_threads.is_set():
                break

    # Wrapper for saver thread; Saves the queue to the file
    def _saver(self):
        logging.info("* Saving Thread Starting")
        while True:
            if not self.queued_frames.empty():
                dequeued_frames = []
                for i in range(self.queued_frames.qsize()):
                    dequeued_frames.append(self.queued_frames.get())

                if not path.isfile(self.FILENAME):
                    wf = wave.open(self.FILENAME, 'wb')
                    wf.setnchannels(self.CHANNELS)
                    wf.setsampwidth(self.pa_instance.get_sample_size(self.FORMAT))
                    wf.setframerate(self.RATE)
                    wf.writeframes(b''.join(dequeued_frames))
                    wf.close()
                else:                                
                    read_file = wave.open(self.FILENAME, 'rb')
                    read_file_data = read_file.readframes(read_file.getnframes())
                    read_file.close()
                        
                    wf = wave.open(self.FILENAME, 'wb')
                    wf.setnchannels(self.CHANNELS)
                    wf.setsampwidth(self.pa_instance.get_sample_size(self.FORMAT))
                    wf.setframerate(self.RATE)

                    wf.writeframes(read_file_data)
                    wf.writeframes(b''.join(dequeued_frames))
                    wf.close()

            # check for stop
            if self.stop_threads.is_set():
                break

            time.sleep(5)

    # Start the recording function
    def start_recording(self):
        logging.info("* Recording")

        self.running_stream = self.pa_instance.open(
            input_device_index=self.DEVICE_ID,
            format=self.FORMAT,
            channels=self.CHANNELS,
            rate=self.RATE,
            input=True,
            frames_per_buffer=self.CHUNK
        )
        
        self.recording_thread = threading.Thread(target=self._recorder)
        self.recording_thread.start()

        self.saving_thread = threading.Thread(target=self._saver)
        self.saving_thread.start()
    
    # Stop the recording function
    def stop_recording(self):
        self.stop_threads.set()
        self.recording_thread.join()
        self.saving_thread.join()
        self.running_stream.stop_stream()
        self.running_stream.close()
        self.pa_instance.terminate()
        
        logging.info("* Done recording")


class GracefulExitCatcher:
  def __init__(self, stop_callback):
    self.stop = False

    # The function to run when the exit signal is caught
    self.stop_callback = stop_callback

    # Update what happens when these signals are caught
    signal.signal(signal.SIGINT, self.exit_gracefully)
    signal.signal(signal.SIGTERM, self.exit_gracefully)

  def exit_gracefully(self, *args):
    logging.info("* Stop signal caught...")
    
    # Stop the main loop
    self.stop = True

    # Run the given callback function
    self.stop_callback()




if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("deviceId", type=int, help="The ID of the audio device to use")
    parser.add_argument("filename", type=str, help="The filepath/filename of the output file")    
    args = parser.parse_args()

    logging.debug("Arguments:", args)
    
    recorder = DiscordRecorder(args.deviceId, FILENAME=args.filename)

    exit_catcher = GracefulExitCatcher(recorder.stop_recording)

    recorder.start_recording()

    while not exit_catcher.stop:
        time.sleep(1)