#include <SPI.h>
#include <TFT_eSPI.h>
#include "ESP32_NOW.h"
#include "WiFi.h"
#include <esp_mac.h>  // For the MAC2STR and MACSTR macros
#include <vector>
#include <JPEGDecoder.h>
#include <SD.h>
#include <FS.h>

// Image paths and corresponding values
const char *images1[] = {"/images/05_smiley.jpg", "/images/04_smiley.jpg", "/images/03_smiley.jpg", "/images/02_smiley.jpg", "/images/01_smiley.jpg"};
const int values1[] = {5, 4, 3, 2, 1};

const char *images2[] = {"/images/50_smiley.jpg", "/images/40_smiley.jpg", "/images/30_smiley.jpg", "/images/20_smiley.jpg", "/images/10_smiley.jpg"};
const int values2[] = {50, 40, 30, 20, 10};

struct ImageData {
  const char *src;
  int value;
};

ImageData allImages[] = {
    {"/images/smiley_mix11.jpg", 11}, {"/images/smiley_mix12.jpg", 12}, {"/images/smiley_mix13.jpg", 13},
    {"/images/smiley_mix14.jpg", 14}, {"/images/smiley_mix15.jpg", 15}, {"/images/smiley_mix21.jpg", 21},
    {"/images/smiley_mix22.jpg", 22}, {"/images/smiley_mix23.jpg", 23}, {"/images/smiley_mix24.jpg", 24},
    {"/images/smiley_mix25.jpg", 25}, {"/images/smiley_mix31.jpg", 31}, {"/images/smiley_mix32.jpg", 32},
    {"/images/smiley_mix33.jpg", 33}, {"/images/smiley_mix34.jpg", 34}, {"/images/smiley_mix35.jpg", 35},
    {"/images/smiley_mix41.jpg", 41}, {"/images/smiley_mix42.jpg", 42}, {"/images/smiley_mix43.jpg", 43},
    {"/images/smiley_mix44.jpg", 44}, {"/images/smiley_mix45.jpg", 45}, {"/images/smiley_mix51.jpg", 51},
    {"/images/smiley_mix52.jpg", 52}, {"/images/smiley_mix53.jpg", 53}, {"/images/smiley_mix54.jpg", 54},
    {"/images/smiley_mix55.jpg", 55},
};

int currentIndex1 = 2;
int currentIndex2 = 2;
int currentIndexAll = 0; // Initialize the index to 0

const int numImages1 = sizeof(images1) / sizeof(images1[0]);
const int numImages2 = sizeof(images2) / sizeof(images2[0]);
const int numAllImages = sizeof(allImages) / sizeof(allImages[0]);

TFT_eSPI tft = TFT_eSPI();

/* Definitions */
#define SCREEN_WIDTH 320
#define SCREEN_HEIGHT 240
#define FONT_SIZE 4

#define ESPNOW_WIFI_CHANNEL 6

bool updateDisplay = false; // Global flag to indicate display update
bool updateResult = false; // Global flag to indicate display update

/* Classes */

// Creating a new class that inherits from the ESP_NOW_Peer class is required.

class ESP_NOW_Peer_Class : public ESP_NOW_Peer {
public:
  // Constructor of the class
  ESP_NOW_Peer_Class(const uint8_t *mac_addr,
                     uint8_t channel,
                     wifi_interface_t iface,
                     const uint8_t *lmk)
    : ESP_NOW_Peer(mac_addr, channel, iface, lmk) {}

  // Destructor of the class
  ~ESP_NOW_Peer_Class() {}

  // Function to register the control peer
  bool add_peer() {
    if (!add()) {
      log_e("Failed to register the broadcast peer");
      return false;
    }
    return true;
  }

