#include "ssd1306.h"
#include <EEPROM.h>

// Button pins
#define BTN_UP 2
#define BTN_DOWN 3
#define BTN_SELECT 4
#define BTN_BACK 5

// EEPROM addresses
#define EEPROM_KP 0
#define EEPROM_KI 4
#define EEPROM_KD 8
#define EEPROM_SPEED 12
#define EEPROM_DEVICE_NAME 16
#define EEPROM_MAGIC 500

// Configuration constants
#define DEVICE_NAME_MAX_LEN 15
#define MAGIC_NUMBER 0xAB
#define BUTTON_DEBOUNCE_MS 150
#define DISPLAY_UPDATE_MS 50
#define STARTUP_DELAY_MS 2000
#define SAVE_CONFIRM_DELAY_MS 1500

// Parameter limits
#define MIN_FLOAT_PARAM 0.0f
#define MAX_FLOAT_PARAM 10.0f
#define MIN_SPEED 0
#define MAX_SPEED 255
#define DEFAULT_SPEED 150

// Menu states
enum MenuState : uint8_t {
  MAIN_MENU,
  LINE_FOLLOWER_MENU,
  PARAMETER_MENU,
  DEVICE_NAME_MENU,
  RUNNING_STATE
};

// Parameter types for generic adjustment
enum ParamType : uint8_t {
  PARAM_KP,
  PARAM_KI,
  PARAM_KD,
  PARAM_SPEED
};

// Menu items stored in PROGMEM to save RAM
const char mainMenuItem0[] PROGMEM = "Line Follower";
const char mainMenuItem1[] PROGMEM = "Parameters";
const char mainMenuItem2[] PROGMEM = "Device Name";
const char mainMenuItem3[] PROGMEM = "Start";

const char* const mainMenuItems[] PROGMEM = {
  mainMenuItem0, mainMenuItem1, mainMenuItem2, mainMenuItem3
};

const char paramMenuItem0[] PROGMEM = "Kp";
const char paramMenuItem1[] PROGMEM = "Ki";
const char paramMenuItem2[] PROGMEM = "Kd";
const char paramMenuItem3[] PROGMEM = "Speed";
const char paramMenuItem4[] PROGMEM = "Save EEPROM";

const char* const paramMenuItems[] PROGMEM = {
  paramMenuItem0, paramMenuItem1, paramMenuItem2, paramMenuItem3, paramMenuItem4
};

// Global state
struct SystemState {
  MenuState currentState;
  uint8_t currentMenuItem;
  uint8_t maxMenuItem;
  unsigned long lastButtonPress;
  bool needsDisplayUpdate;  // Flag untuk update tampilan
} state = {MAIN_MENU, 0, 3, 0, true};

// PID parameters structure for better organization
struct PIDParams {
  float kp;
  float ki;
  float kd;
  int baseSpeed;
} params = {1.0f, 0.0f, 0.1f, DEFAULT_SPEED};

// Device configuration
char deviceName[DEVICE_NAME_MAX_LEN + 1] = "LineBot-001";

// Button handling with improved debouncing
inline bool isButtonPressed(uint8_t pin) {
  static uint8_t lastPin = 0;
  static unsigned long lastPressTime = 0;
  
  if (digitalRead(pin) == LOW) {
    unsigned long currentTime = millis();
    if (pin != lastPin || (currentTime - lastPressTime) > BUTTON_DEBOUNCE_MS) {
      lastPin = pin;
      lastPressTime = currentTime;
      return true;
    }
  }
  return false;
}

// Helper function to read string from PROGMEM
void readProgmemString(const char* const* table, uint8_t index, char* buffer, uint8_t maxLen) {
  strcpy_P(buffer, (char*)pgm_read_word(&(table[index])));
}

// Optimized parameter validation
inline float validateFloat(float value, float minVal = MIN_FLOAT_PARAM, float maxVal = MAX_FLOAT_PARAM) {
  return constrain(value, minVal, maxVal);
}

inline int validateInt(int value, int minVal = MIN_SPEED, int maxVal = MAX_SPEED) {
  return constrain(value, minVal, maxVal);
}

void setup() {
  Serial.begin(115200); // Increased baud rate for better performance
  Serial.println(F("Starting LineBot OS..."));
  
  // Initialize display
  ssd1306_128x64_i2c_init();
  ssd1306_clearScreen();
  ssd1306_setFixedFont(ssd1306xled_font6x8);
  ssd1306_printFixed(10, 20, "Display OK!", STYLE_NORMAL);
  delay(1000);
  
  // Initialize buttons (INPUT_PULLUP is more efficient than external resistors)
  const uint8_t buttons[] = {BTN_UP, BTN_DOWN, BTN_SELECT, BTN_BACK};
  for (uint8_t i = 0; i < 4; i++) {
    pinMode(buttons[i], INPUT_PULLUP);
  }
  
  Serial.println(F("Loading parameters..."));
  loadFromEEPROM();
  
  showStartupScreen();
  delay(STARTUP_DELAY_MS);
  
  Serial.println(F("LineBot OS Ready!"));
}

