#include <Wire.h>
#include <Adafruit_GFX.h>
#include <Adafruit_SSD1306.h>
#include <WiFi.h>
#include <esp_now.h>
#include <Preferences.h>

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

// --- BITMAP SECTION ---
const unsigned char ramoon_logo [] PROGMEM = {
	0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 
	0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 
	0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 
	0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 
	0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 
	0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0xff, 0xc0, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 
	0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x47, 0xff, 0xf3, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 
	0x00, 0x00, 0x00, 0x00, 0x00, 0x01, 0xff, 0xff, 0xff, 0xc0, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 
	0x00, 0x00, 0x00, 0x00, 0x00, 0x03, 0x3f, 0xff, 0xfe, 0xe0, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 
	0x00, 0x00, 0x00, 0x00, 0x00, 0x05, 0xdf, 0xff, 0xfb, 0xb0, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 
	0x00, 0x00, 0x00, 0x00, 0x00, 0x0b, 0x2f, 0xff, 0xf0, 0xf8, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 
	0x00, 0x00, 0x00, 0x00, 0x00, 0x0e, 0x77, 0xff, 0xf6, 0x68, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 
	0x00, 0x00, 0x00, 0x00, 0x00, 0x0e, 0x27, 0xff, 0xe4, 0x28, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 
	0x00, 0x00, 0x00, 0x00, 0x00, 0x0e, 0x07, 0xff, 0xe0, 0x38, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 
	0x00, 0x00, 0x00, 0x00, 0x00, 0x0e, 0x0f, 0xff, 0xf0, 0x28, 0x00, 0x00, 0x78, 0x00, 0x00, 0x00, 
	0x00, 0x00, 0x00, 0x00, 0x00, 0x1a, 0x0f, 0xff, 0xf0, 0x7c, 0x3f, 0xff, 0xe0, 0x00, 0x00, 0x00, 
	0x00, 0x00, 0x00, 0x00, 0x00, 0x1d, 0x1f, 0xff, 0xf8, 0xbe, 0x3f, 0xff, 0xc0, 0x00, 0x00, 0x00, 
	0x00, 0x00, 0x03, 0xf8, 0x1f, 0x3f, 0xff, 0xff, 0xff, 0xfe, 0x3f, 0xff, 0x00, 0x00, 0x00, 0x00, 
	0x00, 0x00, 0x00, 0xff, 0xfe, 0x3f, 0xff, 0xff, 0xff, 0xfe, 0x3f, 0xff, 0x00, 0x00, 0x00, 0x00, 
	0x00, 0x00, 0x00, 0x3f, 0xff, 0x3e, 0xfe, 0xfd, 0xff, 0x7e, 0x3f, 0xfe, 0x00, 0x00, 0x00, 0x00, 
	0x00, 0x00, 0x00, 0x1f, 0xff, 0x1f, 0x7f, 0xff, 0xfe, 0xfe, 0x3f, 0xfe, 0x00, 0x00, 0x00, 0x00, 
	0x00, 0x00, 0x00, 0x0f, 0xff, 0x1f, 0xcf, 0xff, 0xf3, 0xfc, 0x7f, 0xfe, 0x00, 0x00, 0x00, 0x00, 
	0x00, 0x00, 0x00, 0x0f, 0xff, 0x8f, 0xf0, 0xff, 0x1f, 0xf8, 0xff, 0xfc, 0x00, 0x00, 0x00, 0x00, 
	0x00, 0x00, 0x00, 0x0f, 0xff, 0xc3, 0xe7, 0x01, 0xff, 0xf1, 0xff, 0xe0, 0x00, 0x00, 0x00, 0x00, 
	0x00, 0x00, 0x00, 0x0f, 0xff, 0xe1, 0xef, 0xff, 0xff, 0xc3, 0xff, 0xc0, 0x00, 0x00, 0x00, 0x00, 
	0x00, 0x00, 0x00, 0x00, 0x7f, 0xf8, 0x7f, 0xff, 0xff, 0x8f, 0xff, 0x80, 0x00, 0x00, 0x00, 0x00, 
	0x00, 0x00, 0x00, 0x00, 0x3f, 0xfe, 0x0f, 0xff, 0xfc, 0x3f, 0xff, 0x80, 0x00, 0x00, 0x00, 0x00, 
	0x00, 0x00, 0x00, 0x00, 0x3f, 0xff, 0xe3, 0xff, 0xf1, 0xff, 0xe0, 0x80, 0x00, 0x00, 0x00, 0x00, 
	0x00, 0x00, 0x00, 0x00, 0x1f, 0xff, 0xcd, 0xff, 0xf7, 0xff, 0xc1, 0x00, 0x00, 0x00, 0x00, 0x00, 
	0x00, 0x00, 0x00, 0x00, 0x10, 0x3f, 0xed, 0xff, 0xfb, 0x83, 0x83, 0x80, 0x00, 0x00, 0x00, 0x00, 
	0x00, 0x00, 0x00, 0x00, 0x00, 0x1c, 0x0d, 0xff, 0xfc, 0x60, 0x07, 0x80, 0x00, 0x00, 0x00, 0x00, 
	0x00, 0x00, 0x00, 0x00, 0x00, 0x10, 0x01, 0xff, 0xff, 0x98, 0x0f, 0x80, 0x00, 0x00, 0x00, 0x00, 
	0x00, 0x00, 0x00, 0x00, 0x00, 0x80, 0x89, 0xff, 0xff, 0xf0, 0x7f, 0x00, 0x00, 0x00, 0x00, 0x00, 
	0x00, 0x00, 0x00, 0x00, 0x07, 0xe7, 0xe5, 0xff, 0xfe, 0xff, 0xff, 0x00, 0x00, 0x00, 0x00, 0x00, 
	0x00, 0x00, 0x00, 0x00, 0x0e, 0x7f, 0xf7, 0xff, 0xff, 0xbc, 0x3e, 0x00, 0x00, 0x00, 0x00, 0x00, 
	0x00, 0x00, 0x00, 0x00, 0x1d, 0xff, 0xe6, 0xff, 0xff, 0xdf, 0xdc, 0x00, 0x00, 0x00, 0x00, 0x00, 
	0x00, 0x00, 0x00, 0x00, 0x1b, 0xde, 0xcc, 0xff, 0xf7, 0xdf, 0xf8, 0x00, 0x00, 0x00, 0x00, 0x00, 
	0x00, 0x00, 0x00, 0x00, 0x1b, 0xff, 0x1c, 0x7f, 0xfb, 0xff, 0xe0, 0x00, 0x00, 0x00, 0x00, 0x00, 
	0x00, 0x00, 0x00, 0x00, 0x3f, 0xfe, 0x38, 0x1f, 0xfd, 0xff, 0xe0, 0x00, 0x00, 0x00, 0x00, 0x00, 
	0x00, 0x00, 0x00, 0x00, 0x3f, 0xff, 0xf8, 0x07, 0xfd, 0xf8, 0xe0, 0x00, 0x00, 0x00, 0x00, 0x00, 
	0x00, 0x00, 0x00, 0x00, 0x3f, 0xff, 0xf8, 0x0c, 0x7d, 0xe0, 0xe0, 0x00, 0x00, 0x00, 0x00, 0x00, 
	0x00, 0x00, 0x00, 0x00, 0x3b, 0xff, 0x5c, 0x3f, 0x03, 0xc1, 0xf8, 0x00, 0x00, 0x00, 0x00, 0x00, 
	0x00, 0x00, 0x00, 0x00, 0x1b, 0xdf, 0xbf, 0x00, 0x03, 0xc6, 0xfc, 0x00, 0x00, 0x00, 0x00, 0x00, 
	0x00, 0x00, 0x00, 0x00, 0x1d, 0xff, 0xf8, 0x00, 0x3f, 0xe0, 0x40, 0x00, 0x00, 0x00, 0x00, 0x00, 
	0x00, 0x00, 0x00, 0x00, 0x0e, 0x2f, 0xf0, 0x00, 0x3b, 0xf8, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 
	0x00, 0x00, 0x00, 0x00, 0x07, 0xe7, 0xe0, 0x07, 0x03, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 
	0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x80, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 
	0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 
	0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 
	0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 
	0x00, 0x00, 0x00, 0x00, 0x3f, 0x03, 0x83, 0x03, 0x1f, 0x81, 0xf0, 0xe3, 0x00, 0x00, 0x00, 0x00, 
	0x00, 0x00, 0x00, 0x00, 0x3f, 0xc7, 0xc3, 0x87, 0x3f, 0xe7, 0xbc, 0xf3, 0x00, 0x00, 0x00, 0x00, 
	0x00, 0x00, 0x00, 0x00, 0x39, 0xc7, 0xc3, 0x8f, 0x70, 0xe6, 0x1c, 0xfb, 0x00, 0x00, 0x00, 0x00, 
	0x00, 0x00, 0x00, 0x00, 0x39, 0xc6, 0xe3, 0xdf, 0x70, 0x6e, 0x0c, 0xff, 0x00, 0x00, 0x00, 0x00, 
	0x00, 0x00, 0x00, 0x00, 0x3f, 0xce, 0xe3, 0xfb, 0x70, 0x6e, 0x0c, 0xdf, 0x00, 0x00, 0x00, 0x00, 
	0x00, 0x00, 0x00, 0x00, 0x3f, 0x8f, 0xf3, 0x73, 0x38, 0xe7, 0x1c, 0xcf, 0x00, 0x00, 0x00, 0x00, 
	0x00, 0x00, 0x00, 0x00, 0x3b, 0x9c, 0x73, 0x33, 0x3f, 0xc3, 0xf8, 0xc7, 0x00, 0x00, 0x00, 0x00, 
	0x00, 0x00, 0x00, 0x00, 0x39, 0xd8, 0x3b, 0x03, 0x0f, 0x81, 0xe0, 0xc3, 0x00, 0x00, 0x00, 0x00, 
	0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 
	0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 
	0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 
	0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 
	0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 
	0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00
};

