#include <Arduino.h>
#include <U8g2lib.h>
#include "stats_font1x1.h"
#include "stats_font2x2.h"
#include "stats_font4x4.h"
#include "glyphs.h"
#include "stats_dto.h"

#ifdef U8X8_HAVE_HW_I2C
#include <Wire.h>
#endif

//#define DEBUG
#ifdef DEBUG
  #define D
#else
  #define D for(;0;)
#endif

#define SCROLL_SPEED 3
#define NO_DATA_TIMEOUT 5 * 1000

// I²C pins        SCL, SDA
#define DISP1_PINS 27,  26
#define DISP2_PINS 25,  33

#define LID_PIN 18


#define USE_ESP32_I2C_HAL
#ifdef USE_ESP32_I2C_HAL
  #include "u8g2_esp32_hal.h"
  #include "u8g2_esp32_hal.c"

  class U8G2_CUSTOM_SSD1306_128X32_UNIVISION_F_ESP32_HAL_I2C : public U8G2 {
    public: U8G2_CUSTOM_SSD1306_128X32_UNIVISION_F_ESP32_HAL_I2C(const u8g2_cb_t *rotation, uint8_t clock, uint8_t data, uint8_t reset = U8X8_PIN_NONE, uint8_t i2c_port = U8X8_PIN_NONE) : U8G2() {
      u8g2_esp32_hal_t esp_hal = U8G2_ESP32_HAL_DEFAULT;
      esp_hal.sda = (gpio_num_t) data;
      esp_hal.scl = (gpio_num_t) clock;
      if (reset != U8X8_PIN_NONE)
        esp_hal.reset = (gpio_num_t) reset;
      if (i2c_port != U8X8_PIN_NONE)
        esp_hal.master_num = (i2c_port_t) i2c_port;

      u8g2_esp32_hal_init((u8x8_t *)(&u8g2), esp_hal);

      u8g2_Setup_ssd1306_i2c_128x32_univision_f(&u8g2, rotation, u8g2_esp32_i2c_byte_cb, u8g2_esp32_gpio_and_delay_cb);
    }
    public: ~U8G2_CUSTOM_SSD1306_128X32_UNIVISION_F_ESP32_HAL_I2C() {
      u8g2_esp32_hal_free((u8x8_t *)(&u8g2));
    }
  };

  U8G2_CUSTOM_SSD1306_128X32_UNIVISION_F_ESP32_HAL_I2C u8g2_d1(U8G2_R0, DISP1_PINS, U8X8_PIN_NONE, 0);
  U8G2_CUSTOM_SSD1306_128X32_UNIVISION_F_ESP32_HAL_I2C u8g2_d2(U8G2_R0, DISP2_PINS, U8X8_PIN_NONE, 1);
#else
  U8G2_SSD1306_128X32_UNIVISION_F_SW_I2C u8g2_d1(U8G2_R0, DISP1_PINS, U8X8_PIN_NONE);
  U8G2_SSD1306_128X32_UNIVISION_F_SW_I2C u8g2_d2(U8G2_R0, DISP2_PINS, U8X8_PIN_NONE);
#endif


u8g2_uint_t f1x1h = 8;
u8g2_uint_t f2x2h = 16;
u8g2_uint_t f4x4h = 32;

char tempLine0Format[] = GLYPH_AMD "%2d" GLYPH_DEGC " "
                   GLYPH_VR_MOS "%2d" GLYPH_DEGC " "
                   GLYPH_GPU "%2d" GLYPH_DEGC " "
                   GLYPH_CHIPSET "%2d" GLYPH_DEGC " "
                   GLYPH_MOBO "%2d" GLYPH_DEGC " "
                   GLYPH_SSD "%2d" GLYPH_DEGC " ";
                   
char tempLine1Format[] = GLYPH_AMD " %s" GLYPH_HZ "  "
                   GLYPH_LOAD " %2d.%02d %4d%%";


char fansLine0Format[] = GLYPH_AIO "%4d" GLYPH_RPM " "
                   GLYPH_GPU "%4d" GLYPH_RPM " "
                   GLYPH_CHIPSET "%4d" GLYPH_RPM " "
                   GLYPH_FAN_TOP "%4d" GLYPH_RPM " "
                   GLYPH_FAN_BOTTOM "%4d" " "
                   GLYPH_FAN_BACK "%4d" GLYPH_RPM " ";

char fansLine1Format[] = GLYPH_RAM " %2d%%  %s  %4dM" GLYPH_HZ;

char waiting_string[] = "     Waiting for data from the computer...    ";

u8g2_uint_t waiting_width;
u8g2_uint_t waiting_offset = 0;
u8g2_uint_t waiting_offset1;

u8g2_uint_t temp_page_offset = 0;
u8g2_uint_t fans_page_offset = 0;

