#include <Arduino.h>
#include <WiFi.h>
#include <WiFiUdp.h>
#include <WiFiManager.h>

// -------------------------------
// ESPUI + AsyncWebServer
// -------------------------------
#include <AsyncTCP.h>
#include <ESPAsyncWebServer.h>
#include <ESPUI.h>

// -------------------------------
// Preferences
// -------------------------------
#include <Preferences.h>
Preferences prefs;

AsyncWebServer server(80);

// -------------------------------
// Firmware version
// -------------------------------
const char* FIRMWARE_VERSION = "V2.4_UDP_LEDC_WM_SLEW";

// -------------------------------
// UDP
// -------------------------------
WiFiUDP udp;
const int listenPort = 12345;   // default / legacy constant
int udpPort = 12345;            // runtime port, loaded from Preferences

// -------------------------------
// PWM setup (LEDC, ESP32 Core 3.x)
// -------------------------------
const uint8_t NUM_CHANNELS = 8;

uint8_t pwmPins[NUM_CHANNELS] = {
    26, // D0
    22, // D1
    21, // D2
    17, // D3
    16, // D4
    5,  // D5
    18, // D6
    19  // D7
    // 23 (D8) remains as a spare
};

const uint32_t pwmFrequency = 25000; // 25 kHz
const uint8_t pwmResolution = 10;    // 10-bit resolution (0–1023)

// -------------------------------
// Channel labels for debugging
// -------------------------------
const char* channelLabels[NUM_CHANNELS] = {
  "CPU Load",
  "CPU Temp",
  "RAM Usage",
  "GPU Load",
  "GPU Temp",
  "VRAM Usage",
  "Reserved 6",
  "Reserved 7"
};

const char* channelUnits[NUM_CHANNELS] = {
  "%",     // CPU Load
  "°C",    // CPU Temp
  "%",     // RAM Usage
  "%",     // GPU Load
  "°C",    // GPU Temp
  "%",     // VRAM Usage
  "",      // Reserved
  ""       // Reserved
};

// -------------------------------
// Channel labels ESPUI
// -------------------------------
const char* channelDropdownLabels[NUM_CHANNELS] = {
    "CH0 (CPU Load)",
    "CH1 (CPU Temp)",
    "CH2 (RAM Usage)",
    "CH3 (GPU Load)",
    "CH4 (GPU Temp)",
    "CH5 (VRAM Usage)",
    "CH6 (Reserved 6)",
    "CH7 (Reserved 7)"
};

// -------------------------------
// Calibration tables
// -------------------------------
float logicalPoints[5] = {0, 25, 50, 75, 100};

float calibratedPoints[NUM_CHANNELS][5] = {
  {0.0f, 25.0f, 50.0f, 75.0f, 99.0f},
  {0.0f, 24.0f, 49.0f, 74.0f, 98.0f},
  {0.0f, 26.0f, 51.0f, 76.0f, 99.0f},
  {0.0f, 25.0f, 50.0f, 75.0f, 97.0f},
  {0.0f, 25.0f, 50.0f, 75.0f, 99.0f},
  {0.0f, 24.0f, 50.0f, 74.0f, 98.0f},
  {0.0f, 25.0f, 49.0f, 75.0f, 97.0f},
  {0.0f, 26.0f, 50.0f, 76.0f, 99.0f}
};

// -------------------------------
// Duty tracking + Slew system
// -------------------------------
float currentDuty[NUM_CHANNELS] = {0.0f};
float targetDuty[NUM_CHANNELS]  = {0.0f};
float slewStartDuty[NUM_CHANNELS] = {0.0f};

unsigned long slewStartTime = 0;
const unsigned long slewDuration = 1000; // 1 second smooth transition

// -------------------------------
// Watchdog (UDP-based)
// -------------------------------
unsigned long lastPacketTime = 0;
const unsigned long watchdogTimeout = 5000;  // 5 seconds
unsigned long lastFadeTime = 0;

