//========================================================
//  BUTTON
//========================================================

  #include  <Button.h>

  Button  BlueBody(0);      // Connect your button between pin 2 and GND
  Button  BlueHead(1);      // Connect your button between pin 3 and GND
  Button  RedBody(2);       // Connect your button between pin 4 and GND
  Button  RedHead(3);       // Connect your button between pin 4 and GND
  Button  YellowButton(5);  // Connect your button between pin 4 and GND

//========================================================
//  OLED Display 
//========================================================

  #include <Wire.h>
  #include <Adafruit_GFX.h>
  #include <Adafruit_SSD1306.h>

  #define SCREEN_WIDTH      128   // OLED display width, in pixels
  #define SCREEN_HEIGHT     64    // OLED display height, in pixels
  #define OLED_RESET        -1    // Reset pin # (or -1 if sharing Arduino reset pin)
  #define SCREEN_ADDRESS    0x3C  //< See datasheet for Address; 0x3D for 128x64, 0x3C for 128x32
  Adafruit_SSD1306          display(SCREEN_WIDTH, SCREEN_HEIGHT, &Wire, OLED_RESET);

  // Start Display
  void displayInit() {

    // SSD1306_SWITCHCAPVCC = generate display voltage from 3.3V internally
    if(!display.begin(SSD1306_SWITCHCAPVCC, SCREEN_ADDRESS)) {
      Serial.println(F("SSD1306 allocation failed"));
      for(;;); // Don't proceed, loop forever
    }

    // Clearing Out Every Thing
    display.clearDisplay();
    display.display();

  }

  // Display Identity
  void displayIdentity() {

    // Display MAC Address
    display.setTextSize(1); // Draw 2X-scale text
    display.setTextColor(SSD1306_WHITE);
    display.setCursor(0, 0);
    display.println();
    display.println("TAKEONE JoyStick");
    display.println();
    display.println(WiFi.localIP());
    display.println();
    display.println(WiFi.macAddress());
    display.display();

  }

  // Display Identity
  void displayStatus(const char* title, const char* line2 = nullptr, const char* line3 = nullptr) {

    display.clearDisplay();
    display.setTextSize(2);
    display.setTextColor(SSD1306_WHITE);
    display.setCursor(0, 0);
    display.println(title);
    display.setTextSize(1);
    
    if(line2) {
      display.setCursor(0, 20);
      display.println(line2);
    }
    
    if(line3) {
      display.setCursor(0, 35);
      display.println(line3);
    }

    display.display();

  }

