#include <Wire.h>
#include <Adafruit_GFX.h>
#include <Adafruit_SSD1306.h>
#include <FS.h>
#include <SPIFFS.h>
#include <WiFiManager.h>
#include <ArduinoJson.h>

#define PRODUCT_NAME    "TAKEONE"

// OLED display setup (on I2C bus 0)
#define OLED_SDA 18
#define OLED_SCL 19
#define SCREEN_WIDTH 128
#define SCREEN_HEIGHT 32
#define OLED_RESET    -1

Adafruit_SSD1306 display(SCREEN_WIDTH, SCREEN_HEIGHT, &Wire, OLED_RESET);

// Second I2C bus for other sensors/devices
#define I2C2_SDA 1
#define I2C2_SCL 2
TwoWire I2C_Two = TwoWire(1);

// MQTT/WiFi config
char mqtt_server[40] = "";
char mqtt_port[6] = "8080";
char api_token[34] = "YOUR_API_TOKEN";
bool shouldSaveConfig = false;

// WiFi loss detection
unsigned long wifiLostTime = 0;
bool portalOpened = false;

// Button reset
#define BOOT_BUTTON 0 // Usually GPIO 0 on ESP32

void saveConfigCallback() {
  
  shouldSaveConfig = true;
  display.clearDisplay();
  display.setCursor(0, 0);
  display.println("Saving config...");
  display.display();

}

void showStatus(const String& line1, const String& line2 = "", const String& line3 = "") {

  display.clearDisplay();
  display.setCursor(0, 0);
  display.println(line1);
  if (!line2.isEmpty()) display.println(line2);
  if (!line3.isEmpty()) display.println(line3);
  display.display();

}

void setup() {

  Serial.begin(115200);

  // Start both I2C buses
  Wire.begin(OLED_SDA, OLED_SCL);
  I2C_Two.begin(I2C2_SDA, I2C2_SCL);

  // OLED init
  if (!display.begin(SSD1306_SWITCHCAPVCC, 0x3C)) {
    Serial.println("SSD1306 allocation failed");
    while (1);
  }
  display.setTextSize(1);
  display.setTextColor(WHITE);
  showStatus("Booting...");

  // Mount SPIFFS
  if (!SPIFFS.begin(true)) {
    showStatus("SPIFFS mount", "failed!");
    delay(2000);
  } else if (SPIFFS.exists("/config.json")) {
    File configFile = SPIFFS.open("/config.json", "r");
    if (configFile) {
      DynamicJsonDocument json(256);
      DeserializationError err = deserializeJson(json, configFile);
      if (!err) {
        strlcpy(mqtt_server, json["mqtt_server"] | "", sizeof(mqtt_server));
        strlcpy(mqtt_port, json["mqtt_port"] | "8080", sizeof(mqtt_port));
        strlcpy(api_token, json["api_token"] | "YOUR_API_TOKEN", sizeof(api_token));
      }
      configFile.close();
    }
  }

  // WiFiManager parameters
  WiFiManagerParameter custom_mqtt_server("server", "MQTT Server", mqtt_server, 40);
  WiFiManagerParameter custom_mqtt_port("port", "MQTT Port", mqtt_port, 6);
  WiFiManagerParameter custom_api_token("apikey", "API Token", api_token, 33);

  WiFiManager wifiManager;
  wifiManager.setSaveConfigCallback(saveConfigCallback);
  wifiManager.addParameter(&custom_mqtt_server);
  wifiManager.addParameter(&custom_mqtt_port);
  wifiManager.addParameter(&custom_api_token);

  showStatus("Connecting WiFi...");
  if (!wifiManager.autoConnect(PRODUCT_NAME)) {
    showStatus("WiFi Connect", "Failed!", "Restarting...");
    delay(3000);
    ESP.restart();
  }

  // Update config values
  strlcpy(mqtt_server, custom_mqtt_server.getValue(), sizeof(mqtt_server));
  strlcpy(mqtt_port, custom_mqtt_port.getValue(), sizeof(mqtt_port));
  strlcpy(api_token, custom_api_token.getValue(), sizeof(api_token));

  // Save config if needed
  if (shouldSaveConfig) {
    DynamicJsonDocument json(256);
    json["mqtt_server"] = mqtt_server;
    json["mqtt_port"] = mqtt_port;
    json["api_token"] = api_token;
    File configFile = SPIFFS.open("/config.json", "w");
    if (configFile) {
      serializeJson(json, configFile);
      configFile.close();
    }
  }

  // Show final info (truncate to fit OLED)
  String ipStr = WiFi.localIP().toString();
  if (ipStr.length() > 15) ipStr = ipStr.substring(0, 15);
  showStatus("IPAD: " + ipStr, "MQTT: " + String(mqtt_server), "Port: " + String(mqtt_port));

  // Button setup
  pinMode(BOOT_BUTTON, INPUT_PULLUP);
}

void loop() {

  // WiFi loss detection
  if (WiFi.status() != WL_CONNECTED) {
    if (wifiLostTime == 0) {
      wifiLostTime = millis();
    } else if (millis() - wifiLostTime > 10000 && !portalOpened) {
      showStatus("WiFi lost", "Opening portal...");
      WiFiManager wifiManager;
      portalOpened = true;
      wifiManager.startConfigPortal(PRODUCT_NAME);
      portalOpened = false;
      wifiLostTime = 0;
    }
  } else {
    wifiLostTime = 0;
  }

  // Button hold check
  static unsigned long buttonPressedTime = 0;
  static bool buttonActionTaken = false;

  if (digitalRead(BOOT_BUTTON) == LOW) {
    if (buttonPressedTime == 0) {
      buttonPressedTime = millis();
    } else if (millis() - buttonPressedTime > 5000 && !buttonActionTaken) {
      showStatus("Erasing WiFi", "Opening portal...");
      WiFi.disconnect(true);
      WiFiManager wifiManager;
      wifiManager.resetSettings();
      wifiManager.startConfigPortal("AutoConnectAP", "12345678");
      buttonActionTaken = true;
    }
  } else {
    buttonPressedTime = 0;
    buttonActionTaken = false;
  }

}