// ESP COACH PRO

#include <Wire.h>
#include <SPI.h>
#include <SD.h>
#include <Adafruit_GFX.h>
#include <Adafruit_SSD1306.h>
#include <MPU6050_tockn.h>
#include "esp_sleep.h"

// ========== PIN DEFINES ==========
const int PIN_BTN_UP    = 32;
const int PIN_BTN_DOWN  = 33;
const int PIN_BTN_ENTER = 25;
const int PIN_BTN_BACK  = 26;

const int PIN_SD_CS     = 5;

const int PIN_I2C_SDA   = 21;
const int PIN_I2C_SCL   = 22;

const int PIN_BATTERY   = 34;

// ========== OLED / MPU ==========
#define SCREEN_WIDTH 128
#define SCREEN_HEIGHT 64
#define OLED_RESET    -1

Adafruit_SSD1306 oled(SCREEN_WIDTH, SCREEN_HEIGHT, &Wire, OLED_RESET);
MPU6050 mpu(Wire);

// ========== BUTTON STATE ==========
struct ButtonState {
  int  pin;
  bool last;
  bool pressed;
};

ButtonState btnUp    = {PIN_BTN_UP,    HIGH, false};
ButtonState btnDown  = {PIN_BTN_DOWN,  HIGH, false};
ButtonState btnEnter = {PIN_BTN_ENTER, HIGH, false};
ButtonState btnBack  = {PIN_BTN_BACK,  HIGH, false};

// ========== ACTIVITY / AUTO-OFF ==========
unsigned long lastActivityMs = 0;

const uint8_t AUTO_OFF_CHOICES = 4;
const uint16_t autoOffMinutes[AUTO_OFF_CHOICES] = {5, 10, 15, 30};
uint8_t  autoOffIndex = 2;
unsigned long autoOffMs = 0;

void recordActivity() {
  lastActivityMs = millis();
}

void updateAutoOffMs() {
  autoOffMs = (unsigned long)autoOffMinutes[autoOffIndex] * 60UL * 1000UL;
}

// ========== ROW DATA ==========
struct RowData {
  unsigned long elapsedMs;
  float t_s;
  float spmAvg;
  float tempC;
  float vbat;
  int   battPct;
};

RowData rowData;

// ========== MODES & SESSION STATES ==========
enum AppMode {
  MODE_MENU,
  MODE_JUST_ROW,
  MODE_RACE_1500,
  MODE_RACE_2000,
  MODE_WORKOUT_MENU,
  MODE_WORKOUT_RUN,
  MODE_STROKE_CT,
  MODE_RATE_ONLY,
  MODE_SETTINGS_MENU,
  MODE_SETTINGS_SENS,
  MODE_SETTINGS_CALIB,
  MODE_SETTINGS_AUTO_OFF,
  MODE_SETTINGS_LOGGING
};

enum SessionState {
  STATE_IDLE,
  STATE_ARMED,
  STATE_RUNNING
};

AppMode      currentMode   = MODE_MENU;
SessionState sessionState  = STATE_IDLE;

// ========== MAIN MENU (SCROLLING) ==========
const uint8_t MENU_ITEMS = 7;
const char* menuItems[MENU_ITEMS] = {
  "JUST ROW",
  "1500 RACE",
  "2000 RACE",
  "WORKOUTS",
  "STROKE COUNT",
  "RATE ONLY",
  "SETTINGS"
};

const uint8_t MENU_VISIBLE_ITEMS = 4;
uint8_t menuIndex    = 0;
uint8_t menuTopIndex = 0;

// ========== SETTINGS MENU (SCROLLING) ==========
const uint8_t SETTINGS_ITEMS = 4;
const char* settingsItems[SETTINGS_ITEMS] = {
  "Sensitivity",
  "Calibrate IMU",
  "Auto-off",
  "Logging"
};
const uint8_t SETTINGS_VISIBLE_ITEMS = 3;
uint8_t settingsIndex     = 0;
uint8_t settingsTopIndex  = 0;

// ========== STROKE / MODEL ==========
float metersPerStroke = 10.0f;
int   sensitivityLevel = 1;

float thresholdHigh = 0.3f;
float thresholdLow  = 0.1f;

bool   waitingForPeak      = true;
unsigned long lastStrokeMillis = 0;
unsigned long strokeCount  = 0;
float  strokeRateSPM       = 0.0f;

const int SPM_BUF_SIZE = 5;
float spmBuf[SPM_BUF_SIZE];
int   spmIdx    = 0;
bool  spmFilled = false;

// 1-minute stroke stats
unsigned long minuteWindowStartMs = 0;
unsigned int  minuteStrokeCount   = 0;
unsigned int  lastMinuteStrokes   = 0;

// session timing
unsigned long sessionStartMillis = 0;

// logging
File logFile;
unsigned long lastLogMillis = 0;
bool loggingEnabled = true;

// ==== SD paths for PRO ====
const char* DATA_DIR_PRO      = "/ESP_COACH/PRO";
const char* LOG_FILE_PRO      = "/ESP_COACH/PRO/log.csv";
const char* SETTINGS_FILE_PRO = "/ESP_COACH/PRO/settings.txt";

void saveSettingsPro() {
  SD.mkdir("/ESP_COACH");
  SD.mkdir(DATA_DIR_PRO);
  File f = SD.open(SETTINGS_FILE_PRO, FILE_WRITE);
  if (!f) return;
  f.print("SENS=");    f.println(sensitivityLevel);
  f.print("AUTOOFF="); f.println(autoOffIndex);
  f.print("LOG=");     f.println(loggingEnabled ? 1 : 0);
  f.close();
}