//========================================================
//  WiFi Connection
//========================================================

  #include <WiFiManager.h>
  #include <SPIFFS.h>
  #include <ArduinoJson.h>

  // TCP client
  WiFiClient client;

  // Global config variables
  char server[40]                   = "";
  char port[6]                      = "8080";
  char deviceName[32]               = "ESP_Device";
  bool shouldSaveConfig             = false;

  // Variables
  unsigned long       lastWiFiCheck = 0;
  unsigned long       lastTCPCheck  = 0;
  const unsigned long checkInterval = 5000; // 5 seconds

  // Call Back Function
  void saveConfigCallback() { 
    shouldSaveConfig = true; 
  }

  // Robust SPIFFS mount/format function
  bool safeMountFS() {
    if (SPIFFS.begin(true)) return true; // true = format if mount fails
    Serial.println("SPIFFS mount/format failed!");
    return false;
  }

  // Build The Configuration
  void setupWiFiAndConfig(const char* apName) {
    displayStatus("Starting...");
    safeMountFS();

    displayStatus("Loading", "Config...");
    bool configLoaded = false;
    if (SPIFFS.exists("/config.json")) {
      File f = SPIFFS.open("/config.json", "r");
      if (f) {
        DynamicJsonDocument j(512);
        if (!deserializeJson(j, f)) {
          strcpy(server, j["server"] | "");
          strcpy(port, j["port"] | "8080");
          strcpy(deviceName, j["deviceName"] | "ESP_Device");
          configLoaded = true;
        }
        f.close();
      }
    }
    if (!configLoaded) {
      strcpy(server, "");
      strcpy(port, "8080");
      strcpy(deviceName, "ESP_Device");
    }

    displayStatus("WiFi Setup", "Open Portal");
    WiFiManagerParameter pServer("server", "Server", server, 40);
    WiFiManagerParameter pPort("port", "Port", port, 6);
    WiFiManagerParameter pDevName("devname", "Device Name", deviceName, 32);
    WiFiManager wm;
    wm.setSaveConfigCallback(saveConfigCallback);
    wm.addParameter(&pServer); wm.addParameter(&pPort); wm.addParameter(&pDevName);

    if (!wm.autoConnect(apName)) {
      displayStatus("Restarting...", "WiFi Failed");
      delay(2000);
      ESP.restart();
    }

    strcpy(server, pServer.getValue());
    strcpy(port, pPort.getValue());
    strcpy(deviceName, pDevName.getValue());

    displayStatus("WiFi OK", WiFi.localIP().toString().c_str());

    if (shouldSaveConfig) {
      File f = SPIFFS.open("/config.json", "w");
      if (f) {
        DynamicJsonDocument j(512);
        j["server"] = server; j["port"] = port; j["deviceName"] = deviceName;
        serializeJson(j, f); f.close();
      }
    }
    delay(1000); // Give user a moment to see the status
  }

  // Reset WiFi
  void resetWiFiAndConfig() {
    // Erase WiFi credentials
    WiFi.disconnect(true, true); // erase WiFi credentials from flash[1]
    delay(1000);

    // Remove config file from SPIFFS
    SPIFFS.begin(true); // Ensure SPIFFS is mounted
    if (SPIFFS.exists("/config.json")) {
      SPIFFS.remove("/config.json");
    }

    // Optionally, restart the device to trigger portal on next boot
    ESP.restart();
  }

  // Monitor Connection
  void monitorConnection() {

    unsigned long now = millis();

    // 1. WiFi connectivity check and reconnect
    if(now - lastWiFiCheck > checkInterval) {
      lastWiFiCheck = now;
      if (WiFi.status() != WL_CONNECTED) {
        Serial.println("WiFi disconnected, trying to reconnect...");
        WiFi.reconnect();
      }
    }

    // 2. TCP connectivity check and reconnect
    if(now - lastTCPCheck > checkInterval) {
      lastTCPCheck = now;
      if (WiFi.status() == WL_CONNECTED) {
        if (!client.connected()) {
          Serial.println("TCP disconnected, trying to reconnect...");
          client.stop(); // Ensure clean state
          if (client.connect(server, atoi(port))) {
            Serial.println("TCP reconnected!");
          } else {
            Serial.println("TCP reconnect failed.");
          }
        }
      }
    }

    // Example: if connected, send something every 10 seconds
    static unsigned long lastSend = 0;
    if(client.connected() && now - lastSend > 10000) {
      client.print("Hello from ");
      client.println(deviceName);
      lastSend = now;
    }

  }

  // Send Device Identity on TCP Connect ---
  void sendDeviceIdentity() {
    DynamicJsonDocument doc(128);
    doc["id"] = WiFi.macAddress();
    doc["type"] = deviceName;
    String out;
    serializeJson(doc, out);
    client.println(out);
  }

  // Send Button Event ---
  void sendButtonMessage(const char* btn) {
    DynamicJsonDocument doc(64);
    doc["button"] = btn;
    String out;
    serializeJson(doc, out);
    client.println(out);
  }

  // Send Cancel Event ---
  void sendCancelMessage() {
    DynamicJsonDocument doc(64);
    doc["button"] = "cancel";
    String out;
    serializeJson(doc, out);
    client.println(out);
  }

  // // --- Yellow Button Handler ---
  // void checkYellowButton() {
  //   bool pressed = (digitalRead(BTN_YELLOW) == LOW);

  //   if (pressed && !yellowBtnWasPressed) {
  //     yellowBtnPressStart = millis();
  //     yellowLongActionDone = false;
  //     yellowBtnWasPressed = true;
  //   }
  //   if (!pressed && yellowBtnWasPressed) {
  //     unsigned long pressDuration = millis() - yellowBtnPressStart;
  //     if (pressDuration < 500) {
  //       sendCancelMessage();
  //     }
  //     yellowBtnWasPressed = false;
  //   }
  //   if (pressed && !yellowLongActionDone && (millis() - yellowBtnPressStart > 3000)) {
  //     // Long press: WiFi reset
  //     // Display message if you have an OLED
  //     // displayStatus("WiFi Reset", "Hold...");
  //     resetWiFiAndConfig();
  //     yellowLongActionDone = true;
  //   }
  // }

  // Action Button Handler
  void checkActionButtons() {
    
    if(BlueBody.pressed()) { 
      sendButtonMessage("blue_body")
    };
    
    if(BlueHead.pressed()) { 
      sendButtonMessage("blue_head") 
    };
    
    if(RedBody.pressed()) { 
      sendButtonMessage("red_body") 
    };
    
    if(RedHead.pressed()) { 
      sendButtonMessage("red_head") 
    };
    
  }