// -------------------------------
// Animation tuning
// -------------------------------
const float FADE_IN_FACTOR  = 0.999f;   // boot-up 0 → 100%
const float FADE_OUT_FACTOR = 0.999f;   // watchdog 100% → 0
const unsigned long FADE_INTERVAL = 1;  // ms between fade steps

// -------------------------------
// Connection state machine
// -------------------------------
enum ConnectionState {
    STATE_DISCONNECTED,        // fade to zero
    STATE_CONNECTING,          // fade in 0 → 100%
    STATE_WAIT_FOR_FIRST_PACKET, // hold at 100%, wait for first UDP
    STATE_CONNECTED            // normal UDP-driven slew
};

ConnectionState connectionState = STATE_DISCONNECTED;

// -------------------------------
// ESPUI controls
// -------------------------------
uint16_t portInput;   // ID of the UDP port Number control

// --- Calibration UI Controls ---
uint16_t calChannelDropdown;
uint16_t calInputs[5];
uint16_t calTestValueInput;
uint16_t calSaveButton;
uint16_t calOverrideSwitch;

uint8_t selectedCalChannel = 0;
bool overrideActive[NUM_CHANNELS] = {false};

uint16_t connectionStatusLabel;

// -------------------------------
// Calibration interpolation
// -------------------------------
float applyCalibration(uint8_t ch, float logicalDuty) {
  if (logicalDuty <= 0.0f) return calibratedPoints[ch][0];
  if (logicalDuty >= 100.0f) return calibratedPoints[ch][4];

  for (int i = 0; i < 4; i++) {
    if (logicalDuty >= logicalPoints[i] && logicalDuty <= logicalPoints[i+1]) {
      float x0 = logicalPoints[i];
      float x1 = logicalPoints[i+1];
      float y0 = calibratedPoints[ch][i];
      float y1 = calibratedPoints[ch][i+1];
      float t = (logicalDuty - x0) / (x1 - x0);
      return y0 + t * (y1 - y0);
    }
  }
  return calibratedPoints[ch][4];
}

// -------------------------------
// Calibration UI helpers & callbacks
// -------------------------------
void refreshCalibrationUI() {
  for (int i = 0; i < 5; i++) {
    ESPUI.updateControlValue(calInputs[i], String(calibratedPoints[selectedCalChannel][i], 2));
  }
}

void calChannelCallback(Control *sender, int type) {
    // Turn override OFF when switching channels
    for (int ch = 0; ch < NUM_CHANNELS; ch++) {
        overrideActive[ch] = false;
    }
    ESPUI.updateControlValue(calOverrideSwitch, "0");

    selectedCalChannel = sender->value.toInt();
    Serial.print("Calibration channel changed to ");
    Serial.println(selectedCalChannel);

    refreshCalibrationUI();
}

void calPointCallback(Control *sender, int type) {
  int index = sender->id - calInputs[0];
  if (index >= 0 && index < 5) {
    float val = sender->value.toFloat();
    calibratedPoints[selectedCalChannel][index] = val;
    Serial.printf("Cal[%d][%d] = %.2f\n", selectedCalChannel, index, val);
  }
}

void calTestCallback(Control *sender, int type) {
    if (!overrideActive[selectedCalChannel]) return;

    float logical = sender->value.toFloat();   // 0–100 integer

    if (logical < 0) logical = 0;
    if (logical > 100) logical = 100;

    float calibrated = applyCalibration(selectedCalChannel, logical);
    int duty = (int)((calibrated / 100.0f) * ((1 << pwmResolution) - 1));

    ledcWrite(pwmPins[selectedCalChannel], duty);

    Serial.printf("Override update CH%d: logical=%.2f calibrated=%.2f duty=%d\n",
                  selectedCalChannel, logical, calibrated, duty);
}

void calSaveCallback(Control *sender, int type) {
  Serial.printf("Saving calibration for CH%d...\n", selectedCalChannel);

  for (int i = 0; i < 5; i++) {
    String key = "cal_" + String(selectedCalChannel) + "_" + String(i);
    prefs.putFloat(key.c_str(), calibratedPoints[selectedCalChannel][i]);
  }

  Serial.println("Calibration saved.");
}

