/*********
  ESP32 Serial LED Controller
*********/

#define WOKWI

// FastLED setup ----------------------
#include <FastLED.h>

#ifdef WOKWI
#define LED_TYPE WS2812B
#define COLOR_ORDER GRB
#define NUM_LEDS 24
#define START_DELAY 500
#define BOARD_NAME "WOKWI SIM"
#define BAUD_RATE 230400
#else
#define LED_TYPE WS2813
#define COLOR_ORDER RGB
#define NUM_LEDS 300
#define START_DELAY 3000
#define BOARD_NAME "ESP32"
#define BAUD_RATE 230400
#endif

#define DATA_PIN 12
CRGB leds[NUM_LEDS];

#define BRIGHTNESS 96
#define FRAMES_PER_SECOND 60
// End FastLED setup ------------------

#define onboardLED 2

struct specialBytes {
  byte cmd = 0x24;    //36  0x24 $
  byte start = 0x3C;  //60  0x3C <
  byte end = 0x3E;    //62  0x3E >
  byte array = 0x41;  //65  0x41 A
  byte arrayC = 0x61; //97  0x61 a
  byte arrayC2 = 0x62;//98  0x62 b
  byte led = 0x4C;    //76  0x4C L
  byte ledC = 0x6C;   //108 0x6C l
  byte ledC2 = 0x6D;  //109 0x6D m
  byte bright = 0x42; //66  0x42 B
  byte update = 0x55; //85  0x55 U
} specByte;

const int maxBytes = 1024;
byte receivedBytes[maxBytes];
int failedMsgs = 0;
bool bCheckSum = true;

void setup() {
  pinMode(onboardLED, OUTPUT);

  Serial.begin(BAUD_RATE);
  Serial.print("<");
  Serial.print(BOARD_NAME);
  Serial.println(" is ready>");

  // FastLED setup ----------------------
  delay(START_DELAY);  // 3 second delay for recovery
  
  // tell FastLED about the LED strip configuration
  FastLED.addLeds<LED_TYPE, DATA_PIN, COLOR_ORDER>(leds, NUM_LEDS).setCorrection(TypicalLEDStrip);
  // set master brightness control
  FastLED.setBrightness(BRIGHTNESS);
}

void loop() {
  //digitalWrite(onboardLED, HIGH);
  //delay(500);
  //digitalWrite(onboardLED, LOW);
  //delay(500);

  // insert a delay to keep the framerate modest
  //FastLED.delay(1000 / FRAMES_PER_SECOND);

  // do some periodic updates
  //EVERY_N_MILLISECONDS( 20 ) { gHue++; } // slowly cycle the "base color" through the rainbow
  //EVERY_N_SECONDS( 10 ) { nextPattern(); } // change patterns periodically

  receiveBytes();
}

void receiveBytes() {
  static bool recvInProgress = false;
  static int ndx = 0;
  static bool bGetCmd = false;
  byte rb;
  
  while (Serial.available() > 0) {
    rb = Serial.read();
    //Serial.print("Byte: "); Serial.println(rb);

    if (rb == specByte.cmd) {
      //Serial.println("Got command byte");
      if (bGetCmd) {
        //Serial.println("escape to extract cmdByte value");
        bGetCmd = false;

        receivedBytes[ndx] = rb; // duplicated code, TODO: make function and use pointer
        ndx++;
        if (ndx >= maxBytes) ndx = maxBytes - 1;
      }
      else bGetCmd = true;
    }
    else if (bGetCmd) {
      bGetCmd = false;

      if (rb == specByte.start) {
        //Serial.println("START transmission");
        if (recvInProgress) {
          Serial.println("Previous transmission incomplete, discard and reset");
          ndx = 0;
        }
        else recvInProgress = true;
      }
      else if (rb == specByte.end) {
        //Serial.println("END transmission");
        receivedBytes[ndx] = '\0'; // terminate the string TODO: remove
        processData(ndx);

        recvInProgress = false;
        ndx = 0;

        return; // exit to start fresh with any new data
      }
    }
    else if (recvInProgress) {
      receivedBytes[ndx] = rb;
      ndx++;
      if (ndx >= maxBytes) ndx = maxBytes - 1;
    }

  }
}

byte checkSum(byte* msg, size_t len) {
  int sum = 0;

  for (int i = 0; i < len; i++) {
    sum += msg[i];
  }
  
  return sum & 255;
}

void processData(int len) {
  if (bCheckSum && checkSum(&receivedBytes[0], len-1) != receivedBytes[len-1]) {
    failedMsgs++;
    Serial.print("CHECKSUM FAILED, discarding message. #failed: ");
    Serial.println(failedMsgs);
    return;
  }

  byte header = receivedBytes[0];
  byte* msg = &receivedBytes[1];
  int msgLen = len - 2; // subtract header and checksum

  // update whole array
  if (header == specByte.array || header == specByte.arrayC || header == specByte.arrayC2) {
    byte comp = (header == specByte.arrayC) ? 1 : 0;
    if (header == specByte.arrayC2) comp = 2;

    updateArray(msg, msgLen, comp);
    updateStrip();
  }

  // update individual leds
  if (header == specByte.led || header == specByte.ledC || header == specByte.ledC2) {
    byte comp = (header == specByte.ledC) ? 1 : 0;
    if (header == specByte.ledC2) comp = 2;

    updateLed(msg, msgLen, comp);
    updateStrip();
  }

  if (header == specByte.bright) {
    changeBrightness(msg[0]);
    updateStrip();
  }
  if (header == specByte.update) {
    updateStrip();
  }
}
// --------------------------------
// LED FUNCTIONS
// --------------------------------
void changeColor(int r, int g, int b) {
  for (int i = 0; i < NUM_LEDS; i++) {
    leds[i].r = r;
    leds[i].g = g;
    leds[i].b = b;
  }
}

