#include <Wire.h>
#include <Adafruit_GFX.h>
#include <Adafruit_SSD1306.h>
#include "RTClib.h"
#include <EEPROM.h>

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

// ---------- RTC ----------
RTC_DS3231 rtc;

// ---------- Encoder ----------
const uint8_t ENC_CLK = 3;
const uint8_t ENC_DT  = 2;
const uint8_t ENC_SW  = 4;

volatile long encoderPos = 0;
long lastEncoderPos = 0;

// ---------- UI ----------
bool inMenu = false;
int menuIndex = 0;
bool editing = false;
int editField = 0;
DateTime editDT;

const int menuItemCount = 4;
const char *menuItems[menuItemCount] = { "Set Time/Date", "Screen Timeout", "Contrast", "Info" };

// EEPROM addresses
#define EEPROM_TIMEOUT_ADDR 0
#define EEPROM_CONTRAST_ADDR 4

unsigned long screenTimeout = 4000;
int contrastValue = 127;

// Button
bool buttonState = HIGH;
bool lastButtonState = HIGH;
unsigned long buttonDownTime = 0;
const unsigned long longPressTime = 800;

// Idle tracking
unsigned long lastInteraction = 0;

// Blinking cursor
bool blinkState = true;
unsigned long lastBlinkTime = 0;
const unsigned long blinkInterval = 500; // ms

// Small clock icon 16x16 in PROGMEM
const unsigned char clock_bmp[] PROGMEM = {
  0x0E,0x00,0x11,0x80,0x20,0x40,0x40,0x40,0x40,0x40,0x20,0x40,0x11,0x80,0x0E,0x00
};

// ---------- Helper ----------
int lastDayOfMonth(int year, int month) {
  if(month==2){
    if((year%4==0 && year%100!=0) || (year%400==0)) return 29;
    else return 28;
  }
  else if(month==4 || month==6 || month==9 || month==11) return 30;
  else return 31;
}

// ---------- Encoder Polling (Wokwi-safe) ----------
void pollEncoder() {
  static int lastCLK = HIGH;
  int clkState = digitalRead(ENC_CLK);
  int dtState  = digitalRead(ENC_DT);

  if (clkState != lastCLK) {
    if (dtState != clkState) encoderPos++;
    else encoderPos--;
    lastInteraction = millis(); // reset idle timer
  }
  lastCLK = clkState;
}

// ---------- Button Handling ----------
void updateButton() {
  bool cur = digitalRead(ENC_SW);
  if (cur != lastButtonState) { delay(5); cur = digitalRead(ENC_SW); }

  if (lastButtonState == HIGH && cur == LOW) buttonDownTime = millis();
  if (lastButtonState == LOW && cur == HIGH) {
    unsigned long held = millis() - buttonDownTime;
    if (held >= longPressTime) {
      if (editing) editing = false;
      else if (inMenu) inMenu = false;
    } else {
      if (!inMenu) { inMenu = true; menuIndex = 0; }
      else {
        if (!editing) {
          if (menuIndex==0) { editing=true; editField=0; editDT=rtc.now(); }
          else if (menuIndex==1) editing=true;
          else if (menuIndex==2) editing=true;
        } else {
          if (menuIndex==0) { editField++; if(editField>5){rtc.adjust(editDT); editing=false;} }
          else if (menuIndex==1){ EEPROM.put(EEPROM_TIMEOUT_ADDR, screenTimeout); editing=false; } 
          else if (menuIndex==2){ setDisplayContrast((uint8_t)contrastValue); EEPROM.put(EEPROM_CONTRAST_ADDR, contrastValue); editing=false; }
        }
      }
    }
    lastInteraction = millis();
  }
  lastButtonState=cur;
  buttonState=cur;
}