// --- PIN DEFINITIONS ---
const int PIN_UP = 2, PIN_DOWN = 3, PIN_ENTER = 4, PIN_BACK = 5;
const int PIN_LEFT = 10, PIN_RIGHT = 6, ENC_SW = 14;
const int ENC_CLK = 12, ENC_DT = 13;
const int JS1_X = 34, JS1_Y = 35, JS2_Z = 32; 

// --- ESP-NOW SETUP ---
uint8_t robotAddress[] = {0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF}; // <--- CHANGE TO YOUR ROBOT MAC
typedef struct struct_message {
  int x; int y; int z; int speed; int posID; bool isSaving;
} struct_message;
struct_message controlData;

// --- MENU DATA ---
const char* mainMenuOptions[] = {"HOMING", "MODE", "SETTINGS", "DEV INFO"};
const char* modeSelectOptions[] = {"MANUAL", "FIXED POS", "TEMP POS"};

// --- STATE & VARIABLES ---
enum MenuState { LOGO, MAIN_MENU, MODE_SELECT, FIXED_POS, TEMP_POS, TRANSITION, MANUAL_CAM, DEV_INFO };
MenuState currentState = LOGO;

int currentSpeed = 5;
int lastClkState;
int menuIndex = 0;
int selectedPos = 1;
int lastConfirmedPos = 1;
unsigned long animStartTime = 0;
unsigned long pressTimer = 0;
bool isFixedMenu = true;