void loop() {
  static unsigned long lastUpdate = 0;
  unsigned long currentTime = millis();
  
  // Handle buttons every loop for responsiveness
  handleButtons();
  
  // Update display at fixed interval to reduce flicker
  if (currentTime - lastUpdate >= DISPLAY_UPDATE_MS) {
    updateDisplay();
    lastUpdate = currentTime;
  }
}

void handleButtons() {
  if (isButtonPressed(BTN_UP)) {
    navigateUp();
  }
  else if (isButtonPressed(BTN_DOWN)) {
    navigateDown();
  }
  else if (isButtonPressed(BTN_SELECT)) {
    selectItem();
  }
  else if (isButtonPressed(BTN_BACK)) {
    goBack();
  }
}

inline void navigateUp() {
  state.currentMenuItem = (state.currentMenuItem == 0) ? state.maxMenuItem : state.currentMenuItem - 1;
}

inline void navigateDown() {
  state.currentMenuItem = (state.currentMenuItem == state.maxMenuItem) ? 0 : state.currentMenuItem + 1;
}

void selectItem() {
  switch (state.currentState) {
    case MAIN_MENU:
      handleMainMenuSelect();
      break;
    case PARAMETER_MENU:
      handleParameterMenuSelect();
      break;
    case LINE_FOLLOWER_MENU:
      break;
    case DEVICE_NAME_MENU:
      break;
  }
}

void goBack() {
  if (state.currentState != MAIN_MENU) {
    state.currentState = MAIN_MENU;
    state.currentMenuItem = 0;
    state.maxMenuItem = 3;
  }
}

void handleMainMenuSelect() {
  switch (state.currentMenuItem) {
    case 0: // Line Follower
      state.currentState = LINE_FOLLOWER_MENU;
      state.currentMenuItem = 0;
      state.maxMenuItem = 0;
      break;
    case 1: // Parameters
      state.currentState = PARAMETER_MENU;
      state.currentMenuItem = 0;
      state.maxMenuItem = 4;
      break;
    case 2: // Device Name
      state.currentState = DEVICE_NAME_MENU;
      state.currentMenuItem = 0;
      state.maxMenuItem = 0;
      break;
    case 3: // Start
      state.currentState = RUNNING_STATE;
      startLineFollower();
      break;
  }
}

void handleParameterMenuSelect() {
  switch (state.currentMenuItem) {
    case 0: // Kp
      adjustFloatParameter(&params.kp, PSTR("Kp"), 0.1f);
      break;
    case 1: // Ki
      adjustFloatParameter(&params.ki, PSTR("Ki"), 0.01f);
      break;
    case 2: // Kd
      adjustFloatParameter(&params.kd, PSTR("Kd"), 0.01f);
      break;
    case 3: // Speed
      adjustSpeedParameter();
      break;
    case 4: // Save to EEPROM
      saveToEEPROM();
      showSaveConfirmation();
      break;
  }
}

// Generic parameter adjustment function
void adjustFloatParameter(float* param, const char* paramName, float increment) {
  char nameBuffer[8];
  strcpy_P(nameBuffer, paramName);
  
  while (true) {
    ssd1306_clearScreen();
    ssd1306_setFixedFont(ssd1306xled_font6x8);
    
    // Display header
    ssd1306_printFixed(0, 0, "Adjust ", STYLE_NORMAL);
    ssd1306_printFixed(42, 0, nameBuffer, STYLE_NORMAL);
    
    // Display instructions
    ssd1306_printFixed(0, 16, "UP/DOWN: Change", STYLE_NORMAL);
    ssd1306_printFixed(0, 24, "SELECT: Save", STYLE_NORMAL);
    ssd1306_printFixed(0, 32, "BACK: Cancel", STYLE_NORMAL);
    
    // Display current value
    char valueStr[10];
    dtostrf(*param, 6, 3, valueStr);
    ssd1306_printFixed(0, 48, "Value: ", STYLE_NORMAL);
    ssd1306_printFixed(42, 48, valueStr, STYLE_NORMAL);
    
    // Handle input
    if (isButtonPressed(BTN_UP)) {
      *param = validateFloat(*param + increment);
    }
    else if (isButtonPressed(BTN_DOWN)) {
      *param = validateFloat(*param - increment);
    }
    else if (isButtonPressed(BTN_SELECT) || isButtonPressed(BTN_BACK)) {
      break;
    }
    
    delay(50); // Reduced delay for better responsiveness
  }
}

