#include <ESP8266WiFi.h>
#include <ESP8266WebServer.h>
#include <ESP8266HTTPClient.h>
#include <LittleFS.h>
#include <ArduinoJson.h>
#include <Wire.h>
#include <Adafruit_GFX.h>
#include <Adafruit_SSD1306.h>

#define SCREEN_WIDTH 128
#define SCREEN_HEIGHT 64

#define TJMAX_CPU 95
#define TJMAX_GPU_HOTSPOT 100


String cpuTemp = "--";
String gpuTemp = "--";
String hotspotTemp = "--";
String cpuUtil = "--";
String gpuUtil = "--";
String ssid = "";
String password = "";
String ipadd = "--";

unsigned long lastSuccessfulFetch = millis();
unsigned long lastSwitch = 0;
bool showUtil = false;
bool dataReady = false;


const int buttonPinFreeze = D6;
const int buttonPinCustomUI = D7;
const int ledPin = D5;

bool freezeScreen = false;
unsigned long buttonPressTime = 0;
bool lastButtonState = HIGH;

bool useCustomUI = false;
bool lastCustomUIButtonState = HIGH;


Adafruit_SSD1306 display(SCREEN_WIDTH, SCREEN_HEIGHT, &Wire, -1);
WiFiClient client;
ESP8266WebServer server(80);

void saveCredentials(String newSSID, String newPASS, String newIPADD) 
{
  File f = LittleFS.open("/wifi.txt", "w");
  if (f) 
  {
    f.println(newSSID);
    f.println(newPASS);
    f.println(newIPADD);
    f.close();
  }
  else
  {
    display.clearDisplay();
    display.setCursor(0, 0);
    display.print("File Error");
    display.display();
    delay(5000);
    Serial.println("File error");
  }
}

bool loadCredentials() 
{
  if (LittleFS.exists("/wifi.txt")) 
  {
    File f = LittleFS.open("/wifi.txt", "r");
    if (f) 
    {
      ssid = f.readStringUntil('\n');
      ssid.trim();
      password = f.readStringUntil('\n');
      password.trim();
      ipadd = f.readStringUntil('\n');
      ipadd.trim();
      f.close();
      return true;
    }
  }
  return false;
}

bool connectToWiFi() 
{
  WiFi.begin(ssid.c_str(), password.c_str());
  unsigned long start = millis();
  display.clearDisplay();
  display.setCursor(0, 0);
  display.print("Connecting to WiFi");
  display.display();
  while (WiFi.status() != WL_CONNECTED && millis() - start < 20000) 
  {
    delay(500);
    display.print(".");
    display.display();
    Serial.println(".");
  }
  return WiFi.status() == WL_CONNECTED;
}

void startAPMode() {
  WiFi.softAP("ESP-Setup", "12345678");
  IPAddress IP = WiFi.softAPIP();
  display.clearDisplay();
  display.setCursor(0, 0);
  display.print("AP IP address: " + IP.toString());
  display.display();
  Serial.println("AP IP address: " + IP.toString());
  delay(5000);

  server.on("/", []() {
    server.send(200, "text/html", R"rawliteral(
      <form action="/save" method="POST">
        SSID: <input name="ssid"><br>
        Password: <input name="pass" type="password"><br>
        IP address: <input name="IPadd"><br>
        <input type="submit" value="Save">
      </form>
    )rawliteral");
  });

  server.on("/save", []() {
    String newSSID = server.arg("ssid");
    String newPASS = server.arg("pass");
    String newIPADD = server.arg("IPadd");
    saveCredentials(newSSID, newPASS, newIPADD);
    server.send(200, "text/plain", "Saved. Restarting...");
    delay(1000);
    ESP.restart();
  });

  server.begin();
}