struct line_buffers {
  char line0[300];
  char line1[30];
  char line2[30];
};

// Two buffers to swap around and avoid blocking rendering
struct line_buffers tempLines_1 = {0};
struct line_buffers tempLines_2 = {0};
struct line_buffers fansLines_1 = {0};
struct line_buffers fansLines_2 = {0};

struct line_buffers *tempLines = &tempLines_1;
struct line_buffers *tempLines_next = &tempLines_2;
struct line_buffers *fansLines = &fansLines_1;
struct line_buffers *fansLines_next = &fansLines_2;

SemaphoreHandle_t stringsMutex;

stats_t stats;

byte bytes_read = 0;
bool stats_ever_received = false;
unsigned long data_expiration = 0;

bool lid = 0;


bool receiveStats() {
  size_t to_read = sizeof(stats_t) - bytes_read;
  
  if (Serial.available() < 4) return false;
  if (bytes_read == 0 && !Serial.find(MAGIC_START_ASSTR, 4)) {
    return false;
  } else if (bytes_read == 0) {
    // Put back magic since Serial.find() strips it out
    stats.magic_start = MAGIC_START;
    bytes_read = 4;
    // It also says it read them even though it doesn't place them in the buffer
    to_read += 4;
  }
  
  D Serial.println(F("Magic found, accepting data"));
  
  byte *buffer = (byte *) &stats;
  bytes_read += Serial.readBytes(buffer + bytes_read, to_read);

  D Serial.print(F("Read "));
  D Serial.print(bytes_read);
  D Serial.println(F(" bytes from serial"));

  if (bytes_read >= sizeof(stats_t)) {

#ifdef DEBUG
    char printbuf[100];
    sprintf(printbuf, "Received %d bytes, wanted %d == sizeof(stats_t)", bytes_read, sizeof(stats_t));
    Serial.println(printbuf);
    Serial.print(F("buffer: "));
    for (byte *ptr = buffer; ptr < buffer + sizeof(stats_t); ptr++) {
      sprintf(printbuf, "%02x ", *ptr);
      Serial.print(printbuf);
    }
    Serial.println();
    
    Serial.println(F("Check magic start"));
#endif

    if (stats.magic_start != MAGIC_START) goto invalidate;

    D Serial.println(F("Check magic end"));
    D Serial.println(stats.magic_end, HEX);

    if (stats.magic_end != MAGIC_END) goto invalidate;

    uint8_t checkxor = 0;
    for (byte *b = (byte *) &stats; b < &(stats.checkxor); b++) {
      checkxor ^= *b;
    }

#ifdef DEBUG
    Serial.println(F("Checking check xor"));
    Serial.print(F("Computed: "));
    Serial.print(checkxor, HEX);
    Serial.print(F(" at offset "));
    Serial.println((uint32_t) ((uint32_t) &(stats.checkxor) - ((uint32_t) &stats)));
    Serial.print(F("Got: "));
    Serial.println(stats.checkxor, HEX);
#endif
    
    if (stats.checkxor != checkxor) goto invalidate;

    D Serial.println(F("Checks passed, data accepted"));

    bytes_read = 0;
    return true;
  } else {
    return false;
  }
  
invalidate:
  D Serial.println(F("Invalid data, ignoring"));
  bytes_read = 0;
  return false;
}


void drawWaitingForData(U8G2 *u8g2, u8g2_uint_t *offset) {
  u8g2->clearBuffer();

  u8g2_uint_t x = *offset;
  do {
    u8g2->setCursor(x, 32);
    u8g2->setFont(stats_font4x4);
    u8g2->print(GLYPH_UNPLUGGED);

    u8g2->setCursor(x + 32, 26);
    u8g2->setFont(u8g2_font_helvR18_tf);
    u8g2->print(waiting_string);
    x += waiting_width;
  } while( x < u8g2->getDisplayWidth() );

  yield();
  u8g2->sendBuffer();

  (*offset) -= SCROLL_SPEED;
  if ( (u8g2_uint_t)(*offset) < (u8g2_uint_t)-waiting_width )  
    (*offset) = 0;
}

void drawWaitingForDataAllDisplays() {
  drawWaitingForData(&u8g2_d1, &waiting_offset);
  drawWaitingForData(&u8g2_d2, &waiting_offset1);
}

void format_freq(char *dest, uint16_t value) {
  if (value < 1000) {
    sprintf(dest, "%4dM", value);
  } else {
    sprintf(dest, "%1d.%03dG", value / 1000, value % 1000);
  }
}