void loadSettingsPro() {
  if (!SD.exists(SETTINGS_FILE_PRO)) return;
  File f = SD.open(SETTINGS_FILE_PRO, FILE_READ);
  if (!f) return;
  while (f.available()) {
    String line = f.readStringUntil('\n');
    line.trim();
    if (line.startsWith("SENS=")) {
      sensitivityLevel = line.substring(5).toInt();
    } else if (line.startsWith("AUTOOFF=")) {
      autoOffIndex = line.substring(8).toInt();
    } else if (line.startsWith("LOG=")) {
      loggingEnabled = (line.substring(4).toInt() != 0);
    }
  }
  f.close();

  if (sensitivityLevel < 0) sensitivityLevel = 0;
  if (sensitivityLevel > 2) sensitivityLevel = 2;
  if (autoOffIndex >= AUTO_OFF_CHOICES) autoOffIndex = AUTO_OFF_CHOICES - 1;
}

// ========== WORKOUTS ==========
struct WorkoutStep {
  uint16_t duration_s;
  bool isWork;
  const char* label;
};

struct WorkoutPlan {
  const char* name;
  uint8_t numSteps;
  const WorkoutStep* steps;
};

const WorkoutStep steps_steady20[] = {
  {20 * 60, true, "20min"}
};

const WorkoutStep steps_8x1[] = {
  {60, true,  "W1"},
  {60, false, "R"},
  {60, true,  "W2"},
  {60, false, "R"},
  {60, true,  "W3"},
  {60, false, "R"},
  {60, true,  "W4"},
  {60, false, "R"},
  {60, true,  "W5"},
  {60, false, "R"},
  {60, true,  "W6"},
  {60, false, "R"},
  {60, true,  "W7"},
  {60, false, "R"},
  {60, true,  "W8"},
  {60, false, "R"}
};

const WorkoutStep steps_pyramid[] = {
  {60,  true,  "1"},
  {60,  false, "R"},
  {120, true,  "2"},
  {60,  false, "R"},
  {180, true,  "3"},
  {60,  false, "R"},
  {240, true,  "4"},
  {60,  false, "R"},
  {300, true,  "5"},
  {60,  false, "R"},
  {240, true,  "4"},
  {60,  false, "R"},
  {180, true,  "3"},
  {60,  false, "R"},
  {120, true,  "2"},
  {60,  false, "R"},
  {60,  true,  "1"}
};

const WorkoutStep steps_6x2[] = {
  {120, true,  "W1"},
  {90,  false, "R"},
  {120, true,  "W2"},
  {90,  false, "R"},
  {120, true,  "W3"},
  {90,  false, "R"},
  {120, true,  "W4"},
  {90,  false, "R"},
  {120, true,  "W5"},
  {90,  false, "R"},
  {120, true,  "W6"},
  {90,  false, "R"}
};

const WorkoutStep steps_ladder[] = {
  {240, true, "20spm"},
  {240, true, "24spm"},
  {240, true, "28spm"}
};

const WorkoutPlan workoutPlans[] = {
  {"20min steady",     (uint8_t)(sizeof(steps_steady20)/sizeof(WorkoutStep)), steps_steady20},
  {"8x1'/1' on/off",   (uint8_t)(sizeof(steps_8x1)/sizeof(WorkoutStep)),      steps_8x1},
  {"Pyramid 1-5-1",    (uint8_t)(sizeof(steps_pyramid)/sizeof(WorkoutStep)),  steps_pyramid},
  {"6x2'/90s rest",    (uint8_t)(sizeof(steps_6x2)/sizeof(WorkoutStep)),      steps_6x2},
  {"Rate ladder 3x4'", (uint8_t)(sizeof(steps_ladder)/sizeof(WorkoutStep)),   steps_ladder}
};

const uint8_t WORKOUT_PLAN_COUNT    = sizeof(workoutPlans)/sizeof(WorkoutPlan);
const uint8_t WORKOUT_VISIBLE_ITEMS = 3;
uint8_t workoutMenuIndex = 0;
uint8_t workoutTopIndex  = 0;

const WorkoutPlan* currentPlan = nullptr;
uint8_t currentStepIdx   = 0;
unsigned long stepStartMs = 0;
bool workoutFinished      = false;

// ========== SENSITIVITY ==========
void applySensitivity() {
  switch (sensitivityLevel) {
    case 0:
      thresholdHigh = 0.2f;
      thresholdLow  = 0.06f;
      break;
    case 2:
      thresholdHigh = 0.4f;
      thresholdLow  = 0.2f;
      break;
    case 1:
    default:
      thresholdHigh = 0.3f;
      thresholdLow  = 0.1f;
      break;
  }
}

// ========== BUTTONS ==========
void updateButton(ButtonState &b) {
  bool state = digitalRead(b.pin);
  bool nowPressed = (b.last == HIGH && state == LOW);
  b.pressed = nowPressed;
  b.last    = state;
  if (nowPressed) {
    recordActivity();
  }
}

void readButtons() {
  updateButton(btnUp);
  updateButton(btnDown);
  updateButton(btnEnter);
  updateButton(btnBack);
}

// ========== TIME & BATTERY ==========
void formatTime(char *buf, size_t len, unsigned long ms) {
  unsigned long sec = ms / 1000;
  unsigned long m = sec / 60;
  unsigned long s = sec % 60;
  snprintf(buf, len, "%02lu:%02lu", m, s);
}