void updateLed(byte* data, size_t len, byte comp) {
  //comp0 = 5 bytes (40bits) 16bits for idx, 8bits each for r,g,b
  //comp1 = 4 bytes (32bits i9 r8 g8 b7) - idx up to 512 (7bits = 0-127)
  //comp2 = 4 bytes (32bits i10 r7 g8 b7) - idx up to 1024
  byte pcktLen = (comp == 0) ? 5 : 4;
  int numLeds = len / pcktLen;
  byte r, g, b;
  int idx;

  for (int i = 0; i < numLeds; i++) {
    int startIdx = i * pcktLen;

    if (comp == 0) {
      idx = data[startIdx+1];
      if (data[startIdx] > 0) idx += data[startIdx] * 256; // combine 2 idx bytes

      r = data[startIdx+2];
      g = data[startIdx+3];
      b = data[startIdx+4];
    }
    else {
      idx = data[startIdx]; // compressed packed data
      r = data[startIdx+1]; //comp1 uncompressed, comp2 compressed packed data
      g = data[startIdx+2]; // always uncompressed
      b = data[startIdx+3]; // compressed packed data
      
      // extract idx & b, extract r if comp == 2 (otherwise already uncompressed) 
      decompress(&idx, &b, (comp == 2) ? &r : nullptr);
    }

    leds[idx].r = r;
    leds[idx].g = g;
    leds[idx].b = b;
    
    //Serial.print("LED: "); Serial.print(idx); Serial.print(" - R: "); Serial.print(leds[idx].r);
    //Serial.print(" G: "); Serial.print(leds[idx].g); Serial.print(" B: "); Serial.println(leds[idx].r);
  }
}

void updateArray(byte* data, size_t len, byte comp) {
  if (len == 3) {
    // TODO: this shortcut doesn't handle differet compressions
    changeColor(data[0], data[1], data[2]);
  }
  else {
    // comp0 = RGB data comes in 3bytes
    // comp1 = RGB data comes in 2bytes
    // comp2 = RGB data comes in 1byte
    byte pcktLen = 3 - comp; // just a coincidence this works (should use a LUT)
    int numLeds = len / pcktLen;
    byte r, g, b;

    for (int i = 0; i < numLeds; i++) {
      int startIdx = i * pcktLen;

      // get 3rgb bytes from 3, 2, or 1 byte container
      if (comp == 0) fromRGB888(&data[startIdx], &r, &g, &b);
      else if (comp == 1) fromRGB565(&data[startIdx], &r, &g, &b);
      else if (comp == 2) fromRGB332(&data[startIdx], &r, &g, &b);

      leds[i].r = r;
      leds[i].g = g;
      leds[i].b = b;

      //Serial.print("C"); Serial.print(comp); Serial.print("|");
      //Serial.print("R: "); Serial.print(r); Serial.print(" B: "); Serial.print(b); Serial.print(" G: "); Serial.println(g);
    }
  }
}

void changeBrightness(byte newBright) {
  //Serial.print("Brightness: "); Serial.println(newBright);

  FastLED.setBrightness(newBright);
}

void updateStrip() {
  FastLED.show();
  //Serial.println("Updating Strip");
}

// --------------------------------
// DECOMPRESSION & UNPACKING
// --------------------------------
const double compFactors[] = { 1, 255, 85, 36.428, 17, 8.225, 4.047, 2.0078 };
const byte bitMasks[] = { 128, 64, 32, 16, 8, 4, 2, 1 };

byte expandBits(byte val, byte numBits) {
  return (byte)round((double)val * compFactors[numBits]);
}

byte extractBits(byte packedByte, byte start, byte numBits) {
  byte mask = 0;
  for (int i = start; i < start + numBits; i++) {
    mask |= bitMasks[i];
  }
  return (packedByte & mask) >> (8 - (start + numBits));
}

void decompress(int* idx, byte* b, byte* r) {
  if ((*b & 128) != 0) *idx |= 256; // if B8 set, set idx9 bit
  *b = expandBits(*b & ~128, 7); // clear 8th bit & decompress

  if (r != nullptr) {
    if ((*r & 128) != 0) *idx |= 512; // if R8 set, set idx10 bit
    *r = expandBits(*r & ~128, 7); // clear 8th bit & decompress
  }
}

// RGB uncompressed 3bytes
void fromRGB888(byte* msg, byte* r, byte* g, byte* b) {
  *r = msg[0];
  *g = msg[1];
  *b = msg[2];
}
// RGB compressed to 2bytes
void fromRGB565(byte* msg, byte* r, byte* g, byte* b) {
  *r = expandBits(extractBits(msg[0], 0, 5), 5);
  *b = expandBits(extractBits(msg[1], 0, 5), 5);

  byte g1 = extractBits(msg[0], 5, 3) << 3;
  byte g2 = extractBits(msg[1], 5, 3);
  *g = expandBits(g1 | g2, 6); // decompress
}
// RGB compressed to 1byte
void fromRGB332(byte* msg, byte* r, byte* g, byte* b) {
  *r = expandBits(extractBits(msg[0], 0, 3), 3);
  *g = expandBits(extractBits(msg[0], 3, 3), 3);
  *b = expandBits(extractBits(msg[0], 6, 2), 2);
}