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

#define SCREEN_WIDTH 128
#define SCREEN_HEIGHT 64
Adafruit_SSD1306 display(SCREEN_WIDTH, SCREEN_HEIGHT, &Wire, -1);

// -------------------- Pins --------------------
const int PIN_POT_VSET   = 34;
const int PIN_POT_ILIM   = 35;
const int PIN_POT_LOAD   = 32;

const int PIN_BTN_OUT    = 25;
const int PIN_BTN_SCREEN = 26;
const int PIN_BTN_UP     = 27;
const int PIN_BTN_DOWN   = 14;

const int PIN_LED_OUT    = 18;
const int PIN_LED_CV     = 19;
const int PIN_LED_CC     = 23;

const int PIN_BUZZER     = 13;

// -------------------- States --------------------
enum SystemState {
  ST_BOOT,
  ST_STANDBY,
  ST_SOFTSTART,
  ST_RUN_CV,
  ST_RUN_CC,
  ST_FAULT,
  ST_COOLDOWN
};

enum FaultCode {
  FLT_NONE,
  FLT_OCP,
  FLT_SCP,
  FLT_OTP
};

SystemState state = ST_BOOT;
FaultCode fault = FLT_NONE;

// -------------------- Measurements --------------------
float vSetPot = 0.0f;
float iLimitPot = 0.0f;
float rLoadPot = 10.0f;

float vTrim = 0.0f;
float vSet = 0.0f;
float iLimit = 1.0f;
float rLoad = 10.0f;

float vRefRamp = 0.0f;
float vOut = 0.0f;
float iOut = 0.0f;
float pOut = 0.0f;
float eff = 0.80f;
float pDiss = 0.0f;
float tempC = 25.0f;
float fanDuty = 0.0f;

bool outputEnabled = false;

// -------------------- UI --------------------
int currentScreen = 0;
const int NUM_SCREENS = 4;

// -------------------- Time Slices --------------------
unsigned long lastInputMs = 0;
unsigned long lastControlMs = 0;
unsigned long lastThermalMs = 0;
unsigned long lastDisplayMs = 0;
unsigned long lastTelemetryMs = 0;

// -------------------- Button debounce --------------------
struct Button {
  int pin;
  bool lastReading;
  bool stableState;
  unsigned long lastChangeTime;
};

Button btnOut    = {PIN_BTN_OUT,    HIGH, HIGH, 0};
Button btnScreen = {PIN_BTN_SCREEN, HIGH, HIGH, 0};
Button btnUp     = {PIN_BTN_UP,     HIGH, HIGH, 0};
Button btnDown   = {PIN_BTN_DOWN,   HIGH, HIGH, 0};

const unsigned long debounceMs = 35;

// -------------------- Fault Log --------------------
struct FaultRecord {
  FaultCode code;
  unsigned long ts;
  float temp;
  float v;
  float i;
};

const int FAULT_LOG_SIZE = 6;
FaultRecord faultLog[FAULT_LOG_SIZE];
int faultLogIndex = 0;
bool faultLogWrapped = false;

// -------------------- Buzzer --------------------
bool buzzerActive = false;
unsigned long buzzerUntil = 0;

// -------------------- Helpers --------------------
float mapFloat(float x, float in_min, float in_max, float out_min, float out_max) {
  return (x - in_min) * (out_max - out_min) / (in_max - in_min) + out_min;
}

const char* stateToStr(SystemState s) {
  switch (s) {
    case ST_BOOT: return "BOOT";
    case ST_STANDBY: return "STBY";
    case ST_SOFTSTART: return "SSRT";
    case ST_RUN_CV: return "CV";
    case ST_RUN_CC: return "CC";
    case ST_FAULT: return "FAULT";
    case ST_COOLDOWN: return "COOL";
    default: return "?";
  }
}

const char* faultToStr(FaultCode f) {
  switch (f) {
    case FLT_NONE: return "NONE";
    case FLT_OCP:  return "OCP";
    case FLT_SCP:  return "SHORT";
    case FLT_OTP:  return "OTP";
    default: return "?";
  }
}