// ---------- Encoder Movement ----------
void handleEncoderMovement() {
  long pos = encoderPos;
  if (pos != lastEncoderPos) {
    int8_t dir = (pos>lastEncoderPos)?1:-1;
    lastEncoderPos = pos;
    lastInteraction = millis(); // reset idle timer

    if (!inMenu) return;
    else if (!editing) {
      menuIndex += dir; if(menuIndex<0)menuIndex=menuItemCount-1; if(menuIndex>=menuItemCount)menuIndex=0;
    } else {
      if(menuIndex==0){
        switch(editField){
          case 0: // Year 2000-2099
            {int y = editDT.year()+dir; if(y<2000)y=2099; if(y>2099)y=2000; editDT=DateTime(y, editDT.month(), editDT.day(), editDT.hour(), editDT.minute(), editDT.second());} break;
          case 1: // Month 1-12
            {int m = editDT.month()+dir; if(m<1)m=12; if(m>12)m=1;
             int d = editDT.day(); int maxD = lastDayOfMonth(editDT.year(), m); if(d>maxD)d=maxD;
             editDT=DateTime(editDT.year(), m, d, editDT.hour(), editDT.minute(), editDT.second());} break;
          case 2: // Day 1-maxDay
            {int maxD = lastDayOfMonth(editDT.year(), editDT.month());
             int d = editDT.day()+dir; if(d<1)d=maxD; if(d>maxD)d=1;
             editDT=DateTime(editDT.year(), editDT.month(), d, editDT.hour(), editDT.minute(), editDT.second());} break;
          case 3: // Hour 0-23
            {int h = editDT.hour()+dir; if(h<0)h=23; if(h>23)h=0;
             editDT=DateTime(editDT.year(), editDT.month(), editDT.day(), h, editDT.minute(), editDT.second());} break;
          case 4: // Minute 0-59
            {int mi = editDT.minute()+dir; if(mi<0)mi=59; if(mi>59)mi=0;
             editDT=DateTime(editDT.year(), editDT.month(), editDT.day(), editDT.hour(), mi, editDT.second());} break;
          case 5: // Second 0-59
            {int s = editDT.second()+dir; if(s<0)s=59; if(s>59)s=0;
             editDT=DateTime(editDT.year(), editDT.month(), editDT.day(), editDT.hour(), editDT.minute(), s);} break;
        }
      } else if(menuIndex==1){
        if(dir>0)screenTimeout+=1000; else screenTimeout-=1000;
        if(screenTimeout<1000)screenTimeout=1000; if(screenTimeout>60000)screenTimeout=60000;
      } else if(menuIndex==2){
        if(dir>0)contrastValue+=8; else contrastValue-=8;
        if(contrastValue<0)contrastValue=0; if(contrastValue>255)contrastValue=255;
        setDisplayContrast((uint8_t)contrastValue);
      }
    }
  }
}

// ---------- Display Utilities ----------
void setDisplayContrast(uint8_t v){
#if defined(SSD1306_SETCONTRAST)
  display.ssd1306_command(SSD1306_SETCONTRAST);
  display.ssd1306_command(v);
#endif
}

void splashScreen() {
  display.clearDisplay();
  display.setTextSize(1); display.setTextColor(SSD1306_WHITE);
  display.setCursor(20,10); display.print("Welcome!");
  for(int i=0; i<4; i++){
    display.clearDisplay();
    display.setCursor(20,10); display.print("Welcome!");
    display.drawBitmap(54,30, clock_bmp, 16,16, SSD1306_WHITE);
    if(i%2==0) display.fillRect(54,30,16,16, SSD1306_WHITE);
    display.display(); delay(400);
  }
  display.clearDisplay(); display.display();
}

// ---------- Drawing ----------
void drawMainScreen(const DateTime &now){
  display.clearDisplay();
  display.drawRect(0,0,SCREEN_WIDTH,SCREEN_HEIGHT,SSD1306_WHITE);
  display.drawBitmap(6,6,clock_bmp,16,16,SSD1306_WHITE);
  display.setTextSize(3); display.setTextColor(SSD1306_WHITE); display.setCursor(28,4);
  char buf[9]; sprintf(buf,"%02u:%02u",now.hour(),now.minute()); display.print(buf);
  display.setTextSize(1); display.setCursor(28+72,20); sprintf(buf,"%02u",now.second()); display.print(buf);
  display.setCursor(10,44); sprintf(buf,"%02u/%02u/%04u",now.day(),now.month(),now.year()); display.print(buf);
  display.setCursor(80,54); display.print("DS3231 | OLED");
  display.display();
}

