//  Audio Echo and OLED Display Demo 
//    Record from INMP441 I2S Microphone, and
//    Playback to MAX98357A I2S Audio Amplifier
//
//  Targeted to ESP Edge Hardware Platform from HackerBox 0126:
//  https://hackerboxes.com/products/hackerbox-0126-imitation-game
//
//  Audio record and playback approach was adapted from:
//  https://github.com/derdacavga/esp32-voice-recorder/tree/main

#include <Wire.h>
#include <Adafruit_GFX.h>
#include <Adafruit_SSD1306.h>
#include "driver/i2s.h"

#define BUTTON_PIN 3

// INMP441 config
#define I2S_BCLK_MIC 47
#define I2S_LRC_MIC 41
#define I2S_DIN_MIC 14

// MAX98357A config
#define I2S_BCLK_SPK 48
#define I2S_LRC_SPK 21
#define I2S_DOUT_SPK 42

// OLED config
#define SCREEN_WIDTH 128
#define SCREEN_HEIGHT 64
#define OLED_RESET    -1
#define I2C_SDA 2
#define I2C_SCL 1
Adafruit_SSD1306 display(SCREEN_WIDTH, SCREEN_HEIGHT, &Wire, OLED_RESET);

// Audio config
const int SAMPLE_RATE = 16000;
const int BITS_PER_SAMPLE = 16;
const int RECORD_DURATION_SEC = 5; 
const int AUDIO_BUFFER_SIZE = RECORD_DURATION_SEC * SAMPLE_RATE * (BITS_PER_SAMPLE / 8);
int16_t *audio_buffer = NULL;

void record_and_playback();

void setup() {
  Serial.begin(115200);
  delay(1000);

  pinMode(BUTTON_PIN, INPUT_PULLDOWN);

  Wire.begin(I2C_SDA, I2C_SCL);
  if (!display.begin(SSD1306_SWITCHCAPVCC, 0x3C)) {
    Serial.println("OLED init failed!");
    while (1);
  }

  displayMessage("Press to record");

  Serial.println("\n\n");
  Serial.println("Microphone Test");
  Serial.println("Press button to record");

  if (!psramFound()) { 
    Serial.println("PSRAM ERROR!");
    while (1)
      ;
  }

  i2s_config_t i2s_rx_config = {
    .mode = (i2s_mode_t)(I2S_MODE_MASTER | I2S_MODE_RX),
    .sample_rate = SAMPLE_RATE,
    .bits_per_sample = (i2s_bits_per_sample_t)BITS_PER_SAMPLE,
    .channel_format = I2S_CHANNEL_FMT_ONLY_RIGHT,
    .communication_format = I2S_COMM_FORMAT_I2S,
    .intr_alloc_flags = ESP_INTR_FLAG_LEVEL1,
    .dma_buf_count = 8,
    .dma_buf_len = 1024
  };
  i2s_pin_config_t i2s_rx_pin_config = {
    .bck_io_num = I2S_BCLK_MIC,
    .ws_io_num = I2S_LRC_MIC,
    .data_out_num = I2S_PIN_NO_CHANGE,
    .data_in_num = I2S_DIN_MIC
  };
  i2s_driver_install(I2S_NUM_0, &i2s_rx_config, 0, NULL);
  i2s_set_pin(I2S_NUM_0, &i2s_rx_pin_config);

  i2s_config_t i2s_tx_config = {
    .mode = (i2s_mode_t)(I2S_MODE_MASTER | I2S_MODE_TX),
    .sample_rate = SAMPLE_RATE,
    .bits_per_sample = (i2s_bits_per_sample_t)BITS_PER_SAMPLE,
    .channel_format = I2S_CHANNEL_FMT_ONLY_RIGHT,
    .communication_format = I2S_COMM_FORMAT_I2S,
    .intr_alloc_flags = ESP_INTR_FLAG_LEVEL1,
    .dma_buf_count = 8,
    .dma_buf_len = 1024
  };
  i2s_pin_config_t i2s_tx_pin_config = {
    .bck_io_num = I2S_BCLK_SPK,
    .ws_io_num = I2S_LRC_SPK,
    .data_out_num = I2S_DOUT_SPK,
    .data_in_num = I2S_PIN_NO_CHANGE
  };
  i2s_driver_install(I2S_NUM_1, &i2s_tx_config, 0, NULL);
  i2s_set_pin(I2S_NUM_1, &i2s_tx_pin_config);
}

void loop() {
  if (digitalRead(BUTTON_PIN) == HIGH) {
    delay(50);
    while (digitalRead(BUTTON_PIN) == HIGH)
      ;
    delay(100);
    record_and_playback();
  }
}

void displayMessage(String msg) {
  display.clearDisplay();
  display.setTextSize(1);
  display.setTextColor(SSD1306_WHITE);
  display.setCursor(0, 0);
  display.println(msg);
  display.display();
}

void record_and_playback() {
  Serial.println("Allocating memory\nin PSRAM...");

  audio_buffer = (int16_t *)heap_caps_malloc(AUDIO_BUFFER_SIZE, MALLOC_CAP_SPIRAM);
  if (audio_buffer == NULL) {
    Serial.println("ERROR: Could not allocate memory in PSRAM");
    return;
  }
  Serial.println("Buffer created in PSRAM");

  displayMessage("Recording...");
  Serial.printf("Recording for %d seconds...\n", RECORD_DURATION_SEC);

  delay(250);

  size_t bytes_read;
  i2s_read(I2S_NUM_0, audio_buffer, AUDIO_BUFFER_SIZE, &bytes_read, portMAX_DELAY);

  Serial.printf("%d bytes of audio data recorded\n", bytes_read);

  Serial.println("Playback recorded audio...");
  displayMessage("Playback...");

  size_t bytes_written;
  delay(20);
  i2s_write(I2S_NUM_1, audio_buffer, bytes_read, &bytes_written, portMAX_DELAY);

  Serial.printf("%d bytes of audio data played back\n", bytes_written);

  free(audio_buffer);
  audio_buffer = NULL;
  Serial.println("Test Finished");
  Serial.println("\nPress button to record again\n");
  displayMessage("Press to repeat");
}