void logFault(FaultCode code) {
  faultLog[faultLogIndex].code = code;
  faultLog[faultLogIndex].ts = millis();
  faultLog[faultLogIndex].temp = tempC;
  faultLog[faultLogIndex].v = vOut;
  faultLog[faultLogIndex].i = iOut;
  faultLogIndex++;
  if (faultLogIndex >= FAULT_LOG_SIZE) {
    faultLogIndex = 0;
    faultLogWrapped = true;
  }
}

void buzzerStart(unsigned long durationMs) {
  buzzerActive = true;
  buzzerUntil = millis() + durationMs;
  digitalWrite(PIN_BUZZER, HIGH);
}

void buzzerStop() {
  buzzerActive = false;
  digitalWrite(PIN_BUZZER, LOW);
}

void updateBuzzer() {
  if (buzzerActive && millis() >= buzzerUntil) {
    buzzerStop();
  }
}

void enterFault(FaultCode f) {
  fault = f;
  state = ST_FAULT;
  outputEnabled = false;
  vOut = 0.0f;
  iOut = 0.0f;
  pOut = 0.0f;
  pDiss = 0.0f;
  logFault(f);
  buzzerStart(700);
}

// -------------------- Button handling --------------------
bool updateButton(Button &b) {
  bool reading = digitalRead(b.pin);

  if (reading != b.lastReading) {
    b.lastChangeTime = millis();
    b.lastReading = reading;
  }

  if ((millis() - b.lastChangeTime) > debounceMs) {
    if (reading != b.stableState) {
      b.stableState = reading;
      if (b.stableState == LOW) {
        return true;
      }
    }
  }
  return false;
}

// -------------------- Reading Inputs --------------------
void readInputs() {
  int rawV = analogRead(PIN_POT_VSET);
  int rawI = analogRead(PIN_POT_ILIM);
  int rawR = analogRead(PIN_POT_LOAD);

  vSetPot = mapFloat(rawV, 0, 4095, 0.0f, 24.0f);
  iLimitPot = mapFloat(rawI, 0, 4095, 0.10f, 3.00f);
  rLoadPot = mapFloat(rawR, 0, 4095, 0.5f, 50.0f);

  if (updateButton(btnOut)) {
    if (state == ST_STANDBY) {
      outputEnabled = true;
      vRefRamp = 0.0f;
      state = ST_SOFTSTART;
      buzzerStart(80);
    } else if (state == ST_RUN_CV || state == ST_RUN_CC || state == ST_SOFTSTART) {
      outputEnabled = false;
      state = ST_STANDBY;
      vOut = 0.0f;
      iOut = 0.0f;
      pOut = 0.0f;
      pDiss = 0.0f;
      buzzerStart(60);
    } else if (state == ST_FAULT) {
      if (tempC < 60.0f) {
        fault = FLT_NONE;
        state = ST_STANDBY;
        buzzerStart(100);
      }
    }
  }

  if (updateButton(btnScreen)) {
    currentScreen++;
    if (currentScreen >= NUM_SCREENS) currentScreen = 0;
  }

  if (updateButton(btnUp)) {
    vTrim += 0.10f;
    if (vTrim > 3.0f) vTrim = 3.0f;
  }

  if (updateButton(btnDown)) {
    vTrim -= 0.10f;
    if (vTrim < -3.0f) vTrim = -3.0f;
  }

  vSet = vSetPot + vTrim;
  if (vSet < 0.0f) vSet = 0.0f;
  if (vSet > 24.0f) vSet = 24.0f;

  iLimit = iLimitPot;
  rLoad = rLoadPot;
}

// -------------------- Control / Plant --------------------
float computeFoldbackLimit(float loadR, float nominalLimit) {
  if (loadR > 3.0f) return nominalLimit;
  if (loadR > 1.5f) return nominalLimit * 0.85f;
  if (loadR > 0.8f) return nominalLimit * 0.65f;
  return nominalLimit * 0.45f;
}