void drawMenu() {
  display.clearDisplay(); display.drawRect(0,0,SCREEN_WIDTH,SCREEN_HEIGHT,SSD1306_WHITE);
  display.setTextSize(1); display.setCursor(8,6); display.print("MENU");
  for(int i=0;i<menuItemCount;i++){
    int y=18+i*12;
    if(i==menuIndex){ display.fillRect(6,y-2,SCREEN_WIDTH-12,12,SSD1306_WHITE); display.setTextColor(SSD1306_BLACK); display.setCursor(10,y); display.print(menuItems[i]); display.setTextColor(SSD1306_WHITE);}
    else{ display.setCursor(10,y); display.print(menuItems[i]); }
  }
  if(editing){ display.setCursor(8,SCREEN_HEIGHT-10); display.print("Editing - press to advance"); }
  display.display();
}

void drawEditing() {
  display.clearDisplay(); display.drawRect(0,0,SCREEN_WIDTH,SCREEN_HEIGHT,SSD1306_WHITE);
  display.setTextSize(1);
  if(menuIndex==0){
    display.setCursor(8,6); display.print("Set Time / Date");
    display.setCursor(8,20); if(editField==0 && blinkState)display.print("> "); else display.print("  "); display.print("Year: "); display.print(editDT.year());
    display.setCursor(8,32); if(editField==1 && blinkState)display.print("> "); else display.print("  "); display.print("Month: "); display.print(editDT.month());
    display.setCursor(8,44); if(editField==2 && blinkState)display.print("> "); else display.print("  "); display.print("Day: "); display.print(editDT.day());
    display.setCursor(72,20); if(editField==3 && blinkState)display.print("> "); else display.print("  "); display.print("Hour: "); display.print(editDT.hour());
    display.setCursor(72,32); if(editField==4 && blinkState)display.print("> "); else display.print("  "); display.print("Min: "); display.print(editDT.minute());
    display.setCursor(72,44); if(editField==5 && blinkState)display.print("> "); else display.print("  "); display.print("Sec: "); display.print(editDT.second());
  } else if(menuIndex==1){
    display.setCursor(8,6); display.print("Screen Timeout"); display.setCursor(8,26); display.print(screenTimeout/1000); display.print(" s");
  } else if(menuIndex==2){
    display.setCursor(8,6); display.print("Contrast"); display.setCursor(8,26); display.print(contrastValue);
  }
  display.display();
}

void drawInfo() {
  display.clearDisplay(); display.drawRect(0,0,SCREEN_WIDTH,SCREEN_HEIGHT,SSD1306_WHITE);
  display.setTextSize(1); display.setCursor(8,8); display.print("Info");
  display.setCursor(8,20); display.print("SSD1306 + DS3231 Nano Demo");
  display.setCursor(8,32); display.print("Rotate=nav, press=select");
  display.setCursor(8,44); display.print("Long press = back/close");
  display.display();
}

// ---------- Setup ----------
void setup() {
  pinMode(ENC_CLK, INPUT_PULLUP);
  pinMode(ENC_DT, INPUT_PULLUP);
  pinMode(ENC_SW, INPUT_PULLUP);

  Wire.begin();
  if(!display.begin(SSD1306_SWITCHCAPVCC, 0x3C)){ while(1); }
  display.clearDisplay(); display.display();

  splashScreen();
  if(!rtc.begin()){ while(1); }

  EEPROM.get(EEPROM_TIMEOUT_ADDR, screenTimeout);
  if(screenTimeout<1000 || screenTimeout>60000) screenTimeout=4000;
  EEPROM.get(EEPROM_CONTRAST_ADDR, contrastValue);
  if(contrastValue>255) contrastValue=127;
  setDisplayContrast((uint8_t)contrastValue);

  lastInteraction = millis();
}

// ---------- Loop ----------
void loop() {
  unsigned long nowMillis = millis();
  if (nowMillis - lastBlinkTime >= blinkInterval) { blinkState = !blinkState; lastBlinkTime = nowMillis; }

  updateButton();
  pollEncoder();          
  handleEncoderMovement();

  if(inMenu && (millis() - lastInteraction > screenTimeout)) { inMenu=false; editing=false; }

  if(inMenu){
    if(editing) drawEditing();
    else if(menuIndex==3) drawInfo();
    else drawMenu();
  } else drawMainScreen(rtc.now());

  delay(20);
}