#include <WiFi.h>
#include <WiFiUdp.h>
#define UDP_TX_PACKET_MAX_SIZE 1024

const char * networkName = "foo";
const char * networkPswd = "bar";

#define INCLUDE_vTaskSuspend 1

WiFiUDP Udp;
const int udpPort = 8888;

char packetBuffer[UDP_TX_PACKET_MAX_SIZE];
#define LED_PER_PACKET 340
#define INCLUDE_vTaskSuspend 1

IPAddress remote_address;
int remote_port = NULL;
int led_number = NULL;

unsigned short int frame_number;
unsigned short int packet_number;

portMUX_TYPE myMutex = portMUX_INITIALIZER_UNLOCKED;

#include <dcll.h>
Dcll *dcll;

TaskHandle_t Receive_Udp_Task;

void write_ack(){
    if(remote_address && remote_port){
        Serial.println("WRITE ACK");
        Udp.beginPacket(remote_address, remote_port);
        Udp.print("ACK");
        Udp.endPacket();
    }
}

void write_frame(){
    Udp.beginPacket(remote_address, remote_port);
    Udp.print("FRAME");
    Udp.endPacket();
}

void handle_registration_packet(char* packetBuffer){
  led_number = *((unsigned int*)(packetBuffer+1));
  remote_address = Udp.remoteIP();
  remote_port = Udp.remotePort();
}

void handle_frame_packet(char* packetBuffer){
  write_ack();
  frame_number = (unsigned short int)*((char*)(packetBuffer+1));
  packet_number = (unsigned short int)*((char*)(packetBuffer+2));

  int led_offset = 340 * packet_number;
  int r_off = 3;
  for(int i=0;i<340;i++){
    if(led_offset+i < led_number){
      dcll->leds[led_offset+i].setRGB(
        packetBuffer[r_off],
        packetBuffer[r_off+1],
        packetBuffer[r_off+2]
      );
    }

    r_off = r_off + 3;
  }

  if(packet_number == 0){
  }else if(packet_number == ceil(led_number / LED_PER_PACKET)){
    dcll->swap_buffer();
    write_frame();
  }
}

void setup() {
  Serial.begin(115200);
  // put your setup code here, to run once:
  WiFi.begin(networkName, networkPswd);
  while (WiFi.status() != WL_CONNECTED) {
    delay(500);
    Serial.print(".");
  }
  // Print local IP address and start web server
  Serial.println("");
  Serial.println("WiFi connected.");
  Serial.println("IP address: ");
  Serial.println(WiFi.localIP());

  Udp.begin(udpPort);

  //delay(1000);
  dcll = new Dcll(1);

  xTaskCreatePinnedToCore(
    Receive_Udp_Task_Code,
    "Receive_Udp_Task",
    1000,
    NULL,
    0,
    &Receive_Udp_Task,
    0
  );
}

void Receive_Udp_Task_Code(void * parameter) {
  while(true){
    //Serial.println("RECEIVE UDP TASK");
    int packetSize = Udp.parsePacket();

    if(packetSize){
      //Serial.println("GOT PACKET");
      Udp.read(packetBuffer, packetSize);

      vTaskDelay(40 / portTICK_PERIOD_MS);

      if(packetBuffer[0] == 1){
        handle_registration_packet(packetBuffer);
      }else if(packetBuffer[0] == 2){
        Serial.println("HANDLE PACKET FRAME");
        handle_frame_packet(packetBuffer);
      }

      /*
      if(packetBuffer[0] == 1){
        handle_registration_packet(packetBuffer);
      }else if(packetBuffer[0] == 2){
        handle_frame_packet(packetBuffer);
      }
      */
      write_ack();
      vTaskDelay(40 / portTICK_PERIOD_MS);
    }
  }
}


void loop() {
  /*
  int packetSize = Udp.parsePacket();

  if(packetSize){
    Udp.read(packetBuffer, packetSize);

    if(packetBuffer[0] == 1){
      handle_registration_packet(packetBuffer);
    }else if(packetBuffer[0] == 2){
      handle_frame_packet(packetBuffer);
    }
  }
  */
  delay(100000);
}