float computeEfficiency(float pout, float temp) {
  float base = 0.76f + 0.08f * (pout / 30.0f);
  if (base > 0.88f) base = 0.88f;
  float tempPenalty = 0.0f;
  if (temp > 60.0f) tempPenalty = (temp - 60.0f) * 0.0025f;
  float result = base - tempPenalty;
  if (result < 0.55f) result = 0.55f;
  return result;
}

void solveOutputForReference(float vRef) {
  float safeR = rLoad;
  if (safeR < 0.05f) safeR = 0.05f;

  float requiredCurrent = vRef / safeR;

  float targetV = 0.0f;
  float targetI = 0.0f;

  if (requiredCurrent <= iLimit) {
    targetV = vRef;
    targetI = targetV / safeR;
    if (state != ST_SOFTSTART) state = ST_RUN_CV;
  } else {
    float limitedI = computeFoldbackLimit(safeR, iLimit);
    targetI = limitedI;
    targetV = targetI * safeR;
    if (state != ST_SOFTSTART) state = ST_RUN_CC;
  }

  static float prevV = 0.0f;
  static float prevI = 0.0f;

  vOut = prevV + 0.25f * (targetV - prevV);
  iOut = prevI + 0.25f * (targetI - prevI);

  prevV = vOut;
  prevI = iOut;

  pOut = vOut * iOut;
  eff = computeEfficiency(pOut, tempC);
  float pIn = (eff > 0.05f) ? (pOut / eff) : pOut;
  pDiss = pIn - pOut;
  if (pDiss < 0.0f) pDiss = 0.0f;
}

void updateControl() {
  switch (state) {
    case ST_BOOT:
      state = ST_STANDBY;
      break;

    case ST_STANDBY:
      vOut = 0.0f;
      iOut = 0.0f;
      pOut = 0.0f;
      pDiss = 0.0f;
      break;

    case ST_SOFTSTART:
      vRefRamp += 0.10f;
      if (vRefRamp > vSet) vRefRamp = vSet;
      solveOutputForReference(vRefRamp);

      if (vRefRamp >= vSet - 0.01f) {
        if ((vSet / max(rLoad, 0.05f)) <= iLimit) state = ST_RUN_CV;
        else state = ST_RUN_CC;
      }
      break;

    case ST_RUN_CV:
    case ST_RUN_CC:
      solveOutputForReference(vSet);
      break;

    case ST_FAULT:
      vOut = 0.0f;
      iOut = 0.0f;
      pOut = 0.0f;
      pDiss = 0.0f;
      break;

    case ST_COOLDOWN:
      vOut = 0.0f;
      iOut = 0.0f;
      pOut = 0.0f;
      pDiss = 0.0f;
      if (tempC < 50.0f) {
        fault = FLT_NONE;
        state = ST_STANDBY;
      }
      break;
  }
}

// -------------------- Protection --------------------
void updateProtection() {
  if (state == ST_STANDBY || state == ST_BOOT || state == ST_FAULT) return;

  if (rLoad < 0.18f) {
    enterFault(FLT_SCP);
    return;
  }

  if (tempC >= 95.0f) {
    fault = FLT_OTP;
    logFault(FLT_OTP);
    buzzerStart(1000);
    state = ST_COOLDOWN;
    outputEnabled = false;
    vOut = 0.0f;
    iOut = 0.0f;
    pOut = 0.0f;
    pDiss = 0.0f;
    return;
  }

  if (iOut > (iLimit * 1.15f)) {
    enterFault(FLT_OCP);
    return;
  }
}

// -------------------- Thermal --------------------
void updateThermal(float dtSeconds) {
  if (tempC > 70.0f) fanDuty = 1.0f;
  else if (tempC > 55.0f) fanDuty = 0.55f;
  else if (tempC > 40.0f) fanDuty = 0.25f;
  else fanDuty = 0.0f;

  float heatIn = 0.09f * pDiss;
  float activeCooling = 0.18f * fanDuty;
  float passiveCooling = 0.020f;

  tempC += (heatIn - activeCooling - passiveCooling) * dtSeconds * 55.0f;

  if (tempC < 25.0f) tempC = 25.0f;
}