// Storage Arrays
int fixedPosCoords[11]; 
int tempPosCoords[11];  

void setup() {
  int inputPins[] = {PIN_UP, PIN_DOWN, PIN_ENTER, PIN_BACK, PIN_LEFT, PIN_RIGHT, ENC_SW};
  for(int p : inputPins) pinMode(p, INPUT_PULLUP);
  
  pinMode(ENC_CLK, INPUT);
  pinMode(ENC_DT, INPUT);
  lastClkState = digitalRead(ENC_CLK);

  WiFi.mode(WIFI_STA);
  if (esp_now_init() == ESP_OK) {
    esp_now_peer_info_t peerInfo = {};
    memcpy(peerInfo.peer_addr, robotAddress, 6);
    peerInfo.channel = 0; peerInfo.encrypt = false;
    esp_now_add_peer(&peerInfo);
  }

  preferences.begin("rob-pos", false);
  for(int i=1; i<=10; i++) {
    fixedPosCoords[i] = preferences.getInt(("p" + String(i)).c_str(), 0);
  }

  Wire.begin(8, 9); 
  display.begin(SSD1306_SWITCHCAPVCC, 0x3C);
  animStartTime = millis();
}

// --- LOGIC HELPERS ---

void updateJoystickAndSend() {
  controlData.x = analogRead(JS1_X);
  controlData.y = analogRead(JS1_Y);
  controlData.z = analogRead(JS2_Z);
  controlData.speed = currentSpeed;
  esp_now_send(robotAddress, (uint8_t *) &controlData, sizeof(controlData));
}

