//  https://grok.com/share/bGVnYWN5_5e8c1122-1202-4c87-9f07-23ee02c94b4e

#include <Arduino.h>
#include <U8g2lib.h>
#include <Preferences.h>
#include <RotaryEncoder.h>
#include <Wire.h>

// ================== PIN DEFINITIONS ==================
#define OLED_SDA 21
#define OLED_SCL 22
#define ENCODER_CLK 32
#define ENCODER_DT 33
#define ENCODER_BTN 25
#define PWM_CURRENT_PIN 26
#define VOLTAGE_SENSE_PIN 34
#define CURRENT_SENSE_PIN 35
#define OUTPUT_ENABLE_PIN 27

U8G2_SSD1306_128X64_NONAME_F_HW_I2C u8g2(U8G2_R0, U8X8_PIN_NONE);
RotaryEncoder encoder(ENCODER_CLK, ENCODER_DT, RotaryEncoder::LatchMode::FOUR3);
Preferences prefs;

// ================== ENUMS & STRUCTS (from original) ==================
enum Readout {
  READOUT_NONE, READOUT_CURRENT_SETPOINT, READOUT_CURRENT_USAGE,
  READOUT_VOLTAGE, READOUT_POWER, READOUT_RESISTANCE,
  READOUT_TOTAL_CURRENT, READOUT_TOTAL_POWER
};

enum MenuState {
  STATE_SPLASH, STATE_MAIN_MENU, STATE_CC_LOAD, STATE_CONFIG_READOUTS,
  STATE_CHOOSE_READOUT, STATE_SET_MIN_VOLTAGE, STATE_SET_CONTRAST,
  STATE_CALIBRATE, STATE_RESET_TOTALS
};

MenuState currentState = STATE_SPLASH;
int menuSelection = 0;
int readoutPosition = 0; // 1=main, 2=left, 3=right for config

struct DisplayConfig {
  Readout readouts[3] = {READOUT_CURRENT_SETPOINT, READOUT_CURRENT_USAGE, READOUT_VOLTAGE};
} displayConfig;

float currentSetpoint_mA = 0.0;
float measuredCurrent_mA = 0.0;
float measuredVoltage_V = 0.0;
float minVoltage_V = 1.0;
bool outputEnabled = false;
float total_mAh = 0.0;
float total_mWh = 0.0;
unsigned long lastIntegralTime = 0;
unsigned long lastUIUpdate = 0;
unsigned long buttonPressTime = 0;
bool calibrating = false;

// Calibration
struct Calibration {
  float voltage_offset = 0.0;
  float voltage_gain = 0.0122;
  float current_offset = 0.0;
  float current_gain = 10.0;
} cal;

enum CalStep { CAL_NONE, CAL_OFFSETS, CAL_VOLTAGE, CAL_CURRENT };
CalStep calStep = CAL_NONE;
float calRefVoltage = 0.0;
float calRefCurrent = 0.0;

// ================== SETTINGS ==================
void loadSettings() {
  prefs.begin("reload", false);
  currentSetpoint_mA = prefs.getFloat("setpoint", 0);
  minVoltage_V = prefs.getFloat("minV", 1.0);
  // Load cal, displayConfig...
}

void saveSettings() {
  prefs.putFloat("setpoint", currentSetpoint_mA);
  prefs.putFloat("minV", minVoltage_V);
}

// ================== MEASUREMENT & CONTROL (simplified) ==================
void updateMeasurements() {
  int rawV = analogRead(VOLTAGE_SENSE_PIN);
  int rawI = analogRead(CURRENT_SENSE_PIN);
  measuredVoltage_V = (rawV * cal.voltage_gain) + cal.voltage_offset;
  measuredCurrent_mA = (rawI * cal.current_gain) + cal.current_offset;

  if (measuredVoltage_V < minVoltage_V && outputEnabled) outputEnabled = false;
}

void updateTotals() {
  unsigned long now = millis();
  float dt = (now - lastIntegralTime) / 3600000.0f;
  lastIntegralTime = now;
  if (outputEnabled && measuredCurrent_mA > 10) {
    total_mAh += measuredCurrent_mA * dt;
    total_mWh += measuredCurrent_mA * measuredVoltage_V * dt;
  }
}

// ================== Updated Control Function ==================
void updateCurrentControl() {
  digitalWrite(OUTPUT_ENABLE_PIN, outputEnabled ? HIGH : LOW);

  if (!outputEnabled) {
    ledcWrite(PWM_CURRENT_PIN, 0);
    return;
  }

  // Scale setpoint (0-6000 mA) to 8-bit PWM (0-255)
  int pwmVal = map(currentSetpoint_mA, 0, 6000, 0, 255);
  pwmVal = constrain(pwmVal, 0, 255);
  ledcWrite(PWM_CURRENT_PIN, pwmVal);
}