// -------------------- LEDs --------------------
void updateLeds() {
  digitalWrite(PIN_LED_OUT, (state == ST_SOFTSTART || state == ST_RUN_CV || state == ST_RUN_CC) ? HIGH : LOW);
  digitalWrite(PIN_LED_CV,  (state == ST_RUN_CV) ? HIGH : LOW);
  digitalWrite(PIN_LED_CC,  (state == ST_RUN_CC) ? HIGH : LOW);
}

// -------------------- OLED Drawing --------------------
void drawHeader(const char* title) {
  display.fillRect(0, 0, 128, 10, SSD1306_WHITE);
  display.setTextColor(SSD1306_BLACK);
  display.setCursor(2, 1);
  display.print(title);
  display.setCursor(92, 1);
  display.print(stateToStr(state));
  display.setTextColor(SSD1306_WHITE);
}

void drawBar(int x, int y, int w, int h, float value, float maxVal) {
  display.drawRect(x, y, w, h, SSD1306_WHITE);
  int fill = 0;
  if (maxVal > 0.001f) fill = (int)((value / maxVal) * (w - 2));
  if (fill < 0) fill = 0;
  if (fill > w - 2) fill = w - 2;
  display.fillRect(x + 1, y + 1, fill, h - 2, SSD1306_WHITE);
}

void screenMain() {
  drawHeader("SMART PSU");

  display.setCursor(0, 14);
  display.print("SET ");
  display.print(vSet, 2);
  display.print("V");

  display.setCursor(72, 14);
  display.print(iLimit, 2);
  display.print("A");

  display.setCursor(0, 28);
  display.print("OUT ");
  display.print(vOut, 2);
  display.print("V");

  display.setCursor(0, 40);
  display.print("CUR ");
  display.print(iOut, 2);
  display.print("A");

  display.setCursor(0, 52);
  display.print("LOAD ");
  display.print(rLoad, 1);
  display.print("R");
}

void screenBars() {
  drawHeader("OUTPUT");

  display.setCursor(0, 14);
  display.print("VOUT");
  drawBar(34, 14, 90, 8, vOut, 24.0f);

  display.setCursor(0, 28);
  display.print("IOUT");
  drawBar(34, 28, 90, 8, iOut, 3.0f);

  display.setCursor(0, 42);
  display.print("POUT");
  drawBar(34, 42, 90, 8, pOut, 30.0f);

  display.setCursor(0, 54);
  if (state == ST_RUN_CC) display.print("MODE CC");
  else if (state == ST_RUN_CV) display.print("MODE CV");
  else display.print("MODE ---");
}

void screenThermal() {
  drawHeader("THERMAL");

  display.setCursor(0, 14);
  display.print("TEMP ");
  display.print(tempC, 1);
  display.print("C");

  display.setCursor(0, 26);
  display.print("FAN  ");
  display.print((int)(fanDuty * 100));
  display.print("%");

  display.setCursor(0, 38);
  display.print("LOSS ");
  display.print(pDiss, 2);
  display.print("W");

  display.setCursor(0, 50);
  display.print("EFF  ");
  display.print((int)(eff * 100));
  display.print("%");

  drawBar(70, 14, 54, 10, tempC, 100.0f);
  drawBar(70, 28, 54, 10, fanDuty * 100.0f, 100.0f);
}

void screenFaults() {
  drawHeader("FAULT LOG");

  if (!faultLogWrapped && faultLogIndex == 0) {
    display.setCursor(0, 18);
    display.print("No faults logged");
    return;
  }

  int shown = 0;
  int idx = faultLogIndex - 1;
  if (idx < 0) idx = FAULT_LOG_SIZE - 1;

  while (shown < 3) {
    if (!faultLogWrapped && idx < 0) break;

    FaultRecord &r = faultLog[idx];
    display.setCursor(0, 14 + shown * 16);
    display.print(faultToStr(r.code));
    display.print(" ");
    display.print(r.temp, 0);
    display.print("C ");
    display.print(r.i, 1);
    display.print("A");

    idx--;
    if (idx < 0) {
      if (faultLogWrapped) idx = FAULT_LOG_SIZE - 1;
      else break;
    }
    shown++;
  }
}