bool isEnterPressed() {
  return (digitalRead(PIN_ENTER) == LOW || digitalRead(ENC_SW) == LOW);
}

void savePosition() {
  // Store current X as a coordinate example
  if (isFixedMenu) {
    fixedPosCoords[selectedPos] = controlData.x; 
    preferences.putInt(("p" + String(selectedPos)).c_str(), controlData.x);
  } else {
    tempPosCoords[selectedPos] = controlData.x;
  }
  display.clearDisplay();
  display.setCursor(20, 30);
  display.print("POSITION SAVED!");
  display.display();
  delay(1000);
}

void handlePositionSelection() {
  bool btnDown = isEnterPressed();
  if (btnDown) {
    if (pressTimer == 0) pressTimer = millis();
    if (millis() - pressTimer > 2000) {
       savePosition();
       pressTimer = 0; 
       delay(500); 
    }
  } else if (pressTimer > 0) {
    if (millis() - pressTimer < 2000) {
      if (selectedPos != lastConfirmedPos) {
        currentState = TRANSITION;
        animStartTime = millis();
        controlData.posID = selectedPos;
        controlData.isSaving = false;
      }
    }
    pressTimer = 0;
  }
}

// --- MAIN LOOP ---

void loop() {
  updateJoystickAndSend();
  
  // Encoder Speed Logic
  int currentStateClk = digitalRead(ENC_CLK);
  if (currentStateClk != lastClkState && currentStateClk == 1) {
    if (digitalRead(ENC_DT) != currentStateClk) { if (currentSpeed < 10) currentSpeed++; }
    else { if (currentSpeed > 1) currentSpeed--; }
  }
  lastClkState = currentStateClk;

  display.clearDisplay();

  if (currentState == LOGO) {
    display.drawBitmap(0, 0, ramoon_logo, 128, 64, WHITE);
    if (millis() - animStartTime > 3000) currentState = MAIN_MENU;
  } else {
    // Top Info Bar
    display.setTextSize(1); display.setTextColor(WHITE);
    display.setCursor(0, 0);
    display.print("R:85% B:90% SPD:"); display.print(currentSpeed);
    display.drawFastHLine(0, 12, 128, WHITE);

    switch (currentState) {
      case MAIN_MENU:
        renderListMenu(mainMenuOptions, 4);
        if (isEnterPressed()) {
          if (menuIndex == 0) { currentState = TRANSITION; animStartTime = millis(); }
          else if (menuIndex == 1) { currentState = MODE_SELECT; menuIndex = 0; }
          else if (menuIndex == 3) currentState = DEV_INFO;
          delay(200);
        }
        break;

      case MODE_SELECT:
        renderListMenu(modeSelectOptions, 3);
        if (isEnterPressed()) {
          if (menuIndex == 0) currentState = MANUAL_CAM;
          if (menuIndex == 1) { currentState = FIXED_POS; isFixedMenu = true; }
          if (menuIndex == 2) { currentState = TEMP_POS; isFixedMenu = false; }
          delay(200);
        }
        if (digitalRead(PIN_BACK) == LOW) currentState = MAIN_MENU;
        break;

      case FIXED_POS:
      case TEMP_POS:
        drawGridPositions();
        handlePositionSelection();
        if (digitalRead(PIN_BACK) == LOW) currentState = MODE_SELECT;
        break;

      case MANUAL_CAM:
        display.setCursor(30, 35); display.print("MOVE THE CAM");
        if (digitalRead(PIN_BACK) == LOW) currentState = MODE_SELECT;
        break;

      case TRANSITION:
        runRepeatingTransition();
        break;

      case DEV_INFO:
        display.drawBitmap(0, 0, ramoon_logo, 128, 64, WHITE);
        if (digitalRead(PIN_BACK) == LOW) currentState = MAIN_MENU;
        break;
    }
  }
  display.display();
}