void calOverrideSwitchCallback(Control *sender, int type) {
    bool enabled = sender->value.toInt() == 1;

    if (enabled) {
        Serial.println("Override enabled.");

        // Enable override only for the selected channel
        for (int ch = 0; ch < NUM_CHANNELS; ch++) {
            overrideActive[ch] = (ch == selectedCalChannel);
        }

        // Immediately apply the test value
        float logical = ESPUI.getControl(calTestValueInput)->value.toFloat();
        if (logical < 0) logical = 0;
        if (logical > 100) logical = 100;

        float calibrated = applyCalibration(selectedCalChannel, logical);
        int duty = (int)((calibrated / 100.0f) * ((1 << pwmResolution) - 1));
        ledcWrite(pwmPins[selectedCalChannel], duty);

        Serial.printf("Override driving CH%d: logical=%.2f calibrated=%.2f duty=%d\n",
                      selectedCalChannel, logical, calibrated, duty);

    } else {
        Serial.println("Override disabled. Returning to UDP control.");

        // Disable all overrides
        for (int ch = 0; ch < NUM_CHANNELS; ch++) {
            overrideActive[ch] = false;
        }
    }
}

void updateConnectionStatusUI(ConnectionState state) {
    const char* text = "Unknown";

    switch (state) {
        case STATE_DISCONNECTED:
            text = "Disconnected";
            break;
        case STATE_CONNECTING:
            text = "Connecting...";
            break;
        case STATE_WAIT_FOR_FIRST_PACKET:
            text = "Waiting for first UDP packet";
            break;
        case STATE_CONNECTED:
            text = "Connected";
            break;
    }

    ESPUI.updateControlValue(connectionStatusLabel, text);
}