void drawTempWithDegree(int x, int y, const String& label, const String& value, uint8_t textSize = 1) {
  display.setTextSize(textSize);
  display.setCursor(x, y);
  display.print(label);

  int labelWidth = label.length() * 6 * textSize;
  int tempX = x + labelWidth;
  display.setCursor(tempX, y);

  int index = value.indexOf("°");
  String tempNum = (index != -1) ? value.substring(0, index) : value;

  display.print(tempNum);

  int circleX = tempX + tempNum.length() * 6 * textSize + 2;
  int circleY = y + 2 * textSize;
  display.fillCircle(circleX, circleY, 1, SSD1306_WHITE);

  display.setCursor(circleX + 4, y);
  display.print("C");
}//AI generated function

void fetch(JsonArray children) 
{
  bool validDataFound = false;

  for (JsonObject node : children) 
  {
    const char* text = node["Text"];
    const char* value = node["Value"];

    if (text && value) 
    {
      String label = String(text);
      String val = String(value);

      if (label == "Core Average" && val.indexOf("°C") != -1) 
      {
        cpuTemp = val;
        validDataFound = true;
      } 
      else if (label == "GPU Core" && val.indexOf("°C") != -1) 
      {
        gpuTemp = val;
        validDataFound = true;
      } 
      else if (label == "GPU Hot Spot" && val.indexOf("°C") != -1) 
      {
        hotspotTemp = val;
        validDataFound = true;
      } 
      else if (label == "CPU Total" && val.indexOf("%") != -1) 
      {
        cpuUtil = val;
        validDataFound = true;
      } 
      else if (label == "GPU Core" && val.indexOf("%") != -1) 
      {
        gpuUtil = val;
        validDataFound = true;
      }
    }

    if (node.containsKey("Children")) 
    {
      fetch(node["Children"].as<JsonArray>());
    }
  }

  if (validDataFound) {
    dataReady=true;
    lastSuccessfulFetch = millis(); // Reset timeout
  }
}

void showTemps() 
{
  display.clearDisplay();
  if (useCustomUI) 
  {
    display.setTextSize(1);
    display.setCursor(0, 0);
    display.println("CPU Temp:");
    display.setTextSize(2);
    display.setCursor(0, 12);
    display.println(cpuTemp.substring(0, cpuTemp.indexOf("°")) + "oC");

    display.setTextSize(1);
    display.setCursor(0, 36);
    display.println("GPU Temp:");
    display.setTextSize(2);
    display.setCursor(0, 48);
    display.println(gpuTemp.substring(0, gpuTemp.indexOf("°")) + "oC");
  } 
  else 
  {
    display.setTextSize(1);
    drawTempWithDegree(0, 0, "CPU Temp: ", cpuTemp);
    drawTempWithDegree(0, 16, "GPU Temp: ", gpuTemp);
    drawTempWithDegree(0, 32, "Hotspot : ", hotspotTemp);
  }
  display.display();
}

void showUtils() 
{
  display.setTextSize(2);
  display.clearDisplay();
  display.setCursor(0, 0);
  display.println("CPU Util: " + cpuUtil);
  display.println("GPU Util: " + gpuUtil);
  display.display();
}

void checkFreezeButton() 
{
  bool currentState = digitalRead(buttonPinFreeze);

  if (lastButtonState == HIGH && currentState == LOW) 
  {
    buttonPressTime = millis();
  }

  if (lastButtonState == LOW && currentState == HIGH) 
  {
    if (millis() - buttonPressTime >= 1000) 
    {
      freezeScreen = !freezeScreen;
      digitalWrite(ledPin, freezeScreen ? HIGH : LOW);
      Serial.println(freezeScreen ? "Screen frozen." : "Screen resumed.");
    }
  }

  lastButtonState = currentState;
}

void checkCustomUIButton() 
{
  bool currentState = digitalRead(buttonPinCustomUI);

  if (lastCustomUIButtonState == HIGH && currentState == LOW) 
  {
    useCustomUI = !useCustomUI;
  }

  lastCustomUIButtonState = currentState;
}