float avgSPM() {
  int count = spmFilled ? SPM_BUF_SIZE : spmIdx;
  if (count == 0) return strokeRateSPM;
  float sum = 0;
  for (int i = 0; i < count; i++) sum += spmBuf[i];
  return sum / count;
}

void addSPM(float v) {
  spmBuf[spmIdx] = v;
  spmIdx++;
  if (spmIdx >= SPM_BUF_SIZE) {
    spmIdx = 0;
    spmFilled = true;
  }
}

float readBatteryVoltage() {
  int raw = analogRead(PIN_BATTERY);
  float v = (raw / 4095.0f) * 3.3f;
  v *= 2.0f;
  return v;
}

int batteryPercent(float vbat) {
  float pct = (vbat - 3.3f) / (4.2f - 3.3f) * 100.0f;
  if (pct < 0) pct = 0;
  if (pct > 100) pct = 100;
  return (int)(pct + 0.5f);
}

// ========== AUTO-OFF & POWER ==========
void enterDeepSleep() {
  if (logFile) {
    logFile.flush();
    logFile.close();
  }
  saveSettingsPro();

  oled.clearDisplay();
  oled.setTextSize(1);
  oled.setCursor(18, 24);
  oled.print("Powering off...");
  oled.display();
  delay(700);

  esp_sleep_enable_ext0_wakeup((gpio_num_t)PIN_BTN_ENTER, 0);
  oled.ssd1306_command(SSD1306_DISPLAYOFF);
  esp_deep_sleep_start();
}

void checkAutoOff() {
  unsigned long now = millis();
  if (now - lastActivityMs > autoOffMs) {
    enterDeepSleep();
  }
}

// BACK long press (1.5 s) → OFF
unsigned long backHoldStartMs = 0;
bool backHoldActive = false;

void handleBackLongPress() {
  int state = digitalRead(PIN_BTN_BACK);
  unsigned long now = millis();

  if (state == LOW) {
    if (!backHoldActive) {
      backHoldActive = true;
      backHoldStartMs = now;
    } else {
      if (now - backHoldStartMs >= 1500) {
        enterDeepSleep();
      }
    }
  } else {
    backHoldActive = false;
  }
}

// ========== BOOT & CALIBRATION ==========
void showBootScreen() {
  oled.clearDisplay();
  oled.setTextColor(SSD1306_WHITE);

  oled.setTextSize(2);
  oled.setCursor(8, 10);
  oled.print("ESP");
  oled.setCursor(8, 32);
  oled.print("COACH PRO");
  oled.display();

  int x = 10, y = 52, w = 108, h = 8;
  oled.drawRoundRect(x, y, w, h, 3, SSD1306_WHITE);
  oled.display();

  for (int i = 0; i <= w - 4; i += 4) {
    oled.fillRect(x + 2, y + 2, i, h - 4, SSD1306_WHITE);
    oled.display();
    delay(70);
  }
}

void calibrateIMUWithMessage() {
  oled.clearDisplay();
  oled.setTextSize(1);
  oled.setCursor(0, 8);
  oled.print("Hold still for");
  oled.setCursor(0, 18);
  oled.print("calibration");

  int x = 10, y = 36, w = 108, h = 10;
  oled.drawRoundRect(x, y, w, h, 3, SSD1306_WHITE);
  oled.display();

  for (int i = 0; i <= w - 4; i += 12) {
    oled.fillRect(x + 2, y + 2, i, h - 4, SSD1306_WHITE);
    oled.display();
    delay(30);
  }

  mpu.calcGyroOffsets(true);
}

// ========== SENSORS & STROKES ==========
bool updateMPUAndStroke() {
  mpu.update();
  float ay = mpu.getAccY();
  bool strokeDetected = false;

  if (waitingForPeak) {
    if (ay > thresholdHigh) {
      unsigned long now = millis();
      if (lastStrokeMillis != 0) {
        unsigned long dt = now - lastStrokeMillis;
        if (dt > 250 && dt < 4000) {
          float spm = 60000.0f / dt;
          strokeRateSPM = spm;
          addSPM(spm);
        }
      }
      lastStrokeMillis = now;
      strokeCount++;
      strokeDetected = true;

      if (minuteWindowStartMs == 0) minuteWindowStartMs = now;
      if (now - minuteWindowStartMs >= 60000) {
        lastMinuteStrokes = minuteStrokeCount;
        minuteStrokeCount = 0;
        minuteWindowStartMs = now;
      }
      minuteStrokeCount++;

      waitingForPeak = false;
      recordActivity();
    }
  } else {
    if (ay < thresholdLow) {
      waitingForPeak = true;
    }
  }

  return strokeDetected;
}

void updateRowData() {
  unsigned long now = millis();
  if (sessionState == STATE_RUNNING) {
    rowData.elapsedMs = now - sessionStartMillis;
  } else {
    rowData.elapsedMs = 0;
  }
  rowData.t_s      = rowData.elapsedMs / 1000.0f;
  rowData.spmAvg   = avgSPM();
  rowData.tempC    = mpu.getTemp();
  rowData.vbat     = readBatteryVoltage();
  rowData.battPct  = batteryPercent(rowData.vbat);
}