void screenFaultLive() {
  display.clearDisplay();
  display.fillRect(0, 0, 128, 64, SSD1306_WHITE);
  display.setTextColor(SSD1306_BLACK);
  display.setTextSize(2);
  display.setCursor(18, 10);
  display.print("FAULT");

  display.setTextSize(1);
  display.setCursor(38, 36);
  display.print(faultToStr(fault));

  display.setCursor(12, 52);
  display.print("Press OUT to reset");

  display.setTextSize(1);
  display.display();
}

void updateDisplay() {
  if (state == ST_FAULT) {
    screenFaultLive();
    return;
  }

  display.clearDisplay();
  display.setTextSize(1);
  display.setTextColor(SSD1306_WHITE);

  switch (currentScreen) {
    case 0: screenMain(); break;
    case 1: screenBars(); break;
    case 2: screenThermal(); break;
    case 3: screenFaults(); break;
  }

  display.display();
}

// -------------------- Telemetry --------------------
void sendTelemetry() {
  Serial.print("t=");
  Serial.print(millis());
  Serial.print(",state=");
  Serial.print(stateToStr(state));
  Serial.print(",fault=");
  Serial.print(faultToStr(fault));
  Serial.print(",vset=");
  Serial.print(vSet, 2);
  Serial.print(",ilim=");
  Serial.print(iLimit, 2);
  Serial.print(",rload=");
  Serial.print(rLoad, 2);
  Serial.print(",vout=");
  Serial.print(vOut, 2);
  Serial.print(",iout=");
  Serial.print(iOut, 2);
  Serial.print(",pout=");
  Serial.print(pOut, 2);
  Serial.print(",pdiss=");
  Serial.print(pDiss, 2);
  Serial.print(",temp=");
  Serial.print(tempC, 1);
  Serial.print(",fan=");
  Serial.println((int)(fanDuty * 100));
}

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

  pinMode(PIN_BTN_OUT, INPUT_PULLUP);
  pinMode(PIN_BTN_SCREEN, INPUT_PULLUP);
  pinMode(PIN_BTN_UP, INPUT_PULLUP);
  pinMode(PIN_BTN_DOWN, INPUT_PULLUP);

  pinMode(PIN_LED_OUT, OUTPUT);
  pinMode(PIN_LED_CV, OUTPUT);
  pinMode(PIN_LED_CC, OUTPUT);
  pinMode(PIN_BUZZER, OUTPUT);

  digitalWrite(PIN_LED_OUT, LOW);
  digitalWrite(PIN_LED_CV, LOW);
  digitalWrite(PIN_LED_CC, LOW);
  digitalWrite(PIN_BUZZER, LOW);

  Wire.begin(21, 22);

  if (!display.begin(SSD1306_SWITCHCAPVCC, 0x3C)) {
    while (true) {}
  }

  display.clearDisplay();
  display.setTextSize(1);
  display.setTextColor(SSD1306_WHITE);
  display.setCursor(18, 18);
  display.println("SMART PSU");
  display.setCursor(10, 34);
  display.println("INITIALIZING...");
  display.display();

  delay(900);
  state = ST_STANDBY;
}

// -------------------- Loop --------------------
void loop() {
  unsigned long now = millis();

  if (now - lastInputMs >= 20) {
    lastInputMs = now;
    readInputs();
  }

  if (now - lastControlMs >= 20) {
    lastControlMs = now;
    updateControl();
    updateProtection();
    updateLeds();
    updateBuzzer();
  }

  if (now - lastThermalMs >= 50) {
    float dt = (now - lastThermalMs) / 1000.0f;
    lastThermalMs = now;
    updateThermal(dt);
  }

  if (now - lastDisplayMs >= 100) {
    lastDisplayMs = now;
    updateDisplay();
  }

  if (now - lastTelemetryMs >= 250) {
    lastTelemetryMs = now;
    sendTelemetry();
  }
}