// -------------------------------
// Setup
// -------------------------------
void setup() {
  Serial.begin(115200);
  delay(300);

  Serial.println("Booting Analog System Monitor (UDP + LEDC + WiFiManager + Slew)");
  Serial.print("Firmware: ");
  Serial.println(FIRMWARE_VERSION);

  // LEDC PWM init
  for (int ch = 0; ch < NUM_CHANNELS; ch++) {
    bool ok = ledcAttach(pwmPins[ch], pwmFrequency, pwmResolution);
    if (!ok) {
      Serial.print("LEDC attach failed on pin ");
      Serial.println(pwmPins[ch]);
    }
    ledcWrite(pwmPins[ch], 0);
  }

  // WiFi Manager
  WiFiManager wm;
  wm.setHostname("AnalogMonitor");
  wm.setTimeout(180);

  Serial.println("Starting WiFiManager...");
  bool res = wm.autoConnect("AnalogMonitor-Setup");

  if (!res) {
    Serial.println("WiFi failed or timed out. Rebooting...");
    delay(2000);
    ESP.restart();
  }

  Serial.println("WiFi connected!");
  Serial.print("IP: ");
  Serial.println(WiFi.localIP());

  // Preferences: load UDP port (default = listenPort)
  prefs.begin("analogmon", false);
  udpPort = prefs.getInt("udpPort", listenPort);

  // Load calibration from flash (if present)
  for (int ch = 0; ch < NUM_CHANNELS; ch++) {
    for (int i = 0; i < 5; i++) {
      String key = "cal_" + String(ch) + "_" + String(i);
      float val = prefs.getFloat(key.c_str(), calibratedPoints[ch][i]);
      calibratedPoints[ch][i] = val;
    }
  }

  // Start UDP with runtime port
  udp.begin(udpPort);
  Serial.print("Listening on UDP port ");
  Serial.println(udpPort);

  lastPacketTime = millis();

  // -------------------------------
  // ESPUI Web Interface
  // -------------------------------
  ESPUI.setVerbosity(Verbosity::Verbose);

  // Create tabs
  uint16_t tabSettings = ESPUI.addControl(ControlType::Tab, "Settings", "Settings");
  uint16_t tabLighting = ESPUI.addControl(ControlType::Tab, "Lighting", "Lighting");
  uint16_t tabCalibration = ESPUI.addControl(ControlType::Tab, "Calibration", "Calibration");

  // Restart button callback
  auto restartCallback = [](Control *sender, int type) {
      ESP.restart();
  };

  // Port input callback
  auto portInputCallback = [](Control *sender, int type) {
      Serial.print("Port input changed to: ");
      Serial.println(sender->value);
  };

  // Save & Apply callback
  auto savePortCallback = [](Control *sender, int type) {
      if (type != B_UP) return;   // Prevent double-trigger

      Control* c = ESPUI.getControl(portInput);
      int newPort = c->value.toInt();

      if (newPort < 1024 || newPort > 65535) {
          Serial.println("Invalid port (1024–65535)");
          return;
      }

      prefs.putInt("udpPort", newPort);
      udpPort = newPort;

      udp.stop();
      udp.begin(udpPort);

      Serial.print("New UDP port applied: ");
      Serial.println(udpPort);

      ESPUI.updateControlValue(portInput, String(newPort));
  };

  // ------------------------------------------------------
  // Row simulation: Port input + Save button
  // ------------------------------------------------------

// -------------------------------------------
// Connection Status section
// -------------------------------------------
  ESPUI.addControl(
      ControlType::Separator,
      "Connection Status",
      "",
      ControlColor::None,
      tabSettings
  );

  // Live-updating connection status label
  connectionStatusLabel = ESPUI.addControl(
      ControlType::Label,
      "Status",
      "Disconnected",
      ControlColor::Wetasphalt,
      tabSettings
  );

  // -------------------------------------------
  // UDP Telemetry Connection Settings section
  // -------------------------------------------

  ESPUI.addControl(
      ControlType::Separator,
      "UDP Telemetry Connection Settings",
      "",
      ControlColor::None,
      tabSettings
  );

  // UDP Port Number input
  portInput = ESPUI.addControl(
      ControlType::Number,
      "UDP Port",
      String(udpPort),
      ControlColor::Peterriver,
      tabSettings,
      portInputCallback
  );

  // Save & Apply button
  ESPUI.addControl(
      ControlType::Button,
      "UDP Port",
      "Save & Apply",
      ControlColor::Emerald,
      tabSettings,
      savePortCallback
  );

  // Existing separator (leave as-is)
  ESPUI.addControl(
      ControlType::Separator,
      "",
      "",
      ControlColor::None,
      tabSettings
  );

  // Restart button (unchanged)
  ESPUI.addControl(
      ControlType::Button,
      "Restart ESP32",
      "Restart",
      ControlColor::Alizarin,
      tabSettings,
      restartCallback
  );

  // Lighting tab placeholder
  ESPUI.addControl(
      ControlType::Label,
      "Lighting Placeholder",
      "Coming soon...",
      ControlColor::Emerald,
      tabLighting
  );

  // -------------------------------
  // Calibration tab UI
  // -------------------------------

  // Channel selector
  calChannelDropdown = ESPUI.addControl(
      ControlType::Select,
      "Selected Channel",
      "0",
      ControlColor::Peterriver,
      tabCalibration,
      calChannelCallback
  );

  // Add options 0–7
  for (int i = 0; i < NUM_CHANNELS; i++) {
      char value[8];
      snprintf(value, sizeof(value), "%d", i);

      ESPUI.addControl(
          ControlType::Option,
          channelDropdownLabels[i],   // static label
          value,                      // static-ish value (OK)
          ControlColor::None,
          calChannelDropdown
      );
  }

  ESPUI.addControl(
      ControlType::Separator,
      "",
      "",
      ControlColor::None,
      tabCalibration
  );

  // Calibration inputs
  const char* calNames[5] = {"0%", "25%", "50%", "75%", "100%"};

  for (int i = 0; i < 5; i++) {
    calInputs[i] = ESPUI.addControl(
        ControlType::Number,
        calNames[i],
        String(calibratedPoints[0][i], 2),
        ControlColor::Wetasphalt,
        tabCalibration,
        calPointCallback
    );
  }

  ESPUI.addControl(
      ControlType::Separator,
      "",
      "",
      ControlColor::None,
      tabCalibration
  );

  // Test value input
  calTestValueInput = ESPUI.addControl(
      ControlType::Slider,
      "Test Value",
      "50",
      ControlColor::Carrot,
      tabCalibration,
      calTestCallback
  );

  calOverrideSwitch = ESPUI.addControl(
      ControlType::Switcher,
      "Override",
      "0",  // default OFF
      ControlColor::Alizarin,
      tabCalibration,
      calOverrideSwitchCallback
  );

  ESPUI.addControl(
      ControlType::Separator,
      "",
      "",
      ControlColor::None,
      tabCalibration
  );

  // Save button
  calSaveButton = ESPUI.addControl(
      ControlType::Button,
      "Save Calibration",
      "Save",
      ControlColor::Emerald,
      tabCalibration,
      calSaveCallback
  );

  // Start ESPUI
  ESPUI.sliderContinuous = true;   // enables live slider updates
  ESPUI.begin("Analog Monitor UI");
}