// ========== LOGGING ==========
void logDataIfNeeded(const char* modeName) {
  if (!logFile || !loggingEnabled) return;
  unsigned long now = millis();
  if (now - lastLogMillis < 1000) return;
  lastLogMillis = now;

  logFile.print(rowData.t_s, 1);     logFile.print(",");
  logFile.print(modeName);          logFile.print(",");
  logFile.print(rowData.spmAvg, 1); logFile.print(",");
  logFile.print(strokeCount);       logFile.print(",");
  logFile.print(lastMinuteStrokes); logFile.print(",");
  logFile.println(rowData.tempC, 1);
}

// ========== SESSION RESET ==========
void resetSessionCounters() {
  strokeCount          = 0;
  strokeRateSPM        = 0.0f;
  waitingForPeak       = true;
  lastStrokeMillis     = 0;
  spmIdx               = 0;
  spmFilled            = false;
  minuteWindowStartMs  = 0;
  minuteStrokeCount    = 0;
  lastMinuteStrokes    = 0;
}

void startNewSession() {
  resetSessionCounters();
  sessionState        = STATE_IDLE;
  sessionStartMillis  = 0;
  workoutFinished     = false;
  currentStepIdx      = 0;
  stepStartMs         = 0;
}

// ========== UI HELPERS ==========
void drawBatteryBarSmall() {
  int pct = rowData.battPct;
  int w = 20;
  int h = 6;
  int x = 128 - w - 2;
  int y = 2;
  oled.drawRect(x, y, w, h, SSD1306_WHITE);
  int fill = (pct * (w - 2)) / 100;
  oled.fillRect(x + 1, y + 1, fill, h - 2, SSD1306_WHITE);
}

// ========== UI: MAIN MENU ==========
void drawMenu(int arrowOffset = 0) {
  oled.clearDisplay();
  oled.setTextSize(1);
  oled.setTextColor(SSD1306_WHITE);

  oled.drawRoundRect(0, 0, 128, 64, 6, SSD1306_WHITE);
  oled.setCursor(4, 2);
  oled.print("ESP COACH PRO");

  for (uint8_t i = 0; i < MENU_VISIBLE_ITEMS; i++) {
    uint8_t itemIndex = menuTopIndex + i;
    if (itemIndex >= MENU_ITEMS) break;

    int y = 14 + 12 * i;

    if (itemIndex == menuIndex) {
      oled.fillRoundRect(2, y - 1, 124, 10, 3, SSD1306_WHITE);
      oled.setTextColor(SSD1306_BLACK);
    } else {
      oled.setTextColor(SSD1306_WHITE);
    }

    int xArrow = 6;
    if (itemIndex == menuIndex) {
      xArrow += arrowOffset;
    }

    oled.setCursor(xArrow, y);
    if (itemIndex == menuIndex) {
      oled.print(">");
    } else {
      oled.print(" ");
    }
    oled.print(menuItems[itemIndex]);
  }

  oled.setTextColor(SSD1306_WHITE);
  oled.display();
}

// ========== UI SCREENS ==========
void drawJustRowScreen() {
  oled.clearDisplay();
  oled.setTextColor(SSD1306_WHITE);

  oled.drawRoundRect(0, 0, 128, 64, 6, SSD1306_WHITE);
  drawBatteryBarSmall();

  oled.setTextSize(1);
  oled.setCursor(4, 2);
  oled.print("JUST ROW");

  if (sessionState == STATE_IDLE) {
    oled.setCursor(4, 20);
    oled.print("UP: arm");
    oled.setCursor(4, 32);
    oled.print("First stroke starts");
    oled.setCursor(4, 44);
    oled.print("clock & counting");
    oled.display();
    return;
  }

  if (sessionState == STATE_ARMED) {
    oled.setCursor(4, 20);
    oled.print("ARMED...");
    oled.setCursor(4, 32);
    oled.print("Row to begin");
    oled.display();
    return;
  }

  char timeBuf[8];
  formatTime(timeBuf, sizeof(timeBuf), rowData.elapsedMs);

  oled.setTextSize(2);
  oled.setCursor(8, 16);
  oled.print("R");
  oled.print((int)(rowData.spmAvg + 0.5f));

  oled.setTextSize(1);
  oled.setCursor(80, 16);
  oled.print("MIN:");
  oled.setCursor(80, 26);
  oled.print(lastMinuteStrokes);

  oled.setCursor(80, 38);
  oled.print("CNT:");
  oled.setCursor(80, 48);
  oled.print((int)strokeCount);

  oled.setCursor(4, 52);
  oled.print("TIME ");
  oled.print(timeBuf);

  oled.display();
}

void drawRace1500Screen() {
  oled.clearDisplay();
  oled.setTextColor(SSD1306_WHITE);

  oled.drawRoundRect(0, 0, 128, 64, 6, SSD1306_WHITE);
  drawBatteryBarSmall();

  oled.setTextSize(1);
  oled.setCursor(4, 2);
  oled.print("RACE 1500");

  if (sessionState == STATE_IDLE) {
    oled.setCursor(4, 20);
    oled.print("UP: arm");
    oled.setCursor(4, 32);
    oled.print("Row to start race");
    oled.display();
    return;
  }
  if (sessionState == STATE_ARMED) {
    oled.setCursor(4, 20);
    oled.print("ARMED...");
    oled.setCursor(4, 32);
    oled.print("First stroke = GO");
    oled.display();
    return;
  }

  char timeBuf[8];
  formatTime(timeBuf, sizeof(timeBuf), rowData.elapsedMs);

  oled.setTextSize(2);
  oled.setCursor(8, 20);
  oled.print("R");
  oled.print((int)(rowData.spmAvg + 0.5f));

  oled.setTextSize(1);
  oled.setCursor(80, 20);
  oled.print("T:");
  oled.print(timeBuf);

  oled.setCursor(80, 32);
  oled.print("STK:");
  oled.print((int)strokeCount);

  oled.display();
}