  // Function to handle received messages from the control
  void onReceive(const uint8_t *data, size_t len, bool broadcast) {
    Serial.printf("Received a message from control " MACSTR " (%s)\n", MAC2STR(addr()), broadcast ? "broadcast" : "unicast");
    Serial.printf("  Message: %.*s\n", len, data);  // Print the message as received

    // Try to extract the numeric part from the message
    String message = String((char *)data);
    int numberPart = message.substring(message.lastIndexOf(' ') + 1).toInt();

    if (message.startsWith("clockwise")) {
      if (numberPart >= 1 && numberPart <= 5) {
        currentIndex1 = constrain(numberPart - 1, 0, numImages1 - 1); // Update currentIndex1 based on rotary count
        Serial.printf("Updated currentIndex1 to %d\n", currentIndex1);
        updateDisplay = true; // Set flag to update display1 in the main loop
      } else {
        Serial.println("Invalid message format or value.");
      }
    } else if (message.startsWith("counter-clockwise")) {
      if (numberPart >= 1 && numberPart <= 5) {
        currentIndex1 = constrain(numberPart - 1, 0, numImages1 - 1); // Update currentIndex1 based on rotary count
        Serial.printf("Updated currentIndex1 to %d\n", currentIndex1);
        updateDisplay = true; // Set flag to update display1 in the main loop
      } else {
        Serial.println("Invalid message format or value.");
      }
    } else if (message.startsWith("Rockers")) {
      if (numberPart >= 1 && numberPart <= 5) {
        currentIndex2 = constrain(numberPart - 1, 0, numImages2 - 1); // Update currentIndex2 based on rotary count
        Serial.printf("Updated currentIndex2 to %d\n", currentIndex2);
        updateDisplay = true; // Set flag to updateDisplay in the main loop
      } else {
        Serial.println("Invalid message format or value.");
      } 
     } else if (message.equals("SW: pressed")) {
      int sum = values1[currentIndex1] + values2[currentIndex2];
      for (int i = 0; i < numAllImages; ++i) {
        if (allImages[i].value == sum) {
          // Update the currentIndexAll to the index of the matching image
          currentIndexAll = i;  
          Serial.printf("Updated currentIndexAll to %d\n", currentIndexAll);
          updateResult = true; // Set flag to updateResult in the main loop
        }
      }
    } else {
      Serial.println("Invalid message format or value.");
    }
  }
};

/* Global Variables */

// List of all the controls. It will be populated when a new control is registered
std::vector<ESP_NOW_Peer_Class> masters;

/* Callbacks */

// Callback called when an unknown peer sends a message
void register_new_master(const esp_now_recv_info_t *info, const uint8_t *data, int len, void *arg) {
  if (memcmp(info->des_addr, ESP_NOW.BROADCAST_ADDR, 6) == 0) {
    Serial.printf("Unknown peer " MACSTR " sent a broadcast message\n", MAC2STR(info->src_addr));
    Serial.println("Registering the peer as a controller");

    ESP_NOW_Peer_Class new_master(info->src_addr, ESPNOW_WIFI_CHANNEL, WIFI_IF_STA, NULL);

    masters.push_back(new_master);
    if (!masters.back().add_peer()) {
      Serial.println("Failed to register the new controller");
      return;
    }
  } else {
    // The peripherals will only receive broadcast messages
    log_v("Received a unicast message from " MACSTR, MAC2STR(info->src_addr));
    log_v("Ignoring the message");
  }
}

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

    // Initialize SD card
  if (!SD.begin()) {
    Serial.println("SD card initialization failed!");
    return;
  }
  Serial.println("SD card initialized.");

  // Initialize the Wi-Fi module
  WiFi.mode(WIFI_STA);
  WiFi.setChannel(ESPNOW_WIFI_CHANNEL);
  while (!WiFi.STA.started()) delay(100);

  Serial.println("ESP-NOW Example - Broadcast Peripherals");
  Serial.println("Wi-Fi parameters:");
  Serial.println("  Mode: STA");
  Serial.println("  MAC Address: " + WiFi.macAddress());
  Serial.printf("  Channel: %d\n", ESPNOW_WIFI_CHANNEL);

  // Initialize the ESP-NOW protocol
  if (!ESP_NOW.begin()) {
    Serial.println("Failed to initialize ESP-NOW");
    Serial.println("Rebooting in 5 seconds...");
    delay(5000);
    ESP.restart();
  }

  // Register the new peer callback
  ESP_NOW.onNewPeer(register_new_master, NULL);

  Serial.println("Setup complete. Waiting for a control to broadcast a message...");

  tft.init();          // Start the TFT display
  tft.setRotation(2);  // Set the TFT display rotation in landscape mode

  // Clear the screen before writing to it
  tft.fillScreen(TFT_WHITE);
  tft.setTextColor(TFT_BLACK, TFT_WHITE);

  // Display images based on currentIndex
  displayJpeg(images1[currentIndex1], 0);
  displayJpeg(images2[currentIndex2], 1);


  // Set X and Y coordinates for center of display
    int centerX = SCREEN_HEIGHT / 2;
    int centerY = SCREEN_WIDTH / 2;

  tft.drawCentreString("SHOULD", centerX, centerY - 100, FONT_SIZE);
  tft.drawCentreString("YOU", centerX, centerY - 80, FONT_SIZE);
  tft.drawCentreString("TALK", centerX, centerY - 60, FONT_SIZE);
  tft.drawCentreString("TO ME", centerX, centerY - 40, FONT_SIZE);
  tft.drawCentreString("-O-", centerX, centerY - 20, FONT_SIZE);
  tft.drawCentreString("METER", centerX, centerY, FONT_SIZE);

  tft.drawCentreString("Communication is ready", centerX, centerY + 40, 2);

  tft.drawCentreString("By Alexis", 200, 300, 2);
}