void adjustSpeedParameter() {
  while (true) {
    ssd1306_clearScreen();
    ssd1306_setFixedFont(ssd1306xled_font6x8);
    
    ssd1306_printFixed(0, 0, "Adjust Speed", STYLE_NORMAL);
    ssd1306_printFixed(0, 16, "UP/DOWN: Change", STYLE_NORMAL);
    ssd1306_printFixed(0, 24, "SELECT: Save", STYLE_NORMAL);
    ssd1306_printFixed(0, 32, "BACK: Cancel", STYLE_NORMAL);
    
    // Display current speed
    char speedStr[6];
    itoa(params.baseSpeed, speedStr, 10);
    ssd1306_printFixed(0, 48, "Speed: ", STYLE_NORMAL);
    ssd1306_printFixed(42, 48, speedStr, STYLE_NORMAL);
    
    if (isButtonPressed(BTN_UP)) {
      params.baseSpeed = validateInt(params.baseSpeed + 10);
    }
    else if (isButtonPressed(BTN_DOWN)) {
      params.baseSpeed = validateInt(params.baseSpeed - 10);
    }
    else if (isButtonPressed(BTN_SELECT) || isButtonPressed(BTN_BACK)) {
      break;
    }
    
    delay(50);
  }
}

void updateDisplay() {
  ssd1306_clearScreen();
  ssd1306_setFixedFont(ssd1306xled_font6x8);
  
  switch (state.currentState) {
    case MAIN_MENU:
      showMainMenu();
      break;
    case PARAMETER_MENU:
      showParameterMenu();
      break;
    case LINE_FOLLOWER_MENU:
      showLineFollowerMenu();
      break;
    case DEVICE_NAME_MENU:
      showDeviceNameMenu();
      break;
    case RUNNING_STATE:
      showRunningScreen();
      break;
  }
}

void showStartupScreen() {
  ssd1306_clearScreen();
  ssd1306_setFixedFont(ssd1306xled_font6x8);
  
  ssd1306_printFixed(5, 5, "LineBot OS", STYLE_BOLD);
  ssd1306_printFixed(10, 25, "Version 1.0", STYLE_NORMAL);
  ssd1306_printFixed(10, 35, deviceName, STYLE_NORMAL);
  ssd1306_printFixed(10, 50, "Starting...", STYLE_NORMAL);
}

void showMainMenu() {
  ssd1306_printFixed(0, 0, "=== MAIN MENU ===", STYLE_BOLD);
  
  char buffer[20];
  for (uint8_t i = 0; i <= state.maxMenuItem; i++) {
    uint8_t y = 16 + i * 10;
    readProgmemString(mainMenuItems, i, buffer, sizeof(buffer));
    
    if (i == state.currentMenuItem) {
      ssd1306_printFixed(0, y, "> ", STYLE_NORMAL);
      ssd1306_printFixed(12, y, buffer, STYLE_BOLD);
    } else {
      ssd1306_printFixed(12, y, buffer, STYLE_NORMAL);
    }
  }
}

void showParameterMenu() {
  ssd1306_printFixed(0, 0, "== PARAMETERS ==", STYLE_BOLD);
  
  char buffer[20];
  char valueStr[10];
  
  for (uint8_t i = 0; i <= state.maxMenuItem; i++) {
    uint8_t y = 12 + i * 9;
    
    // Show cursor
    ssd1306_printFixed(0, y, (i == state.currentMenuItem) ? ">" : " ", STYLE_NORMAL);
    
    if (i < 4) {
      readProgmemString(paramMenuItems, i, buffer, sizeof(buffer));
      ssd1306_printFixed(8, y, buffer, STYLE_NORMAL);
      ssd1306_printFixed(24, y, ":", STYLE_NORMAL);
      
      // Get parameter value
      switch (i) {
        case 0: dtostrf(params.kp, 5, 2, valueStr); break;
        case 1: dtostrf(params.ki, 5, 3, valueStr); break;
        case 2: dtostrf(params.kd, 5, 3, valueStr); break;
        case 3: itoa(params.baseSpeed, valueStr, 10); break;
      }
      ssd1306_printFixed(32, y, valueStr, STYLE_NORMAL);
    } else {
      readProgmemString(paramMenuItems, i, buffer, sizeof(buffer));
      ssd1306_printFixed(8, y, buffer, STYLE_NORMAL);
    }
  }
}

void showLineFollowerMenu() {
  ssd1306_printFixed(0, 0, "= LINE FOLLOWER =", STYLE_BOLD);
  ssd1306_printFixed(0, 16, "Sensor Status:", STYLE_NORMAL);
  ssd1306_printFixed(0, 32, "L   M   R", STYLE_NORMAL);
  ssd1306_printFixed(0, 42, "0   1   0", STYLE_NORMAL);
  ssd1306_printFixed(0, 56, "BACK to return", STYLE_NORMAL);
}