// ================== UI HELPERS & FORMATTING ==================
String formatReadout(Readout r, bool big = false) {
  char buf[16];
  float val = 0;
  const char* unit = "";
  switch (r) {
    case READOUT_CURRENT_SETPOINT:
      val = currentSetpoint_mA / 1000.0; unit = "A";
      break;
    case READOUT_CURRENT_USAGE:
      val = measuredCurrent_mA / 1000.0;
      unit = "A";
      break;
    case READOUT_VOLTAGE:
      val = measuredVoltage_V;
      unit = "V";
      break;
    case READOUT_POWER:
      val = measuredCurrent_mA * measuredVoltage_V / 1000.0;
      unit = "W";
      break;
    case READOUT_RESISTANCE:
      val = (measuredCurrent_mA > 10) ? measuredVoltage_V * 1000.0 / measuredCurrent_mA : 0;
      unit = "Ω";
      break;
    case READOUT_TOTAL_CURRENT:
      val = total_mAh / 1000.0;
      unit = "Ah";
      break;
    case READOUT_TOTAL_POWER:
      val = total_mWh / 1000.0;
      unit = "Wh";
      break;
    default: return "----";
  }
  if (big) sprintf(buf, "%.2f", val); else sprintf(buf, "%.2f", val);
  //Serial.println(String(buf) + unit);
  return String(buf) + unit;
}

void drawBigNumber(float val, const char* unit) {
  u8g2.setFont(u8g2_font_logisoso42_tn);
  char buf[12]; sprintf(buf, "%.2f", val);
  u8g2.drawStr(4, 52, buf);
  u8g2.setFont(u8g2_font_6x12_tr);
  u8g2.drawStr(100, 40, unit);
}

// Splash (simple version)
void drawSplash() {
  u8g2.clearBuffer();
  u8g2.setFont(u8g2_font_logisoso32_tn);
  u8g2.drawStr(20, 40, "Re:load");
  u8g2.setFont(u8g2_font_6x12_tr);
  u8g2.drawStr(30, 55, "ESP32 Port");
  u8g2.sendBuffer();
}

// ================== MENU DRAWING ==================
const char* readoutNames[] = {
  "None", "Set Current", "Act. Current", "Voltage", "Power",
  "Resistance", "Total mAh", "Total mWh"
};

void drawMenu(const char* title, const char** items, int count, int sel) {
  u8g2.clearBuffer();
  u8g2.setFont(u8g2_font_6x12_tr);
  if (title) u8g2.drawStr(0, 12, title);

  for (int i = 0; i < 5 && i < count; i++) {
    int idx = (sel / 4 * 4) + i;
    if (idx >= count) break;
    int y = 24 + i * 12;
    u8g2.drawStr(10, y, items[idx]);
    if (idx == sel) u8g2.drawStr(0, y, ">");
  }
  u8g2.sendBuffer();
}

void updateDisplay() {
  if (millis() - lastUIUpdate < 50) return;
  lastUIUpdate = millis();
  u8g2.clearBuffer();

  if (currentState == STATE_SPLASH) {
    drawSplash();
    return;
  }

  if (currentState == STATE_CC_LOAD) {
    drawBigNumber(measuredCurrent_mA / 1000.0, "A");
    u8g2.setFont(u8g2_font_6x12_tr);
    u8g2.drawStr(90, 12, "SET");
    u8g2.drawStr(90, 24, formatReadout(READOUT_CURRENT_SETPOINT, false).c_str());
    // Left / Right readouts from displayConfig
    u8g2.drawStr(0, 42, formatReadout(displayConfig.readouts[1]).c_str());
    u8g2.drawStr(70, 42, formatReadout(displayConfig.readouts[2]).c_str());
    u8g2.drawStr(0, 63, outputEnabled ? "ON" : "OFF");
  }
  else if (currentState == STATE_MAIN_MENU) {
    const char* items[] = {"C/C Load", "Readouts", "Min Voltage", "Reset Totals", "Contrast", "Calibrate", "EXIT"};
    drawMenu("Main Menu", items, 7, menuSelection);
  }
  else if (currentState == STATE_CONFIG_READOUTS) {
    const char* items[] = {"Main disp.", "Left disp.", "Right disp.", "EXIT"};
    drawMenu("Readouts", items, 4, menuSelection);
  }
  else if (currentState == STATE_CHOOSE_READOUT) {
    drawMenu("Choose value", readoutNames, 8, menuSelection);
  }
  else if (currentState == STATE_SET_MIN_VOLTAGE) {
    u8g2.setFont(u8g2_font_6x12_tr);
    u8g2.drawStr(0, 20, "Min Voltage (V):");
    u8g2.drawStr(0, 40, String(minVoltage_V, 2).c_str());
  }
  else if (currentState == STATE_CALIBRATE) {
    u8g2.setFont(u8g2_font_6x12_tr);
    u8g2.drawStr(0, 20, "Calibration");
    if (calStep == CAL_OFFSETS) u8g2.drawStr(0, 40, "Step 1: Offsets (no load)");
    else if (calStep == CAL_VOLTAGE) u8g2.drawStr(0, 40, "Step 2: Voltage");
    else if (calStep == CAL_CURRENT) u8g2.drawStr(0, 40, "Step 3: Current");
  }
  u8g2.sendBuffer();
}

