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

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

// I2C
static const int PIN_I2C_SDA = 4;
static const int PIN_I2C_SCL = 5;

// Buttons
static const int PIN_BTN_UP   = 6;
static const int PIN_BTN_DOWN = 7;
static const int PIN_BTN_ARM  = 8;

// PWM
static const int PIN_PWM_OUT  = 2;
static const int PWM_FREQ     = 5000;
static const int PWM_RES_BITS = 10;
static const int PWM_MAX_DUTY = (1 << PWM_RES_BITS) - 1;

// Limits / UX
static const int HARD_CAP_PCT  = 80;   // hard max for safety
static const int STEP_PCT      = 10;   // tap step
static const int DEBOUNCE_MS   = 25;

// ARM behavior
static const int ARM_HOLD_MS   = 1200; // long press to ARM
static const int RAMP_TOTAL_MS = 700;  // soft-start duration
static const int RAMP_STEP_MS  = 10;

// State
int targetPct  = 0;
int appliedPct = 0;
bool armed     = false;

// Simple debounce for UP/DOWN
bool lastUp = true, lastDown = true;
unsigned long lastUpChange = 0, lastDownChange = 0;

// Robust debounced button for ARM
struct DebouncedButton {
  int pin;
  bool lastStable = true; // pullup idle HIGH
  bool lastRead = true;
  unsigned long lastChangeMs = 0;
  bool fellEvent = false;
  bool roseEvent = false;

  void begin(int p) {
    pin = p;
    pinMode(pin, INPUT_PULLUP);
  }

  void update() {
    bool r = digitalRead(pin);
    fellEvent = false;
    roseEvent = false;

    if (r != lastRead) {
      lastRead = r;
      lastChangeMs = millis();
    }

    if ((millis() - lastChangeMs) > DEBOUNCE_MS && lastStable != lastRead) {
      bool prev = lastStable;
      lastStable = lastRead;
      if (prev == HIGH && lastStable == LOW) fellEvent = true;   // press
      if (prev == LOW  && lastStable == HIGH) roseEvent = true;  // release
    }
  }

  bool isPressed() const { return lastStable == LOW; }
  bool fell() const { return fellEvent; }
  bool rose() const { return roseEvent; }
};

DebouncedButton armBtn;

// ARM press timing
unsigned long armPressedAt = 0;
bool armLongHandled = false;

void pwmWriteDuty(int duty) {
  duty = constrain(duty, 0, PWM_MAX_DUTY);
  ledcWrite(PIN_PWM_OUT, duty);
}

void writePwmFromPct(int pct) {
  pct = constrain(pct, 0, HARD_CAP_PCT);
  int duty = armed ? (pct * PWM_MAX_DUTY) / 100 : 0;
  pwmWriteDuty(duty);
}

void rampToPct(int newPct) {
  newPct = constrain(newPct, 0, HARD_CAP_PCT);

  if (!armed) {
    appliedPct = 0;
    writePwmFromPct(0);
    return;
  }

  int start = appliedPct;
  int end = newPct;
  if (start == end) return;

  int steps = max(1, RAMP_TOTAL_MS / RAMP_STEP_MS);
  for (int i = 1; i <= steps; i++) {
    int p = start + (end - start) * i / steps;
    writePwmFromPct(p);
    appliedPct = p;
    delay(RAMP_STEP_MS);
  }
}

bool fellEdge(int pin, bool &lastState, unsigned long &lastChangeMs) {
  bool r = digitalRead(pin); // HIGH idle, LOW pressed
  if (r != lastState && (millis() - lastChangeMs) > DEBOUNCE_MS) {
    lastChangeMs = millis();
    bool prev = lastState;
    lastState = r;
    return (prev == HIGH && r == LOW);
  }
  return false;
}

void drawUI() {
  display.clearDisplay();
  display.setTextColor(SSD1306_WHITE);

  display.setTextSize(1);
  display.setCursor(0, 0);
  display.print("LASER CTRL");

  display.setCursor(78, 0);
  display.print(armed ? "ARMED" : "SAFE");

  display.setTextSize(3);
  display.setCursor(0, 18);
  display.print(targetPct);

  display.setTextSize(2);
  display.print("%");

  display.setTextSize(1);
  display.setCursor(0, 54);
  display.print("UP/DN ");
  display.print(STEP_PCT);
  display.print("%  ARM:hold");

  display.display();
}

void forceSafe() {
  armed = false;
  appliedPct = 0;
  pwmWriteDuty(0);
}

void doArm() {
  armed = true;
  appliedPct = 0;
  writePwmFromPct(0);
  rampToPct(targetPct);
}

void setup() {
  Serial.begin(115200);

  Wire.begin(PIN_I2C_SDA, PIN_I2C_SCL);

  pinMode(PIN_BTN_UP, INPUT_PULLUP);
  pinMode(PIN_BTN_DOWN, INPUT_PULLUP);
  armBtn.begin(PIN_BTN_ARM);

  ledcAttach(PIN_PWM_OUT, PWM_FREQ, PWM_RES_BITS);

  // SAFE on boot
  targetPct = 0;
  forceSafe();

  if (!display.begin(SSD1306_SWITCHCAPVCC, 0x3C)) {
    Serial.println("OLED init failed");
  } else {
    drawUI();
  }

  // Prime ARM button state
  delay(50);
  armBtn.update();
}

void loop() {
  // UP/DOWN adjust target (allowed even in SAFE; output remains 0)
  if (fellEdge(PIN_BTN_UP, lastUp, lastUpChange)) {
    targetPct = min(HARD_CAP_PCT, targetPct + STEP_PCT);
    if (armed) rampToPct(targetPct);
    drawUI();
  }

  if (fellEdge(PIN_BTN_DOWN, lastDown, lastDownChange)) {
    targetPct = max(0, targetPct - STEP_PCT);
    if (armed) rampToPct(targetPct);
    drawUI();
  }

  // ARM button update (ONLY ONCE per loop)
  armBtn.update();

  // Press start
  if (armBtn.fell()) {
    armPressedAt = millis();
    armLongHandled = false;
  }

  // Long press to ARM (only if currently SAFE)
  if (armBtn.isPressed() && !armLongHandled) {
    if (millis() - armPressedAt >= ARM_HOLD_MS) {
      armLongHandled = true;
      if (!armed) doArm();
      drawUI();
    }
  }

  // Release: short press => SAFE (only if ARMED)
  if (armBtn.rose()) {
    unsigned long held = millis() - armPressedAt;
    if (held < ARM_HOLD_MS) {
      if (armed) {
        forceSafe();
        drawUI();
      }
    }
  }
}