// --- UI DRAWING FUNCTIONS ---

void renderListMenu(const char* options[], int count) {
  for(int i = 0; i < count; i++) {
    if(i == menuIndex) {
      display.fillRect(0, 15 + (i*12), 128, 11, WHITE);
      display.setTextColor(BLACK);
    } else { display.setTextColor(WHITE); }
    display.setCursor(5, 17 + (i*12));
    display.print(options[i]);
  }
  if(digitalRead(PIN_DOWN) == LOW) { menuIndex = (menuIndex + 1) % count; delay(150); }
  if(digitalRead(PIN_UP) == LOW) { menuIndex = (menuIndex + count - 1) % count; delay(150); }
}

void drawGridPositions() {
  for(int i = 1; i <= 10; i++) {
    int x = 15 + ((i-1) % 5) * 24;
    int y = 25 + ((i-1) / 5) * 20;
    if(i == lastConfirmedPos) { display.fillCircle(x, y, 9, WHITE); display.setTextColor(BLACK); }
    else if(i == selectedPos) { display.drawCircle(x, y, 9, WHITE); display.drawCircle(x, y, 8, WHITE); display.setTextColor(WHITE); }
    else { display.drawCircle(x, y, 9, WHITE); display.setTextColor(WHITE); }
    display.setCursor(x-3, y-3); display.print(i);
  }
  if(digitalRead(PIN_RIGHT) == LOW) { if(selectedPos < 10) selectedPos++; delay(150); }
  if(digitalRead(PIN_LEFT) == LOW)  { if(selectedPos > 1) selectedPos--; delay(150); }
  if(digitalRead(PIN_DOWN) == LOW)  { if(selectedPos <= 5) selectedPos += 5; delay(150); }
  if(digitalRead(PIN_UP) == LOW)    { if(selectedPos > 5) selectedPos -= 5; delay(150); }
}

void runRepeatingTransition() {
  unsigned long elapsed = millis() - animStartTime;
  display.drawCircle(20, 38, 15, WHITE);
  display.drawCircle(108, 38, 15, WHITE);
  unsigned long repeatTimer = elapsed % 1000;
  int boxX = 20 + (int)((float)repeatTimer / 1000.0 * 88.0); 
  display.fillRect(boxX - 5, 33, 10, 10, WHITE); 
  if (elapsed >= 10000) { 
    lastConfirmedPos = selectedPos; 
    currentState = (isFixedMenu ? FIXED_POS : TEMP_POS); 
    if (currentState == TRANSITION) currentState = MAIN_MENU; // From Homing
  }
}