void showDeviceNameMenu() {
  ssd1306_printFixed(0, 0, "== DEVICE NAME ==", STYLE_BOLD);
  ssd1306_printFixed(0, 16, "Name:", STYLE_NORMAL);
  ssd1306_printFixed(0, 26, deviceName, STYLE_NORMAL);
  ssd1306_printFixed(0, 42, "Edit feature", STYLE_NORMAL);
  ssd1306_printFixed(0, 50, "coming soon...", STYLE_NORMAL);
}

void showRunningScreen() {
  ssd1306_printFixed(0, 0, "=== RUNNING ===", STYLE_BOLD);
  
  // Display parameters efficiently
  char tempStr[8];
  
  itoa(params.baseSpeed, tempStr, 10);
  ssd1306_printFixed(0, 16, "Speed: ", STYLE_NORMAL);
  ssd1306_printFixed(42, 16, tempStr, STYLE_NORMAL);
  
  dtostrf(params.kp, 5, 2, tempStr);
  ssd1306_printFixed(0, 26, "Kp:", STYLE_NORMAL);
  ssd1306_printFixed(18, 26, tempStr, STYLE_NORMAL);
  
  dtostrf(params.ki, 5, 3, tempStr);
  ssd1306_printFixed(0, 34, "Ki:", STYLE_NORMAL);
  ssd1306_printFixed(18, 34, tempStr, STYLE_NORMAL);
  
  dtostrf(params.kd, 5, 3, tempStr);
  ssd1306_printFixed(0, 42, "Kd:", STYLE_NORMAL);
  ssd1306_printFixed(18, 42, tempStr, STYLE_NORMAL);
  
  ssd1306_printFixed(0, 54, "BACK to stop", STYLE_NORMAL);
}

void startLineFollower() {
  Serial.println(F("Line Follower Started!"));
  Serial.print(F("Kp: ")); Serial.print(params.kp);
  Serial.print(F(", Ki: ")); Serial.print(params.ki);
  Serial.print(F(", Kd: ")); Serial.print(params.kd);
  Serial.print(F(", Speed: ")); Serial.println(params.baseSpeed);
}

void saveToEEPROM() {
  Serial.println(F("Saving parameters..."));
  
  EEPROM.put(EEPROM_KP, params.kp);
  EEPROM.put(EEPROM_KI, params.ki);
  EEPROM.put(EEPROM_KD, params.kd);
  EEPROM.put(EEPROM_SPEED, params.baseSpeed);
  
  // Save device name with bounds checking
  uint8_t len = min(strlen(deviceName), DEVICE_NAME_MAX_LEN);
  for (uint8_t i = 0; i < DEVICE_NAME_MAX_LEN + 1; i++) {
    EEPROM.write(EEPROM_DEVICE_NAME + i, (i < len) ? deviceName[i] : 0);
  }
  
  Serial.println(F("Parameters saved!"));
}

void loadFromEEPROM() {
  // Check initialization magic number
  if (EEPROM.read(EEPROM_MAGIC) != MAGIC_NUMBER) {
    Serial.println(F("Initializing EEPROM..."));
    saveToEEPROM();
    EEPROM.write(EEPROM_MAGIC, MAGIC_NUMBER);
    return;
  }
  
  // Load and validate parameters
  float tempKp, tempKi, tempKd;
  int tempSpeed;
  
  EEPROM.get(EEPROM_KP, tempKp);
  EEPROM.get(EEPROM_KI, tempKi);
  EEPROM.get(EEPROM_KD, tempKd);
  EEPROM.get(EEPROM_SPEED, tempSpeed);
  
  params.kp = validateFloat(tempKp);
  params.ki = validateFloat(tempKi);
  params.kd = validateFloat(tempKd);
  params.baseSpeed = validateInt(tempSpeed);
  
  // Load device name safely
  for (uint8_t i = 0; i < DEVICE_NAME_MAX_LEN; i++) {
    uint8_t ch = EEPROM.read(EEPROM_DEVICE_NAME + i);
    deviceName[i] = (ch == 255 || ch == 0) ? 0 : ch;
    if (ch == 0) break;
  }
  deviceName[DEVICE_NAME_MAX_LEN] = 0; // Ensure null termination
  
  // Set default name if empty
  if (deviceName[0] == 0) {
    strcpy(deviceName, "LineBot-001");
  }
  
  Serial.println(F("Parameters loaded"));
}

void showSaveConfirmation() {
  ssd1306_clearScreen();
  ssd1306_setFixedFont(ssd1306xled_font6x8);
  ssd1306_printFixed(0, 20, "Parameters Saved!", STYLE_BOLD);
  ssd1306_printFixed(0, 35, "to EEPROM", STYLE_NORMAL);
  delay(SAVE_CONFIRM_DELAY_MS);
}