void drawRace2000Screen() {
  oled.clearDisplay();
  oled.setTextColor(SSD1306_WHITE);

  oled.drawRoundRect(0, 0, 128, 64, 6, SSD1306_WHITE);
  drawBatteryBarSmall();

  oled.setTextSize(1);
  oled.setCursor(4, 2);
  oled.print("RACE 2000");

  if (sessionState == STATE_IDLE) {
    oled.setCursor(4, 20);
    oled.print("UP: arm");
    oled.setCursor(4, 32);
    oled.print("Row to start race");
    oled.display();
    return;
  }
  if (sessionState == STATE_ARMED) {
    oled.setCursor(4, 20);
    oled.print("ARMED...");
    oled.setCursor(4, 32);
    oled.print("First stroke = GO");
    oled.display();
    return;
  }

  char timeBuf[8];
  formatTime(timeBuf, sizeof(timeBuf), rowData.elapsedMs);

  oled.setTextSize(2);
  oled.setCursor(8, 20);
  oled.print("R");
  oled.print((int)(rowData.spmAvg + 0.5f));

  oled.setTextSize(1);
  oled.setCursor(80, 20);
  oled.print("T:");
  oled.print(timeBuf);

  oled.setCursor(80, 32);
  oled.print("STK:");
  oled.print((int)strokeCount);

  oled.display();
}

void drawStrokeCountScreen() {
  oled.clearDisplay();
  oled.setTextColor(SSD1306_WHITE);

  oled.drawRoundRect(0, 0, 128, 64, 6, SSD1306_WHITE);

  oled.setTextSize(1);
  oled.setCursor(4, 2);
  oled.print("STROKE COUNT");

  if (sessionState == STATE_IDLE) {
    oled.setCursor(4, 20);
    oled.print("UP: arm");
    oled.setCursor(4, 32);
    oled.print("Row to begin");
    oled.display();
    return;
  }
  if (sessionState == STATE_ARMED) {
    oled.setCursor(4, 20);
    oled.print("ARMED...");
    oled.setCursor(4, 32);
    oled.print("Row to count");
    oled.display();
    return;
  }

  oled.setTextSize(3);
  oled.setCursor(24, 24);
  oled.print((int)strokeCount);
  oled.display();
}

void drawRateOnlyScreen() {
  oled.clearDisplay();
  oled.setTextColor(SSD1306_WHITE);

  oled.drawRoundRect(0, 0, 128, 64, 6, SSD1306_WHITE);
  drawBatteryBarSmall();

  oled.setTextSize(1);
  oled.setCursor(4, 2);
  oled.print("RATE ONLY");

  if (sessionState == STATE_IDLE) {
    oled.setCursor(4, 20);
    oled.print("UP: arm");
    oled.setCursor(4, 32);
    oled.print("Row to see rate");
    oled.display();
    return;
  }
  if (sessionState == STATE_ARMED) {
    oled.setCursor(4, 20);
    oled.print("ARMED...");
    oled.setCursor(4, 32);
    oled.print("Row to start");
    oled.display();
    return;
  }

  oled.setTextSize(4);
  oled.setCursor(16, 18);
  oled.print((int)(rowData.spmAvg + 0.5f));
  oled.display();
}

// WORKOUT MENU
void drawWorkoutMenu(int arrowOffset = 0) {
  oled.clearDisplay();
  oled.setTextSize(1);
  oled.setTextColor(SSD1306_WHITE);

  oled.drawRoundRect(0, 0, 128, 64, 6, SSD1306_WHITE);
  oled.setCursor(4, 2);
  oled.print("WORKOUTS (PRO)");

  for (uint8_t i = 0; i < WORKOUT_VISIBLE_ITEMS; i++) {
    uint8_t itemIndex = workoutTopIndex + i;
    if (itemIndex >= WORKOUT_PLAN_COUNT) break;

    int y = 14 + 12 * i;

    if (itemIndex == workoutMenuIndex) {
      oled.fillRoundRect(2, y - 1, 124, 10, 3, SSD1306_WHITE);
      oled.setTextColor(SSD1306_BLACK);
    } else {
      oled.setTextColor(SSD1306_WHITE);
    }

    int xArrow = 6;
    if (itemIndex == workoutMenuIndex) {
      xArrow += arrowOffset;
    }

    oled.setCursor(xArrow, y);
    if (itemIndex == workoutMenuIndex) {
      oled.print(">");
    } else {
      oled.print(" ");
    }
    oled.print(workoutPlans[itemIndex].name);
  }

  oled.setTextColor(SSD1306_WHITE);
  oled.display();
}

