# Original code by TheGreydiamond Copyright 2020-2021
import pyaudio, numpy, keyboard, time
import scipy.io.wavfile as wav
from sys import getsizeof
import os
from os.path import exists
import json
import shutil

##### CONFIG (Now read from file) #####

if(not exists("config.json")):
    shutil.copy("config.json.example", "config.json")
    print("No config found. Created default config")

with open('config.json', 'r') as myfile:
    data = myfile.read()

# parse file
conf = json.loads(data)

RATE = conf["SampleRate"]                                      # Sample rate
RECORD_SECONDS = conf["SecondsToRecord"]                       # Seconds to record (aka. last X seconds)
SourcesToRecord = conf["SourcesToRecord"]                      # Record all sources which contains these strings
SaveDirectory= conf["SaveDirectory"]                           # The directory to save to
shortcut = conf["Shortcut"]                                    # The shortcut to trigger the save

##### CONFIG END #####

CHUNKSIZE = 1024*2

print("Searching for audio sources")
Sources = []
AudioChannels = []
Frames = []
p = pyaudio.PyAudio()                                         # An instance of PyAudio
info = p.get_host_api_info_by_index(0)
numdevices = info.get('deviceCount')                          # The amount of audio devices
samplesI = 0
saves = 0
cleanUp = False
incomingAudioChan = -1
localAudioChan = -1

## Find the actual device pointers
for i in range(0, numdevices):
        if (p.get_device_info_by_host_api_device_index(0, i).get('maxInputChannels')) > 0:
            i2 = 0
            while(i2 < len(SourcesToRecord)):
                if(SourcesToRecord[i2] in p.get_device_info_by_host_api_device_index(0, i).get('name')):
                    print("Using " + p.get_device_info_by_host_api_device_index(0, i).get('name') + " for '" + SourcesToRecord[i2] + "' ID: " + str(i))
                    Sources.append(i)
                    AudioChannels.append(p.open(format=pyaudio.paInt16, channels=1, rate=RATE, input=True, frames_per_buffer=CHUNKSIZE, input_device_index=i))
                i2+=1

## Populate frame buffer
i = 0 
while(i <= len(SourcesToRecord)):
    Frames.append([])
    i += 1

# Save the replay
def on_triggered():
    global cb, saves

    print("Saving replay...")

    # save current working directory and switch to saving path
    cwd = os.getcwd()
    os.chdir(SaveDirectory)

    # save the audio
    cutUpBuffers = []
    i = 0
    name  = time.strftime("%Y-%m-%d_%H-%M-%S-out-")
    while(i < len(Frames) - 1):
        print("Saving file: " + name + str(saves) + str(i) + '.wav', "to", os.getcwd())
        wav.write(name + str(saves) + str(i) + '.wav', RATE, numpy.hstack(Frames[i][-320:]))
        i += 1
    saves+=1

    # return back to original working directory
    os.chdir(cwd)
    print("Saving done")


# add the keyboard hotkey listener
keyboard.add_hotkey(shortcut, on_triggered)

print(" There are " + str(len(SourcesToRecord)) + " SourcesToRecord. Frames: " + str(len(Frames)))

print("Running!")
try:
    while True:
        i = 0
        while(i < len(SourcesToRecord)):
            data = AudioChannels[i].read(CHUNKSIZE)
            # print(i)
            Frames[i].append(numpy.fromstring(data, dtype=numpy.int16))
            i += 1

        samplesI+=1
        if(len(Frames[i]) >= 900):
            cleanUp = True

        if(len(Frames[i]) <= 700):
            cleanUp = False
        
        if(cleanUp):
            while(i < len(SourcesToRecord)):
                Frames[i] = Frames[i].pop(0)
                Frames[i] = Frames[i].pop(0)
                Frames[i] = Frames[i].pop(0)
                i += 1
except KeyboardInterrupt:
    i = 0
    while(i < len(AudioChannels)):
        AudioChannels[i].stop_stream()
        AudioChannels[i].close()
        i += 1
    p.terminate()
    print("Quiting system, saved frame over lifetime: ", end="")
    print(samplesI)