// -------------------------------
// Loop
// -------------------------------
void loop() {

  // -------- UDP parsing --------
  int packetSize = udp.parsePacket();
  if (packetSize > 0) {
    char buf[256];
    int len = udp.read(buf, sizeof(buf) - 1);
    buf[len] = '\0';

    float values[NUM_CHANNELS] = {0};
    int idx = 0;

    char* token = strtok(buf, ",");
    while (token != nullptr && idx < NUM_CHANNELS) {
      values[idx] = atof(token);
      idx++;
      token = strtok(nullptr, ",");
    }

    if (idx == NUM_CHANNELS) {

      if (connectionState == STATE_DISCONNECTED) {
          // First valid packet after being disconnected → start CONNECTING
          Serial.println("STATE CHANGE: DISCONNECTED → CONNECTING (UDP connection established)");
          connectionState = STATE_CONNECTING;
          updateConnectionStatusUI(connectionState);

          // Initialize fade-in: start from 0 on all non-override channels
          for (int ch = 0; ch < NUM_CHANNELS; ch++) {
              if (!overrideActive[ch]) {
                  currentDuty[ch] = 0.0f;
              }
          }

          lastPacketTime = millis();  // prevent watchdog during fade-in
          // Ignore this packet's values during fade-in (Option B)
      }
      else if (connectionState == STATE_CONNECTING) {
          // Ignore UDP values during fade-in, just keep watchdog alive
          lastPacketTime = millis();
      }
      else if (connectionState == STATE_WAIT_FOR_FIRST_PACKET) {
          // First real packet after fade-in completes
          Serial.println("STATE CHANGE: WAIT_FOR_FIRST_PACKET → CONNECTED (first UDP packet received)");

          for (int ch = 0; ch < NUM_CHANNELS; ch++) {
              if (!overrideActive[ch]) {
                  targetDuty[ch] = values[ch];
                  slewStartDuty[ch] = currentDuty[ch]; // currently ~100%
              }
          }

          slewStartTime = millis();
          lastPacketTime = millis();
          connectionState = STATE_CONNECTED;
          updateConnectionStatusUI(connectionState);

          // Debug output
          Serial.println("Received UDP packet (first after fade-in):");
          for (int i = 0; i < NUM_CHANNELS; i++) {
            Serial.print("  CH");
            Serial.print(i);
            Serial.print(" (");
            Serial.print(channelLabels[i]);
            Serial.print("): ");
            Serial.print(values[i], 2);
            if (channelUnits[i][0] != '\0') {
              Serial.print(" ");
              Serial.print(channelUnits[i]);
            }
            Serial.println();
          }
          Serial.println();
      }
      else if (connectionState == STATE_CONNECTED) {
          // Normal UDP-driven update
          for (int ch = 0; ch < NUM_CHANNELS; ch++) {
              if (!overrideActive[ch]) {
                  targetDuty[ch] = values[ch];
                  slewStartDuty[ch] = currentDuty[ch];
              }
          }
          slewStartTime = millis();
          lastPacketTime = millis();

          // Debug output
          Serial.println("Received UDP packet:");
          for (int i = 0; i < NUM_CHANNELS; i++) {
            Serial.print("  CH");
            Serial.print(i);
            Serial.print(" (");
            Serial.print(channelLabels[i]);
            Serial.print("): ");
            Serial.print(values[i], 2);
            if (channelUnits[i][0] != '\0') {
              Serial.print(" ");
              Serial.print(channelUnits[i]);
            }
            Serial.println();
          }
          Serial.println();
      }
    }
  }

  // -------- Connection state machine --------
  unsigned long now = millis();

  switch (connectionState) {

      case STATE_CONNECTED: {
          // Check for lost connection
          if (now - lastPacketTime > watchdogTimeout) {
              Serial.println("STATE CHANGE: CONNECTED → DISCONNECTED (UDP connection lost)");
              connectionState = STATE_DISCONNECTED;
              updateConnectionStatusUI(connectionState);
              break;
          }

          // Normal slew-rate limiting
          float progress = (float)(now - slewStartTime) / (float)slewDuration;
          if (progress > 1.0f) progress = 1.0f;

          for (int ch = 0; ch < NUM_CHANNELS; ch++) {
              if (!overrideActive[ch]) {
                  float newDuty = slewStartDuty[ch] + (targetDuty[ch] - slewStartDuty[ch]) * progress;
                  currentDuty[ch] = newDuty;

                  float calibratedDuty = applyCalibration(ch, newDuty);
                  int duty = (int)((calibratedDuty / 100.0f) * ((1 << pwmResolution) - 1));
                  ledcWrite(pwmPins[ch], duty);
              }
          }
      } break;

      case STATE_CONNECTING: {
          // If we lose packets even while connecting, fall back to DISCONNECTED
          if (now - lastPacketTime > watchdogTimeout) {
              Serial.println("STATE CHANGE: CONNECTING → DISCONNECTED (no packets during fade-in)");
              connectionState = STATE_DISCONNECTED;
              updateConnectionStatusUI(connectionState);
              break;
          }

          // Fade-in animation: 0 → 100% on all non-override channels
          if (now - lastFadeTime >= FADE_INTERVAL) {
              lastFadeTime = now;

              bool allReached = true;

              for (int ch = 0; ch < NUM_CHANNELS; ch++) {
                  if (!overrideActive[ch]) {
                      float cd = currentDuty[ch];

                      // Exponential approach to 100%
                      cd = cd * FADE_IN_FACTOR + 100.0f * (1.0f - FADE_IN_FACTOR);
                      currentDuty[ch] = cd;

                      float calibratedDuty = applyCalibration(ch, cd);
                      int duty = (int)((calibratedDuty / 100.0f) * ((1 << pwmResolution) - 1));
                      ledcWrite(pwmPins[ch], duty);

                      // Check if we're close enough to 100%
                      if (cd < 99.0f) {
                          allReached = false;
                      }
                  }
              }

              if (allReached) {
                  Serial.println("STATE CHANGE: CONNECTING → STATE_WAIT_FOR_FIRST_PACKET (fade-in complete)");
                  connectionState = STATE_WAIT_FOR_FIRST_PACKET;
                  updateConnectionStatusUI(connectionState);
              }
          }
      } break;

      case STATE_WAIT_FOR_FIRST_PACKET: {
          // Hold at ~100%, do nothing until first UDP packet arrives
          // (handled in UDP parsing above)
      } break;

      case STATE_DISCONNECTED: {
          // Watchdog fade-to-zero (always active in this state)
          if (now - lastFadeTime >= FADE_INTERVAL) {
              lastFadeTime = now;

              for (int ch = 0; ch < NUM_CHANNELS; ch++) {

                if (currentDuty[ch] > 0.0f) {

                  currentDuty[ch] *= FADE_OUT_FACTOR;

                  if (currentDuty[ch] < 0.01f)
                    currentDuty[ch] = 0.0f;

                  float calibratedDuty = applyCalibration(ch, currentDuty[ch]);
                  int duty = (int)((calibratedDuty / 100.0f) * ((1 << pwmResolution) - 1));

                  ledcWrite(pwmPins[ch], duty);
                }
              }
          }
      } break;
  }
}