void drawWorkoutRunScreen() {
  oled.clearDisplay();
  oled.setTextColor(SSD1306_WHITE);

  oled.drawRoundRect(0, 0, 128, 64, 6, SSD1306_WHITE);
  drawBatteryBarSmall();

  if (!currentPlan) {
    oled.setTextSize(1);
    oled.setCursor(4, 24);
    oled.print("No plan");
    oled.display();
    return;
  }

  oled.setTextSize(1);
  oled.setCursor(4, 2);
  oled.print(currentPlan->name);

  if (sessionState == STATE_IDLE) {
    oled.setCursor(4, 20);
    oled.print("UP: arm   ENTER:run");
    oled.setCursor(4, 32);
    oled.print("First stroke starts");
    oled.display();
    return;
  }

  if (sessionState == STATE_ARMED) {
    oled.setCursor(4, 20);
    oled.print("ARMED...");
    oled.setCursor(4, 32);
    oled.print("Row to start set");
    oled.display();
    return;
  }

  if (workoutFinished) {
    oled.setCursor(4, 20);
    oled.print("Workout DONE");
    oled.setCursor(4, 32);
    oled.print("BACK: menu");
    oled.display();
    return;
  }

  char timeBuf[8];
  formatTime(timeBuf, sizeof(timeBuf), rowData.elapsedMs);

  const WorkoutStep& step = currentPlan->steps[currentStepIdx];
  unsigned long now = millis();
  unsigned long stepElapsed = (now - stepStartMs) / 1000;
  unsigned long stepTotal   = step.duration_s;
  unsigned long stepRemain  = (stepTotal > stepElapsed) ? (stepTotal - stepElapsed) : 0;

  oled.setCursor(4, 14);
  oled.print("Step ");
  oled.print((int)(currentStepIdx + 1));
  oled.print("/");
  oled.print((int)currentPlan->numSteps);

  oled.setCursor(4, 24);
  oled.print(step.isWork ? "Work " : "Rest ");
  oled.print(step.label);

  oled.setCursor(4, 34);
  oled.print("Tot ");
  oled.print(timeBuf);

  oled.setCursor(4, 44);
  oled.print("Step ");
  oled.print(stepElapsed);
  oled.print("/");
  oled.print(stepTotal);
  oled.print("s (");
  oled.print(stepRemain);
  oled.print("s)");

  oled.setCursor(4, 54);
  oled.print("R ");
  oled.print((int)(rowData.spmAvg + 0.5f));
  oled.print("  ST ");
  oled.print((int)strokeCount);

  oled.display();
}

// SETTINGS UI
void drawSettingsMenu(int arrowOffset = 0) {
  oled.clearDisplay();
  oled.setTextSize(1);
  oled.setTextColor(SSD1306_WHITE);

  oled.drawRoundRect(0, 0, 128, 64, 6, SSD1306_WHITE);

  oled.setCursor(4, 2);
  oled.print("SETTINGS (PRO)");

  for (uint8_t i = 0; i < SETTINGS_VISIBLE_ITEMS; i++) {
    uint8_t itemIndex = settingsTopIndex + i;
    if (itemIndex >= SETTINGS_ITEMS) break;

    int y = 14 + 12 * i;

    if (itemIndex == settingsIndex) {
      oled.fillRoundRect(2, y - 1, 124, 10, 3, SSD1306_WHITE);
      oled.setTextColor(SSD1306_BLACK);
    } else {
      oled.setTextColor(SSD1306_WHITE);
    }

    int xArrow = 6;
    if (itemIndex == settingsIndex) {
      xArrow += arrowOffset;
    }

    oled.setCursor(xArrow, y);
    if (itemIndex == settingsIndex) {
      oled.print(">");
    } else {
      oled.print(" ");
    }
    oled.print(settingsItems[itemIndex]);
  }

  oled.setTextColor(SSD1306_WHITE);
  oled.display();
}

void drawSettingsSens() {
  oled.clearDisplay();
  oled.setTextSize(1);
  oled.setTextColor(SSD1306_WHITE);

  oled.drawRoundRect(0, 0, 128, 64, 6, SSD1306_WHITE);

  oled.setCursor(4, 2);
  oled.print("Sensitivity");

  oled.setTextSize(2);
  oled.setCursor(4, 24);
  oled.print("Lvl ");
  oled.print(sensitivityLevel);

  oled.setTextSize(1);
  oled.setCursor(4, 52);
  oled.print("UP/DN change  BACK=exit");

  oled.display();
}

void drawSettingsCalib() {
  oled.clearDisplay();
  oled.setTextSize(1);
  oled.setTextColor(SSD1306_WHITE);

  oled.drawRoundRect(0, 0, 128, 64, 6, SSD1306_WHITE);

  oled.setCursor(4, 16);
  oled.print("ENTER: calibrate IMU");
  oled.setCursor(4, 28);
  oled.print("BACK: exit");
  oled.display();
}

void drawSettingsAutoOff() {
  oled.clearDisplay();
  oled.setTextSize(1);
  oled.setTextColor(SSD1306_WHITE);

  oled.drawRoundRect(0, 0, 128, 64, 6, SSD1306_WHITE);

  oled.setCursor(4, 2);
  oled.print("Auto-off");

  oled.setTextSize(2);
  oled.setCursor(4, 24);
  oled.print(autoOffMinutes[autoOffIndex]);
  oled.print(" min");

  oled.setTextSize(1);
  oled.setCursor(4, 52);
  oled.print("UP/DN change  BACK=exit");

  oled.display();
}

void drawSettingsLogging() {
  oled.clearDisplay();
  oled.setTextSize(1);
  oled.setTextColor(SSD1306_WHITE);

  oled.drawRoundRect(0, 0, 128, 64, 6, SSD1306_WHITE);

  oled.setCursor(4, 2);
  oled.print("Logging to SD");

  oled.setTextSize(2);
  oled.setCursor(4, 24);
  oled.print(loggingEnabled ? "ON " : "OFF");

  oled.setTextSize(1);
  oled.setCursor(4, 52);
  oled.print("UP/DN/ENT toggles");

  oled.display();
}