void TjmaxCheck() 
{
  int cpuTempVal = cpuTemp.substring(0, cpuTemp.indexOf("°")).toInt();
  int hotspotVal = hotspotTemp.substring(0, hotspotTemp.indexOf("°")).toInt();

  static unsigned long lastBlink = 0;
  static bool ledState = LOW;

  bool exceed = (cpuTempVal >= TJMAX_CPU || hotspotVal >= TJMAX_GPU_HOTSPOT);

  if (exceed) 
  {
    if (millis() - lastBlink >= 500) 
    {
      ledState = !ledState;
      digitalWrite(ledPin, ledState);
      lastBlink = millis();
    }
  } 
  else 
  {
    digitalWrite(ledPin, freezeScreen ? HIGH : LOW);
  }
}

void setup() 
{
  Serial.begin(9600);
  LittleFS.begin();
  Wire.begin(D2, D1);
  
  pinMode(buttonPinFreeze, INPUT_PULLUP);
  pinMode(buttonPinCustomUI, INPUT_PULLUP);
  pinMode(ledPin, OUTPUT);
  digitalWrite(ledPin, LOW);

  if (!display.begin(SSD1306_SWITCHCAPVCC, 0x3C)) 
  {
    Serial.println("SSD1306 error");
    while (true);
  }

  display.clearDisplay();
  display.setTextSize(1);
  display.setTextColor(SSD1306_WHITE);
  display.setCursor(0, 0);
  display.display();

  if (loadCredentials()) 
  {
    display.clearDisplay();
    display.setCursor(0, 0);
    display.print("Trying WiFi");
    display.display();
    Serial.println("Trying WiFi...");
    delay(5000);
    if (connectToWiFi()) 
    {
      display.clearDisplay();
      display.setCursor(0, 0);
      display.println("Connected to " + ssid + "\nIP: " + WiFi.localIP().toString());
      display.display();
      Serial.println("Connected to " + ssid);
      Serial.println("IP: " + WiFi.localIP().toString());
      delay(5000);
    }
  }
  else
  {
    display.clearDisplay();
    display.setCursor(0, 0);
    display.print("Starting AP mode");
    display.display();
    Serial.println("Starting AP mode...");
    startAPMode();
    delay(5000);
  }

 delay(1000);
}

void loop() 
{
  server.handleClient();
  checkFreezeButton();
  checkCustomUIButton();

  if (WiFi.getMode() == WIFI_STA && WiFi.status() == WL_CONNECTED) 
  {
    HTTPClient http;
    http.begin(client, ("https://" + ipadd + ":8085/data.json").c_str());
    int httpCode = http.GET();

    if (httpCode > 0) 
    {
      const size_t capacity = 24576;
      DynamicJsonDocument doc(capacity);
      DeserializationError error = deserializeJson(doc, http.getStream());

      if (!error) 
      {
        JsonArray root = doc["Children"];
        fetch(root);
      } 
      else 
      {
        display.clearDisplay();
        display.setTextSize(1);
        display.setCursor(0, 0);
        display.println("JSON Error");
        display.display();
        Serial.print("JSON Error: ");
        Serial.println(error.c_str());
        delay(5000);
      }
    } 
    else 
    {
      display.clearDisplay();
      display.setTextSize(1);
      display.setCursor(0, 0);
      display.println("Fetch failed.");
      display.display();
      Serial.print("GET failed: ");
      Serial.println(http.errorToString(httpCode));
    }

    http.end();

    if (millis() - lastSuccessfulFetch > 60000) 
    {
      display.clearDisplay();
      display.setTextSize(1);
      display.setCursor(0, 0);
      display.println("Fetch failed > 1 min");
      display.println("Rebooting to AP...");
      display.display();
      delay(10000);
      LittleFS.remove("/wifi.txt"); 
      WiFi.disconnect(true);
      ESP.restart();
    }
  }

  if (dataReady && !freezeScreen && !useCustomUI && millis() - lastSwitch >= 3000) 
  {
    showUtil = !showUtil;
    lastSwitch = millis();
  }

  if (dataReady)
  {
    if (showUtil && !useCustomUI) 
    {
      showUtils();
    } 
    else 
    {
      showTemps();
    }
  }

  TjmaxCheck();
  delay(100);
}