#include <WiFi.h>
#include <WiFiUdp.h>

#define DEBUG

// #define MBS2_1 // AP hosting ESP32-AP2
#define MBS2_2 // Client connected to ESP32-AP1

#ifdef MBS2_1
  #ifdef MBS2_2
   #undef MBS2_2
  #endif
#endif
#ifdef MBS2_2
  #ifdef MBS2_1
   #undef MBS2_1
  #endif
#endif

#ifdef MBS2_1
#define _NAME "MBS2_1"
#define ap_ssid "ESP32-AP2"
#define ap_passwd "PASS0000"
IPAddress apIP(192, 168, 178, 1);
IPAddress subnetMask(255, 255, 255, 0);
#endif

#ifdef MBS2_2
#define _NAME "MBS2_2"
#define base_ssid "ESP32-AP1"
#define base_passwd "PASS0000"
#endif

#define MAX_PACKET_SIZE 1500
#define UDP_PORT 6666

WiFiUDP udp;

IPAddress localIP = IPAddress();

void (*routeMBS)(void);

#pragma region Access Point Functions

#ifdef MBS2_1
void init_wifi_access_point(){
  WiFi.mode(WIFI_AP);
  WiFi.softAPConfig(apIP, apIP, subnetMask);
  WiFi.softAP(ap_ssid, ap_passwd, /* channel */ 1, /* hidden */ false, /* maxConn */ 4);
}

char *receive(){
  int packetSize = udp.parsePacket();
  if(packetSize <= 0) return NULL;
  packetSize = packetSize > MAX_PACKET_SIZE ? MAX_PACKET_SIZE: packetSize;
  char *buffer = (char*)malloc(packetSize + 1);

  udp.read(buffer, packetSize);
  buffer[packetSize] = '\0';

  return buffer;
}

void send(char *buf){
  Serial.write(buf);
  Serial2.write(buf);
}

String readSerialUntil(const String& delimiter) {
  String res = "";
  unsigned long start = millis();

  while (!res.endsWith(delimiter) && (millis() - start < 1000)) {
    if (Serial.available()) {
      char c = Serial.read();
      res += c;
      start = millis();
    } else {
      // yield the CPU so the watchdog won’t fire
      taskYIELD();
    }
  }

  return res;
}

void doSerialReading(){
  if(Serial.available()){
    Serial2.print(readSerialUntil("\n"));
  }
}

void MBS2_1_loop(){
  #ifdef DEBUG
  /*
  To debug loop through instances
  MBS2_1 Serial2 out -> MBS2_1 Serial2 in UDP out
  */
  doSerialReading();
  #endif
  char *recv = receive();
  if(recv != NULL){
    log(String("Received: "));
    log(recv);
    send(recv);
    free(recv);
    recv = NULL;
  }
  free(recv); // NULL free is safe
  vTaskDelay(pdMS_TO_TICKS(1));
}
#endif // #ifdef MBS2_1

#pragma endregion

#pragma region Client (Base Station) Functions

#ifdef MBS2_2
IPAddress apIP(192,168,178,1);

void init_wifi_base_station_connection(){
  WiFi.mode(WIFI_STA);
  WiFi.begin(base_ssid, base_passwd);

  #ifdef DEBUG
  size_t start = millis();
  #endif
  while(WiFi.status() != WL_CONNECTED){ // while esps aren't connected nothing can run either
    #ifdef DEBUG
      log(String("Trying to connect to: ") + String(base_ssid) + String("\nTime: ") + String(millis() - start));
      Serial.flush();
    #endif
    // Serial.println("Trying to connect to: " + this->base_ssid + "\nTime: " + String(millis() - start) + "/" + String(connection_timeout));
    delay(50);
  }

  if(WiFi.status() == WL_CONNECTED){
    localIP = WiFi.localIP();
  }
}

void send(const String& msg){
  udp.beginPacket(apIP, UDP_PORT);
  udp.write((uint8_t*)msg.c_str(), msg.length());
  udp.endPacket();
}

String readSerialUntil(const String& delimiter) {
  String res = "";
  unsigned long start = millis();

  while (!res.endsWith(delimiter) && (millis() - start < 1000)) {
    if (Serial2.available()) {
      char c = Serial2.read();
      res += c;
      start = millis();
    } else {
      // yield the CPU so the watchdog won’t fire
      taskYIELD();
    }
  }

  return res;
}

void MBS2_2_loop(){
  if(Serial2.available()){
    String recv = readSerialUntil("\n");
    log("Received: " + recv);
    send(recv);
  }
  // always give os a tick
  vTaskDelay(pdMS_TO_TICKS(1));
}
#endif // #ifdef MBS2_2

#pragma endregion

void log(String buf){
  #ifdef DEBUG
  Serial.println(buf);
  #endif
}

void log(char *buf){
  #ifdef DEBUG
  Serial.write(buf);
  #endif
}

void setup() {
  #ifdef DEBUG
  Serial.begin(115200);
  #endif
  Serial2.begin(115200, SERIAL_8N1, 18, 17);
  
  #ifdef MBS2_1
  init_wifi_access_point();
  routeMBS = MBS2_1_loop;
  #endif

  #ifdef MBS2_2
  init_wifi_base_station_connection();
  routeMBS = MBS2_2_loop;
  #endif
  
  udp.begin(UDP_PORT);
}

void loop() {
  #ifdef DEBUG
  static size_t start = millis();
  if(millis() - start >= 1000){
    log(String(_NAME));
    start = millis();
  }
  #endif
  routeMBS();
  taskYIELD();
}