void format_bar(
  char *buffer, uint8_t total_tiles, uint16_t full_perc, uint16_t gray_perc, uint16_t perc_max) {

  memset(buffer, 0, total_tiles + 1);
  int32_t gray_tiles = (int32_t) ((double) (gray_perc+1) / (double) perc_max * total_tiles);
  int32_t full_tiles = (int32_t) ((double) (full_perc+1) / (double) perc_max * total_tiles);
  int32_t empty_tiles = total_tiles - 2 - gray_tiles - max(full_tiles - 1, 0);

  // Pick left butt tile
  if (full_tiles == 1 && gray_tiles == 0)
    strcat(buffer, GLYPH_BUTT_LEFT_FULL_END); 
  else if (full_tiles == 0 && gray_tiles == 1)
    strcat(buffer, GLYPH_BUTT_LEFT_GRAY_END);
  else if (full_tiles > 0)
    strcat(buffer, GLYPH_BUTT_LEFT_FULL); 
  else if (gray_tiles > 0)
    strcat(buffer, GLYPH_BUTT_LEFT_GRAY);
  else
    strcat(buffer, GLYPH_BUTT_LEFT_EMPTY);

  // Draw full tiles
  for (char i = 0; i < full_tiles - 2; i++)
    strcat(buffer, GLYPH_MID_FULL);
  
  if (full_tiles > 1 && full_tiles < total_tiles && gray_tiles == 0)
    strcat(buffer, GLYPH_MID_FULL_END);
  else if (full_tiles > 0 && gray_tiles > 0)
    strcat(buffer, GLYPH_MID_FULL);

  // Draw gray tiles
  for (char i = 0; i < gray_tiles - 1 - (full_tiles > 0 ? 0 : 1); i++)
    strcat(buffer, GLYPH_MID_GRAY);
    
  if (gray_tiles > 0 && (full_tiles + gray_tiles) < total_tiles)
    strcat(buffer, GLYPH_MID_GRAY_END);

  // Draw empty tiles
  for (char i = 0; i < empty_tiles; i++)
    strcat(buffer, GLYPH_MID_EMPTY);

  // Pick right butt tile
  if (full_tiles == total_tiles)
    strcat(buffer, GLYPH_BUTT_RIGHT_FULL); 
  else if ((full_tiles + gray_tiles) == total_tiles)
    strcat(buffer, GLYPH_BUTT_RIGHT_GRAY);
  else
    strcat(buffer, GLYPH_BUTT_RIGHT_EMPTY);
}


void format_size(char *buffer, uint32_t megabytes) {
  if (megabytes < 2000) {
    sprintf(buffer, "%4dMB", megabytes);
  } else {
    sprintf(buffer, "%4.2fGB", (float) megabytes / 1024);
  }
}


void drawStatsPage(
  U8G2 *u8g2, struct line_buffers *lines, char *minibanner,
  u8g2_uint_t mb_w, u8g2_uint_t mb_h, u8g2_uint_t *offset) {

  u8g2->setFont(stats_font2x2);
  u8g2_uint_t line0_width = u8g2->getStrWidth(lines->line0);
    
  u8g2->clearBuffer();

  // Line 0
  u8g2_uint_t x = *offset;
  u8g2->setFont(stats_font2x2);
  do {
    u8g2->drawStr(x, f2x2h, lines->line0);
    x += line0_width;
  } while( x < u8g2->getDisplayWidth() );

  // Minibanner
  if (data_expiration > millis()) {
    u8g2->setDrawColor(0);
    u8g2->drawBox(0, 0, mb_w, mb_h);
    u8g2->setDrawColor(1);
    u8g2->drawStr(0, f2x2h, minibanner);
  }

  // Line 1
  u8g2->setFont(stats_font1x1);
  u8g2->drawStr(0, f2x2h + f1x1h + 1, lines->line1);

  // Line 2
  u8g2->drawStr(0, f2x2h + 2*f1x1h, lines->line2);

  // No data maxibanner
  if (data_expiration <= millis()) {
    u8g2->setFont(stats_font4x4);
    u8g2->setDrawColor(0);
    u8g2->drawBox(0, 0, 36, 32);
    u8g2->setDrawColor(1);
    u8g2->drawStr(0, f4x4h, GLYPH_UNPLUGGED);
  }

  yield();

  u8g2->sendBuffer();
  
  (*offset) -= SCROLL_SPEED;
  if ( (u8g2_uint_t) (*offset) < (u8g2_uint_t) -line0_width )  
    (*offset) = 0;
}