void loop() {
  // Update display if the flag is set
  if (updateDisplay) {
    tft.fillScreen(TFT_WHITE);
    tft.setTextColor(TFT_BLACK, TFT_WHITE);
    displayJpeg(images1[currentIndex1], 0);
    displayJpeg(images2[currentIndex2], 1);
    tft.drawCentreString("By Alexis", 200, 300, 2);
    updateDisplay = false; // Reset the flag
  }

  if (updateResult) {
    tft.fillScreen(TFT_WHITE);
    tft.setTextColor(TFT_BLACK, TFT_WHITE);
    displayJpeg(images1[currentIndex1], 0);
    displayJpeg(images2[currentIndex2], 1);
    displayJpeg(allImages[currentIndexAll].src, 2);
    tft.drawCentreString("By Alexis", 200, 300, 2);
    updateResult = false; // Reset the flag
  }
}


void displayJpeg(const char *filename, int imageIndex) {
  Serial.print("Opening file: ");
  Serial.println(filename);

  // Open the file
  File jpegFile = SD.open(filename, FILE_READ);
  if (!jpegFile) {
    Serial.print("Failed to open file ");
    Serial.println(filename);
    return;
  }

  // Decode the JPEG file
  if (JpegDec.decodeSdFile(jpegFile)) {
    Serial.print("Decoding successful for ");
    Serial.println(filename);
    jpegRender(imageIndex);
  } else {
    Serial.print("JPEG file format not supported for ");
    Serial.println(filename);
  }

  // Close the file
  jpegFile.close();
}

void jpegRender(int imageIndex) {
  uint16_t *pImg;
  uint16_t mcu_w = JpegDec.MCUWidth;
  uint16_t mcu_h = JpegDec.MCUHeight;
  uint32_t max_x = JpegDec.width;
  uint32_t max_y = JpegDec.height;

  Serial.print("Rendering image at index ");
  Serial.print(imageIndex);
  Serial.print(", width: ");
  Serial.print(max_x);
  Serial.print(", height: ");
  Serial.println(max_y);

  // Calculate the vertical and horizontal offsets
  int16_t x_offset = 0;
  int16_t y_offset = 0;

  if (imageIndex == 0) {
    // First image positioned at 1/4 of TFT width
    x_offset = tft.width() / 4 - max_x / 2;
  } else if (imageIndex == 1) {
    // Second image positioned at 3/4 of TFT width
    x_offset = 3 * tft.width() / 4 - max_x / 2;
  } else {
    // Center third image at the bottom
    x_offset = (tft.width() - max_x) / 2;
    y_offset = tft.height() - max_y;
  }

  bool swapBytes = tft.getSwapBytes();
  tft.setSwapBytes(true);

  while (JpegDec.read()) {
    pImg = JpegDec.pImage;

    int mcu_x = JpegDec.MCUx * mcu_w;
    int mcu_y = JpegDec.MCUy * mcu_h;

    uint32_t win_w = (mcu_x + mcu_w <= max_x) ? mcu_w : (max_x % mcu_w);
    uint32_t win_h = (mcu_y + mcu_h <= max_y) ? mcu_h : (max_y % mcu_h);

    if ((mcu_x + win_w) <= tft.width() && (mcu_y + win_h) <= tft.height()) {
      tft.pushImage(mcu_x + x_offset, mcu_y + y_offset, win_w, win_h, pImg);
    } else if ((mcu_y + win_h) >= tft.height()) {
      JpegDec.abort();
    }
  }

  tft.setSwapBytes(swapBytes);
}