// ========== GENERIC SIMPLE ROW HANDLER ==========
void handleSimpleRowMode(const char* modeName, void (*drawFn)()) {
  if (btnBack.pressed) {
    currentMode = MODE_MENU;
    startNewSession();
    return;
  }

  if (sessionState == STATE_IDLE) {
    if (btnUp.pressed) {
      sessionState = STATE_ARMED;
    }
    updateRowData();
    drawFn();
    return;
  }

  if (sessionState == STATE_ARMED) {
    bool stroke = updateMPUAndStroke();
    if (stroke) {
      sessionStartMillis = lastStrokeMillis;
      sessionState       = STATE_RUNNING;
    }
    updateRowData();
    drawFn();
    return;
  }

  updateMPUAndStroke();
  updateRowData();
  drawFn();
  logDataIfNeeded(modeName);
}

// ========== MODE HANDLERS ==========
void handleMenu() {
  int8_t prevIndex = menuIndex;

  if (btnUp.pressed   && menuIndex > 0)              menuIndex--;
  if (btnDown.pressed && menuIndex < MENU_ITEMS - 1) menuIndex++;

  if (menuIndex < menuTopIndex) {
    menuTopIndex = menuIndex;
  } else if (menuIndex > menuTopIndex + MENU_VISIBLE_ITEMS - 1) {
    menuTopIndex = menuIndex - (MENU_VISIBLE_ITEMS - 1);
  }

  if (prevIndex != menuIndex) {
    for (int step = 0; step < 4; step++) {
      int offset = step * 2;
      drawMenu(offset);
      delay(20);
    }
  }

  if (btnEnter.pressed) {
    startNewSession();
    switch (menuIndex) {
      case 0: currentMode = MODE_JUST_ROW;      break;
      case 1: currentMode = MODE_RACE_1500;     break;
      case 2: currentMode = MODE_RACE_2000;     break;
      case 3: currentMode = MODE_WORKOUT_MENU;  break;
      case 4: currentMode = MODE_STROKE_CT;     break;
      case 5: currentMode = MODE_RATE_ONLY;     break;
      case 6: currentMode = MODE_SETTINGS_MENU; break;
    }
  }

  drawMenu();
}

void handleJustRow()       { handleSimpleRowMode("JUST",      drawJustRowScreen); }
void handleRace1500()      { handleSimpleRowMode("RACE1500",  drawRace1500Screen); }
void handleRace2000()      { handleSimpleRowMode("RACE2000",  drawRace2000Screen); }
void handleStrokeCount()   { handleSimpleRowMode("STROKECT",  drawStrokeCountScreen); }
void handleRateOnly()      { handleSimpleRowMode("RATEONLY",  drawRateOnlyScreen); }

void handleWorkoutMenu() {
  int8_t prev = workoutMenuIndex;
  if (btnUp.pressed   && workoutMenuIndex > 0)                      workoutMenuIndex--;
  if (btnDown.pressed && workoutMenuIndex < WORKOUT_PLAN_COUNT-1)   workoutMenuIndex++;

  if (workoutMenuIndex < workoutTopIndex) {
    workoutTopIndex = workoutMenuIndex;
  } else if (workoutMenuIndex > workoutTopIndex + WORKOUT_VISIBLE_ITEMS - 1) {
    workoutTopIndex = workoutMenuIndex - (WORKOUT_VISIBLE_ITEMS - 1);
  }

  if (prev != workoutMenuIndex) {
    for (int step = 0; step < 4; step++) {
      int offset = step * 2;
      drawWorkoutMenu(offset);
      delay(20);
    }
  }

  if (btnEnter.pressed) {
    currentPlan   = &workoutPlans[workoutMenuIndex];
    startNewSession();
    currentMode   = MODE_WORKOUT_RUN;
  }
  if (btnBack.pressed) {
    currentMode = MODE_MENU;
    return;
  }

  drawWorkoutMenu();
}

void handleWorkoutRun() {
  if (!currentPlan) {
    currentMode = MODE_WORKOUT_MENU;
    return;
  }

  if (btnBack.pressed) {
    currentMode = MODE_MENU;
    startNewSession();
    return;
  }

  if (sessionState == STATE_IDLE) {
    if (btnUp.pressed || btnEnter.pressed) {
      sessionState = STATE_ARMED;
    }
    updateRowData();
    drawWorkoutRunScreen();
    return;
  }

  if (sessionState == STATE_ARMED) {
    bool stroke = updateMPUAndStroke();
    if (stroke) {
      sessionStartMillis = lastStrokeMillis;
      sessionState       = STATE_RUNNING;
      currentStepIdx     = 0;
      stepStartMs        = millis();
    }
    updateRowData();
    drawWorkoutRunScreen();
    return;
  }

  if (workoutFinished) {
    updateRowData();
    drawWorkoutRunScreen();
    logDataIfNeeded("WORKOUT");
    return;
  }

  updateMPUAndStroke();
  updateRowData();

  unsigned long now = millis();
  const WorkoutStep& step = currentPlan->steps[currentStepIdx];
  unsigned long stepElapsed = (now - stepStartMs) / 1000;

  if (stepElapsed >= step.duration_s) {
    currentStepIdx++;
    if (currentStepIdx >= currentPlan->numSteps) {
      workoutFinished = true;
    } else {
      stepStartMs = now;
    }
  }

  drawWorkoutRunScreen();
  logDataIfNeeded("WORKOUT");
}