void sprintfStrings() {
  uint8_t total_8x8_tiles = u8g2_d1.getDisplayWidth() / 8;
  char tmp_buf[20] = {0};

  // Temps page
  sprintf(tempLines_next->line0, tempLine0Format,
    stats.cpu_temp,
    stats.vr_mos_temp,
    stats.gpu_temp,
    stats.chipset_temp,
    stats.system_temp,
    stats.ssd_temp);

  format_freq(tmp_buf, stats.cpu_freq);
  sprintf(tempLines_next->line1, tempLine1Format,
    tmp_buf, stats.cpu_load_avg / 100, stats.cpu_load_avg % 100, stats.cpu_perc
  );

  format_bar(
    tempLines_next->line2, total_8x8_tiles,
    stats.cpu_perc - stats.cpu_perc_kernel, stats.cpu_perc_kernel, stats.cpu_perc_max);


  // Fans page
  sprintf(fansLines_next->line0, fansLine0Format,
    stats.pump,
    stats.gpu_fan,
    stats.chipset_fan,
    stats.top_fan,
    stats.bottom_fan,
    stats.back_fan);

  format_size(tmp_buf, stats.ram_used);
  sprintf(fansLines_next->line1, fansLine1Format, stats.ram_perc, tmp_buf, 3200);

  format_bar(
    fansLines_next->line2, total_8x8_tiles, stats.ram_perc, stats.ram_perc_buffers, 100);

  struct line_buffers *temp_ptr = tempLines;
  struct line_buffers *fans_ptr = fansLines;

  xSemaphoreTake(stringsMutex, portMAX_DELAY);
  tempLines = tempLines_next;
  fansLines = fansLines_next;
  xSemaphoreGive(stringsMutex);

  tempLines_next = temp_ptr;
  fansLines_next = fans_ptr;
}



void displaysTask(void *pvParameters) {
  Serial.print("Handling displays on core ");
  Serial.println(xPortGetCoreID());

  while (true) {
    bool newLid = digitalRead(LID_PIN);
    if (newLid != lid) {
      lid = newLid;
      u8g2_d1.setPowerSave(!lid);
      u8g2_d2.setPowerSave(!lid);
      Serial.print(F("LID "));
      Serial.println(lid ? F("OPEN") : F("CLOSED"));
    }
    
    if (!stats_ever_received) {
      drawWaitingForDataAllDisplays();
      yield();
      continue;
    }

    xSemaphoreTake(stringsMutex, portMAX_DELAY);
    struct line_buffers *templines_ptr = tempLines;
    struct line_buffers *fanslines_ptr = fansLines;
    xSemaphoreGive(stringsMutex);

    drawStatsPage(&u8g2_d1, templines_ptr, GLYPH_TEMP, 10, 32, &temp_page_offset);
    drawStatsPage(&u8g2_d2, fanslines_ptr, GLYPH_FAN, 16, 32, &fans_page_offset);

    yield();
  }
}

void receiverTask(void *pvParameters) {
  D Serial.print("Handling incoming data on core ");
  D Serial.println(xPortGetCoreID());

  while (true) {
    yield();

    if (receiveStats()) {
      sprintfStrings();
      data_expiration = millis() + NO_DATA_TIMEOUT;
      stats_ever_received = true;
    }

    delay(10);
  }
}


void setup() {
  // Fuck this shit and fuck your "wifi" and "important stuff"
  disableCore0WDT();

  Serial.begin(115200);
  Serial.setTimeout(0);

  // Report initial lid state
  pinMode(LID_PIN, INPUT_PULLUP);
  lid = digitalRead(LID_PIN);
  Serial.print(F("LID "));
  Serial.println(lid ? F("OPEN") : F("CLOSED"));

  stringsMutex = xSemaphoreCreateMutex();

#ifdef USE_ESP32_I2C_HAL
  u8g2_d1.setI2CAddress(0x78);
  u8g2_d2.setI2CAddress(0x78);
#endif
  
  u8g2_d1.begin();
  u8g2_d1.setPowerSave(0);
  u8g2_d1.setFontMode(0);
  u8g2_d1.setAutoPageClear(0);

  u8g2_d1.setFont(u8g2_font_helvR18_tf);
  waiting_width = u8g2_d1.getStrWidth(waiting_string);

  u8g2_d1.setFont(stats_font4x4);
  waiting_width += u8g2_d1.getStrWidth(GLYPH_UNPLUGGED);

  waiting_offset = 0;

  u8g2_d2.begin();
  u8g2_d2.setPowerSave(0);
  u8g2_d2.setFontMode(0);
  u8g2_d2.setAutoPageClear(0);

  u8g2_d2.setFont(u8g2_font_helvR18_tf);
  waiting_offset1 = (u8g2_uint_t) -u8g2_d2.getDisplayWidth();

  delay(500);

  xTaskCreatePinnedToCore(displaysTask, "DisplaysTask", 10000, NULL, 4, NULL, 0);
  xTaskCreatePinnedToCore(receiverTask, "ReceiverTask", 10000, NULL, 4, NULL, 1);
}

void loop() {
  // Avoid task watchdog firing
  vTaskDelete(NULL);
}