// ================== INPUT (with Long Press) ==================
void handleInput() {
  encoder.tick();
  int delta = encoder.getPosition();
  //Serial.print(F("Encoder: "));Serial.println(delta);

  if (delta != 0) {
    if (currentState == STATE_CC_LOAD) {
      currentSetpoint_mA += delta * 10.0; // 10mA steps
      currentSetpoint_mA = constrain(currentSetpoint_mA, 0, 6000);
    } else if (currentState == STATE_SET_MIN_VOLTAGE) {
      minVoltage_V += delta * 0.1;
      minVoltage_V = constrain(minVoltage_V, 0.5, 60);
    } else {
      menuSelection = constrain(menuSelection + delta, 0, 20);
    }
    encoder.setPosition(0);
    saveSettings();
  }

  Serial.print(F("currentSetpoint_mA: ")); Serial.println(currentSetpoint_mA);

  bool btn = digitalRead(ENCODER_BTN);
  if (btn == LOW) {
    if (buttonPressTime == 0) buttonPressTime = millis();
    if (millis() - buttonPressTime > 1000 && currentState != STATE_CC_LOAD) {
      // Long press: exit to main
      currentState = STATE_MAIN_MENU;
      menuSelection = 0;
      buttonPressTime = 0;
    }
  } else if (buttonPressTime > 0) {
    // Short press
    if (millis() - buttonPressTime < 800) {
      handleShortPress();
    }
    buttonPressTime = 0;
  }
}

void handleShortPress() {
  if (currentState == STATE_MAIN_MENU) {
    switch (menuSelection) {
      case 0: currentState = STATE_CC_LOAD; break;
      case 1: currentState = STATE_CONFIG_READOUTS; menuSelection = 0; break;
      case 2: currentState = STATE_SET_MIN_VOLTAGE; break;
      case 3: total_mAh = total_mWh = 0; break;
      case 5: currentState = STATE_CALIBRATE; calStep = CAL_OFFSETS; calibrating = true; break;
      case 6: currentState = STATE_MAIN_MENU; break; // EXIT
    }
  }
  else if (currentState == STATE_CONFIG_READOUTS) {
    if (menuSelection < 3) {
      readoutPosition = menuSelection + 1;
      currentState = STATE_CHOOSE_READOUT;
      menuSelection = 0;
    } else currentState = STATE_MAIN_MENU;
  }
  else if (currentState == STATE_CHOOSE_READOUT) {
    if (menuSelection < 8) {
      displayConfig.readouts[readoutPosition - 1] = (Readout)menuSelection;
    }
    currentState = STATE_CONFIG_READOUTS;
  }
  else if (currentState == STATE_CC_LOAD) {
    outputEnabled = !outputEnabled;
  }
  else if (currentState == STATE_CALIBRATE) {
    // Advance calibration steps (in real code: prompt user for reference values)
    if (calStep == CAL_OFFSETS) {
      cal.voltage_offset = 0;
      cal.current_offset = 0;
      calStep = CAL_VOLTAGE;
    }
    else if (calStep == CAL_VOLTAGE) {
      calStep = CAL_CURRENT;
    }
    else {
      calStep = CAL_NONE;
      currentState = STATE_MAIN_MENU;
      calibrating = false;
    }
  }
  else {
    currentState = STATE_MAIN_MENU;
  }
  menuSelection = 0;
}

// ================== SERIAL & SETUP/LOOP ==================
void handleSerial() { /* same as before */ }

void setup() {
  Serial.begin(115200);
  Serial.println(F("ESP32 Load Setup"));

  u8g2.begin();

  pinMode(ENCODER_BTN, INPUT_PULLUP);
  pinMode(OUTPUT_ENABLE_PIN, OUTPUT);
  //pinMode(PWM_CURRENT_PIN, OUTPUT);

  //ledcSetup(0, 50000, 8);
  // === NEW LEDC API (ESP32 core 3.x) ===
  ledcAttach(PWM_CURRENT_PIN, 1000, 8);  // pin, frequency (50kHz), resolution (8-bit)

  loadSettings();
  lastIntegralTime = millis();
  currentState = STATE_SPLASH;
  delay(2000); // Show splash
  currentState = STATE_MAIN_MENU;
}

void loop() {
  handleInput();
  updateMeasurements();
  updateTotals();
  updateCurrentControl();
  updateDisplay();
  handleSerial();
  delay(10);
}