void handleSettingsMenu() {
  int8_t prev = settingsIndex;
  if (btnUp.pressed   && settingsIndex > 0)                 settingsIndex--;
  if (btnDown.pressed && settingsIndex < SETTINGS_ITEMS-1)  settingsIndex++;

  if (settingsIndex < settingsTopIndex) {
    settingsTopIndex = settingsIndex;
  } else if (settingsIndex > settingsTopIndex + SETTINGS_VISIBLE_ITEMS - 1) {
    settingsTopIndex = settingsIndex - (SETTINGS_VISIBLE_ITEMS - 1);
  }

  if (prev != settingsIndex) {
    for (int step = 0; step < 4; step++) {
      int offset = step * 2;
      drawSettingsMenu(offset);
      delay(20);
    }
  }

  if (btnEnter.pressed) {
    switch (settingsIndex) {
      case 0: currentMode = MODE_SETTINGS_SENS;      break;
      case 1: currentMode = MODE_SETTINGS_CALIB;     break;
      case 2: currentMode = MODE_SETTINGS_AUTO_OFF;  break;
      case 3: currentMode = MODE_SETTINGS_LOGGING;   break;
    }
    return;
  }

  if (btnBack.pressed) {
    currentMode = MODE_MENU;
    return;
  }

  drawSettingsMenu();
}

void handleSettingsSens() {
  bool changed = false;

  if (btnUp.pressed && sensitivityLevel < 2) {
    sensitivityLevel++;
    applySensitivity();
    changed = true;
  }
  if (btnDown.pressed && sensitivityLevel > 0) {
    sensitivityLevel--;
    applySensitivity();
    changed = true;
  }
  if (changed) {
    saveSettingsPro();
  }

  if (btnBack.pressed) {
    saveSettingsPro();
    currentMode = MODE_SETTINGS_MENU;
    return;
  }

  drawSettingsSens();
}

void handleSettingsCalib() {
  if (btnEnter.pressed) {
    calibrateIMUWithMessage();
  }
  if (btnBack.pressed) {
    currentMode = MODE_SETTINGS_MENU;
    return;
  }
  drawSettingsCalib();
}

void handleSettingsAutoOff() {
  bool changed = false;

  if (btnUp.pressed && autoOffIndex < AUTO_OFF_CHOICES - 1) {
    autoOffIndex++;
    updateAutoOffMs();
    changed = true;
  }
  if (btnDown.pressed && autoOffIndex > 0) {
    autoOffIndex--;
    updateAutoOffMs();
    changed = true;
  }
  if (changed) {
    saveSettingsPro();
  }

  if (btnBack.pressed) {
    saveSettingsPro();
    currentMode = MODE_SETTINGS_MENU;
    return;
  }

  drawSettingsAutoOff();
}

void handleSettingsLogging() {
  if (btnUp.pressed || btnDown.pressed || btnEnter.pressed) {
    loggingEnabled = !loggingEnabled;
    saveSettingsPro();
  }
  if (btnBack.pressed) {
    saveSettingsPro();
    currentMode = MODE_SETTINGS_MENU;
    return;
  }

  drawSettingsLogging();
}

// ========== SETUP & LOOP ==========
void setup() {
  Serial.begin(115200);

  pinMode(PIN_BTN_UP,    INPUT_PULLUP);
  pinMode(PIN_BTN_DOWN,  INPUT_PULLUP);
  pinMode(PIN_BTN_ENTER, INPUT_PULLUP);
  pinMode(PIN_BTN_BACK,  INPUT_PULLUP);

  Wire.begin(PIN_I2C_SDA, PIN_I2C_SCL);

  if (!oled.begin(SSD1306_SWITCHCAPVCC, 0x3C)) {
    Serial.println("SSD1306 failed");
    while (1) { delay(10); }
  }
  oled.clearDisplay();
  oled.display();

  mpu.begin();
  applySensitivity();

  if (!SD.begin(PIN_SD_CS)) {
    Serial.println("SD init failed!");
  } else {
    SD.mkdir("/ESP_COACH");
    SD.mkdir(DATA_DIR_PRO);

    loadSettingsPro();
    applySensitivity();
    updateAutoOffMs();

    logFile = SD.open(LOG_FILE_PRO, FILE_WRITE);
    if (logFile) {
      logFile.println("time_s,mode,spm,strokes,lastMin,tempC");
    }
  }

  analogReadResolution(12);

  showBootScreen();
  calibrateIMUWithMessage();

  startNewSession();
  recordActivity();
  drawMenu();
}

void loop() {
  readButtons();
  handleBackLongPress();

  switch (currentMode) {
    case MODE_MENU:             handleMenu();            break;
    case MODE_JUST_ROW:         handleJustRow();         break;
    case MODE_RACE_1500:        handleRace1500();        break;
    case MODE_RACE_2000:        handleRace2000();        break;
    case MODE_WORKOUT_MENU:     handleWorkoutMenu();     break;
    case MODE_WORKOUT_RUN:      handleWorkoutRun();      break;
    case MODE_STROKE_CT:        handleStrokeCount();     break;
    case MODE_RATE_ONLY:        handleRateOnly();        break;
    case MODE_SETTINGS_MENU:    handleSettingsMenu();    break;
    case MODE_SETTINGS_SENS:    handleSettingsSens();    break;
    case MODE_SETTINGS_CALIB:   handleSettingsCalib();   break;
    case MODE_SETTINGS_AUTO_OFF:handleSettingsAutoOff(); break;
    case MODE_SETTINGS_LOGGING: handleSettingsLogging(); break;
  }

  checkAutoOff();
}