#include <Wire.h>
#include <Adafruit_GFX.h>
#include <Adafruit_SSD1306.h>
#include <SPI.h>
#include "SPIFFS.h"
#include <WiFi.h>
#include <WebServer.h>
#include <time.h>
#include <esp_sleep.h>
#include <driver/gpio.h>
#include <esp_timer.h>
#include <ArduinoJson.h>

// ================= DISPLAY CONFIGURATION =================
#define SCREEN_WIDTH 128
#define SCREEN_HEIGHT 64
#define SDA_PIN 8
#define SCL_PIN 9
#define OLED_RESET -1
#define SCREEN_ADDRESS 0x3C
#define OLED_POWER_PIN 10

// ================= BUTTON CONFIGURATION =================
#define BTN_UP     3
#define BTN_DOWN   4
#define BTN_SELECT 5
#define BTN_BACK   6

// ================= WIFI CONFIGURATION =================
const char* ap_ssid = "FPV_CAM";
const char* ap_password = "";
const char* wifi_ssid = "FPV_CAM";
const char* wifi_password = "12345678";

// Timezone: IST (UTC+5:30)
const long gmtOffset_sec = 19800;
const int daylightOffset_sec = 0;
const char* ntpServer = "pool.ntp.org";

// ================= SYSTEM CONSTANTS =================
#define MAX_CHAPTERS 10
#define MAX_TOPICS 20
#define MAX_LINES 500
#define DEBOUNCE_DELAY 200
#define LINES_PER_SCREEN 8
#define VISIBLE_ITEMS 5
#define WATCH_SLEEP_TIMEOUT_MS 30000

// ================= RTC MEMORY =================
RTC_DATA_ATTR int rtc_lastChapter = 0;
RTC_DATA_ATTR int rtc_lastTopic = 0;
RTC_DATA_ATTR int rtc_lastScrollPos = 0;
RTC_DATA_ATTR bool rtc_hasResumeData = false;
RTC_DATA_ATTR unsigned long rtc_bootCount = 0;

// Time stored in RTC
RTC_DATA_ATTR unsigned long rtc_unixTime = 0;
RTC_DATA_ATTR uint64_t rtc_sleepEnterTime = 0;
RTC_DATA_ATTR bool rtc_timeSet = false;

// ================= SYSTEM VARIABLES =================
enum Mode { MODE_WATCH, MODE_READER, MODE_AP, MODE_TIME_SYNC };
Mode currentMode = MODE_WATCH;

enum ScreenID { SCR_MAIN, SCR_CHAPTERS, SCR_TOPICS, SCR_READ };
ScreenID readerScreen = SCR_MAIN;

int readerCursor = 0;
int readerScrollTop = 0;
int selectedChapter = 0;
int selectedTopic = 0;
int scrollPosition = 0;

String chapters[MAX_CHAPTERS];
int chapterCount = 0;
String topics[MAX_TOPICS];
int topicCount = 0;
String currentChapter = "";
String textLines[MAX_LINES];
int lineCount = 0;

int brightness = 200;

// Current time variables
int hours = 12, minutes = 0, seconds = 0;
int day = 1, month = 1, year = 2024;
int weekDay = 3;
unsigned long lastTick = 0;

bool apEnabled = false;
WebServer server(80);
unsigned long lastButtonTime = 0;
unsigned long lastActivityTime = 0;

Adafruit_SSD1306 display(SCREEN_WIDTH, SCREEN_HEIGHT, &Wire, OLED_RESET);

// ================= FUNCTION PROTOTYPES =================
void initSPIFFS();
void loadChapters();
void saveChapters();
void loadTopics(int ch);
void saveTopics(int ch);
void loadContent(int ch, int topic);
bool saveContent(int ch, int topic, String content);
void addChapter(String name);
void deleteChapter(int ch);
void addTopic(int ch, String name);
void deleteTopic(int ch, int topic);
void renameChapter(int ch, String name);
void renameTopic(int ch, int topic, String name);
void setupAPMode();
void setupWebServer();
void syncTimeWithWiFi();
void updateTimeFromUnix();
void saveTimeToRTC();
void loadTimeFromRTC();
void goToSleep();
void showMessage(String msg);
String getWeekday();
int getDaysInMonth(int month, int year);
void displayWatchFace();
void displayMenu();
void displayChapters();
void displayTopics();
void displayContent();
void displayAPScreen();
void handleButtons();

// ================= FIXED TIME CONVERSION FUNCTIONS =================
void convertUnixToTime(unsigned long unixTime) {
  // Add IST offset (5 hours 30 minutes = 19800 seconds)
  unsigned long localTime = unixTime + 19800;
  
  // Extract time components
  seconds = localTime % 60;
  unsigned long totalMinutes = localTime / 60;
  minutes = totalMinutes % 60;
  unsigned long totalHours = totalMinutes / 60;
  hours = totalHours % 24;
  
  // Calculate days since epoch
  unsigned long days = totalHours / 24;
  
  // Unix epoch started on Jan 1, 1970
  // Start calculation from 1970
  year = 1970;
  
  while (days >= 365) {
    int daysInYear = ((year % 4 == 0 && year % 100 != 0) || (year % 400 == 0)) ? 366 : 365;
    if (days >= daysInYear) {
      days -= daysInYear;
      year++;
    } else {
      break;
    }
  }
  
  month = 1;
  while (true) {
    int daysInMonth = getDaysInMonth(month, year);
    if (days >= daysInMonth) {
      days -= daysInMonth;
      month++;
    } else {
      break;
    }
  }
  
  day = days + 1;
  
  // Calculate day of week (Jan 1, 1970 was Thursday = 4)
  // 0 = Sunday, 1 = Monday, 2 = Tuesday, 3 = Wednesday, 4 = Thursday, 5 = Friday, 6 = Saturday
  unsigned long totalDaysSince1970 = (totalHours / 24);
  weekDay = (4 + totalDaysSince1970) % 7;
}

void updateTimeFromUnix() {
  convertUnixToTime(rtc_unixTime);
}

// ================= SETUP =================
void setup() {
  Serial.begin(115200);
  rtc_bootCount++;
  Serial.println("ESP32-C3 Smart Watch - Boot #" + String(rtc_bootCount));
  
  pinMode(BTN_UP, INPUT_PULLUP);
  pinMode(BTN_DOWN, INPUT_PULLUP);
  pinMode(BTN_SELECT, INPUT_PULLUP);
  pinMode(BTN_BACK, INPUT_PULLUP);
  pinMode(OLED_POWER_PIN, OUTPUT);
  digitalWrite(OLED_POWER_PIN, HIGH);
  
  Wire.begin(SDA_PIN, SCL_PIN);
  display.begin(SSD1306_SWITCHCAPVCC, SCREEN_ADDRESS);
  display.setRotation(2);
  display.setTextColor(WHITE);
  display.ssd1306_command(SSD1306_SETCONTRAST);
  display.ssd1306_command(brightness);
  
  loadTimeFromRTC();
  
  if (!SPIFFS.begin(true)) {
    Serial.println("SPIFFS mount failed");
  } else {
    if (!SPIFFS.exists("/chapters.json")) {
      addChapter("Getting Started");
      addChapter("User Guide");
      addChapter("Advanced");
    }
    loadChapters();
  }
  
  display.clearDisplay();
  display.setCursor(25, 20);
  display.println("SMART");
  display.setCursor(20, 40);
  display.println("WATCH");
  display.display();
  delay(1000);
  
  lastActivityTime = millis();
  lastTick = millis();
  updateTimeFromUnix();
  displayWatchFace();
}

void loadTimeFromRTC() {
  if (rtc_timeSet && rtc_unixTime > 0) {
    convertUnixToTime(rtc_unixTime);
    Serial.printf("Time loaded from RTC: %02d:%02d:%02d %02d/%02d/%04d (Unix: %lu)\n", 
                  hours, minutes, seconds, day, month, year, rtc_unixTime);
  } else {
    Serial.println("No time in RTC, will sync with WiFi");
    rtc_unixTime = 0;
  }
}

void saveTimeToRTC() {
  rtc_unixTime = getCurrentUnixTime();
  rtc_timeSet = true;
  Serial.printf("Time saved to RTC: %lu\n", rtc_unixTime);
}

unsigned long getCurrentUnixTime() {
  return rtc_unixTime + ((millis() - lastTick) / 1000);
}

// ================= TIME SYNC FUNCTION =================
void syncTimeWithWiFi() {
  Mode previousMode = currentMode;
  currentMode = MODE_TIME_SYNC;
  
  display.clearDisplay();
  display.setCursor(10, 15);
  display.println("Sync Time");
  display.drawLine(0, 25, 128, 25, WHITE);
  display.setCursor(10, 38);
  display.println("Connecting to WiFi...");
  display.display();
  
  WiFi.mode(WIFI_STA);
  WiFi.begin(wifi_ssid, wifi_password);
  
  int attempts = 0;
  
  while (WiFi.status() != WL_CONNECTED && attempts < 30) {
    delay(500);
    attempts++;
    display.clearDisplay();
    display.setCursor(10, 15);
    display.println("Sync Time");
    display.drawLine(0, 25, 128, 25, WHITE);
    display.setCursor(10, 38);
    display.print("Connecting");
    for (int i = 0; i < (attempts / 3); i++) display.print(".");
    display.display();
  }
  
  if (WiFi.status() == WL_CONNECTED) {
    display.clearDisplay();
    display.setCursor(10, 15);
    display.println("Sync Time");
    display.drawLine(0, 25, 128, 25, WHITE);
    display.setCursor(10, 38);
    display.println("Getting time from NTP...");
    display.display();
    
    configTime(gmtOffset_sec, daylightOffset_sec, ntpServer);
    
    struct tm timeinfo;
    int syncAttempts = 0;
    
    while (!getLocalTime(&timeinfo) && syncAttempts < 20) {
      delay(500);
      syncAttempts++;
      display.clearDisplay();
      display.setCursor(10, 15);
      display.println("Sync Time");
      display.drawLine(0, 25, 128, 25, WHITE);
      display.setCursor(10, 38);
      display.print("Getting time");
      for (int i = 0; i < (syncAttempts / 3); i++) display.print(".");
      display.display();
    }
    
    if (syncAttempts < 20) {
      // Convert tm structure to Unix timestamp
      time_t utcTime = mktime(&timeinfo);
      rtc_unixTime = utcTime;
      rtc_timeSet = true;
      
      updateTimeFromUnix();
      
      display.clearDisplay();
      display.setCursor(10, 15);
      display.println("Sync Time");
      display.drawLine(0, 25, 128, 25, WHITE);
      display.setCursor(20, 38);
      display.println("SUCCESS!");
      display.setCursor(10, 52);
      char timeStr[20];
      sprintf(timeStr, "%02d:%02d:%02d", hours, minutes, seconds);
      display.println(timeStr);
      display.display();
      delay(2000);
      
      Serial.printf("Time synced: %02d:%02d:%02d %02d/%02d/%04d\n", 
                    hours, minutes, seconds, day, month, year);
    } else {
      display.clearDisplay();
      display.setCursor(10, 15);
      display.println("Sync Time");
      display.drawLine(0, 25, 128, 25, WHITE);
      display.setCursor(10, 38);
      display.println("NTP Timeout!");
      display.display();
      delay(2000);
    }
    
    WiFi.disconnect(true);
    WiFi.mode(WIFI_OFF);
  } else {
    display.clearDisplay();
    display.setCursor(10, 15);
    display.println("Sync Time");
    display.drawLine(0, 25, 128, 25, WHITE);
    display.setCursor(10, 38);
    display.println("WiFi Failed!");
    display.display();
    delay(2000);
  }
  
  currentMode = previousMode;
  lastTick = millis();
  
  if (currentMode == MODE_WATCH) {
    displayWatchFace();
  } else if (currentMode == MODE_READER) {
    displayMenu();
  }
}

// ================= TIME COUNTER =================
void updateTimeCounter() {
  unsigned long now = millis();
  if (now - lastTick >= 1000) {
    lastTick = now;
    rtc_unixTime++;
    updateTimeFromUnix();
    
    if (currentMode == MODE_WATCH) {
      displayWatchFace();
    }
  }
}

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

// ================= FILE MANAGEMENT =================
void loadChapters() {
  File file = SPIFFS.open("/chapters.json", FILE_READ);
  if (!file) return;
  
  DynamicJsonDocument doc(2048);
  deserializeJson(doc, file);
  JsonArray arr = doc.as<JsonArray>();
  chapterCount = 0;
  for (JsonVariant v : arr) {
    if (chapterCount < MAX_CHAPTERS) chapters[chapterCount++] = v.as<String>();
  }
  file.close();
}

void saveChapters() {
  DynamicJsonDocument doc(2048);
  JsonArray arr = doc.to<JsonArray>();
  for (int i = 0; i < chapterCount; i++) arr.add(chapters[i]);
  
  File file = SPIFFS.open("/chapters.json", FILE_WRITE);
  serializeJson(doc, file);
  file.close();
}

void loadTopics(int ch) {
  String path = "/ch" + String(ch + 1) + ".json";
  File file = SPIFFS.open(path, FILE_READ);
  topicCount = 0;
  if (!file) return;
  
  DynamicJsonDocument doc(2048);
  deserializeJson(doc, file);
  JsonArray arr = doc.as<JsonArray>();
  for (JsonVariant v : arr) {
    if (topicCount < MAX_TOPICS) topics[topicCount++] = v.as<String>();
  }
  file.close();
}

void saveTopics(int ch) {
  String path = "/ch" + String(ch + 1) + ".json";
  DynamicJsonDocument doc(2048);
  JsonArray arr = doc.to<JsonArray>();
  for (int i = 0; i < topicCount; i++) arr.add(topics[i]);
  
  File file = SPIFFS.open(path, FILE_WRITE);
  serializeJson(doc, file);
  file.close();
}

void loadContent(int ch, int topic) {
  String path = "/ch" + String(ch + 1) + "_" + String(topic + 1) + ".txt";
  File file = SPIFFS.open(path, FILE_READ);
  lineCount = 0;
  
  if (!file) {
    textLines[lineCount++] = "No content";
    return;
  }
  
  while (file.available() && lineCount < MAX_LINES) {
    String line = file.readStringUntil('\n');
    line.trim();
    while (line.length() > 21 && lineCount < MAX_LINES) {
      int wrap = 21;
      for (int i = 21; i > 15; i--) {
        if (i < line.length() && line.charAt(i) == ' ') {
          wrap = i;
          break;
        }
      }
      textLines[lineCount++] = line.substring(0, wrap);
      line = line.substring(wrap);
      line.trim();
    }
    if (line.length() > 0 && lineCount < MAX_LINES) {
      textLines[lineCount++] = line;
    }
  }
  file.close();
}

bool saveContent(int ch, int topic, String content) {
  String path = "/ch" + String(ch + 1) + "_" + String(topic + 1) + ".txt";
  File file = SPIFFS.open(path, FILE_WRITE);
  if (!file) return false;
  file.print(content);
  file.close();
  return true;
}

void addChapter(String name) {
  if (chapterCount >= MAX_CHAPTERS) return;
  chapters[chapterCount++] = name;
  saveChapters();
  
  topicCount = 0;
  for (int i = 1; i <= 3; i++) {
    topics[topicCount++] = "Topic " + String(i);
  }
  saveTopics(chapterCount - 1);
  
  for (int i = 0; i < topicCount; i++) {
    saveContent(chapterCount - 1, i, "Content for " + name + " - " + topics[i] + "\n\nEdit this via WiFi.");
  }
}

void deleteChapter(int ch) {
  if (ch < 0 || ch >= chapterCount) return;
  
  for (int i = ch; i < chapterCount - 1; i++) {
    chapters[i] = chapters[i + 1];
    for (int t = 0; t < MAX_TOPICS; t++) {
      String oldPath = "/ch" + String(i + 2) + "_" + String(t + 1) + ".txt";
      String newPath = "/ch" + String(i + 1) + "_" + String(t + 1) + ".txt";
      if (SPIFFS.exists(oldPath)) {
        File f = SPIFFS.open(oldPath, FILE_READ);
        String content = f.readString();
        f.close();
        File wf = SPIFFS.open(newPath, FILE_WRITE);
        wf.print(content);
        wf.close();
        SPIFFS.remove(oldPath);
      }
    }
    String oldTopic = "/ch" + String(i + 2) + ".json";
    String newTopic = "/ch" + String(i + 1) + ".json";
    if (SPIFFS.exists(oldTopic)) SPIFFS.rename(oldTopic, newTopic);
  }
  chapterCount--;
  saveChapters();
}

void addTopic(int ch, String name) {
  if (ch < 0 || ch >= chapterCount) return;
  loadTopics(ch);
  if (topicCount >= MAX_TOPICS) return;
  topics[topicCount++] = name;
  saveTopics(ch);
  saveContent(ch, topicCount - 1, "Content for " + chapters[ch] + " - " + name + "\n\nEdit this via WiFi.");
}

void deleteTopic(int ch, int topic) {
  if (ch < 0 || ch >= chapterCount) return;
  loadTopics(ch);
  if (topic < 0 || topic >= topicCount) return;
  
  for (int i = topic; i < topicCount - 1; i++) {
    topics[i] = topics[i + 1];
    String oldPath = "/ch" + String(ch + 1) + "_" + String(i + 2) + ".txt";
    String newPath = "/ch" + String(ch + 1) + "_" + String(i + 1) + ".txt";
    if (SPIFFS.exists(oldPath)) {
      File f = SPIFFS.open(oldPath, FILE_READ);
      String content = f.readString();
      f.close();
      File wf = SPIFFS.open(newPath, FILE_WRITE);
      wf.print(content);
      wf.close();
      SPIFFS.remove(oldPath);
    }
  }
  topicCount--;
  saveTopics(ch);
}

void renameChapter(int ch, String name) {
  if (ch >= 0 && ch < chapterCount) {
    chapters[ch] = name;
    saveChapters();
  }
}

void renameTopic(int ch, int topic, String name) {
  loadTopics(ch);
  if (topic >= 0 && topic < topicCount) {
    topics[topic] = name;
    saveTopics(ch);
  }
}

// ================= WEB SERVER =================
void setupAPMode() {
  if (apEnabled) return;
  WiFi.mode(WIFI_AP);
  WiFi.softAP(ap_ssid, ap_password);
  apEnabled = true;
  setupWebServer();
  showMessage("AP Mode: FPV_CAM");
}

void stopAPMode() {
  if (!apEnabled) return;
  WiFi.softAPdisconnect(true);
  apEnabled = false;
}

void setupWebServer() {
  server.on("/", []() {
    String html = "<!DOCTYPE html><html>";
    html += "<head><meta name='viewport' content='width=device-width, initial-scale=1'>";
    html += "<title>E-Reader Manager</title>";
    html += "<style>";
    html += "body{font-family:Arial;background:#1a1a2e;margin:0;padding:20px;color:#eee}";
    html += ".container{max-width:800px;margin:auto}";
    html += ".card{background:#16213e;border-radius:10px;padding:20px;margin-bottom:20px}";
    html += "button{background:#e94560;color:white;padding:8px 15px;border:none;border-radius:5px;cursor:pointer;margin:5px}";
    html += "button:hover{opacity:0.8}";
    html += ".btn-green{background:#4CAF50}";
    html += ".btn-blue{background:#2196F3}";
    html += ".btn-orange{background:#ff9800}";
    html += ".btn-red{background:#f44336}";
    html += "input,textarea{width:100%;padding:10px;margin:10px 0;background:#0f3460;border:1px solid #e94560;border-radius:5px;color:white}";
    html += ".topic-item{background:#0a0a2a;padding:10px;margin:5px 0;border-radius:5px}";
    html += ".modal{display:none;position:fixed;top:0;left:0;width:100%;height:100%;background:rgba(0,0,0,0.8);justify-content:center;align-items:center}";
    html += ".modal-content{background:#16213e;padding:30px;border-radius:10px;max-width:400px}";
    html += "</style>";
    html += "</head><body>";
    html += "<div class='container'>";
    html += "<div class='card'><h2>E-Reader Manager</h2>";
    html += "<button class='btn-green' onclick='showAddChapter()'>Add Chapter</button>";
    html += "<button onclick='location.reload()'>Refresh</button></div>";
    
    for (int i = 0; i < chapterCount; i++) {
      html += "<div class='card'>";
      html += "<h3>" + chapters[i] + "</h3>";
      html += "<button class='btn-blue' onclick='renameChapter(" + String(i) + ",\"" + chapters[i] + "\")'>Rename</button>";
      html += "<button class='btn-red' onclick='deleteChapter(" + String(i) + ",\"" + chapters[i] + "\")'>Delete</button>";
      html += "<button class='btn-green' onclick='showAddTopic(" + String(i) + ")'>Add Topic</button>";
      
      loadTopics(i);
      for (int t = 0; t < topicCount; t++) {
        html += "<div class='topic-item'>";
        html += "<b>" + topics[t] + "</b><br>";
        html += "<button class='btn-orange' onclick='renameTopic(" + String(i) + "," + String(t) + ",\"" + topics[t] + "\")'>Rename</button>";
        html += "<button class='btn-red' onclick='deleteTopic(" + String(i) + "," + String(t) + ",\"" + topics[t] + "\")'>Delete</button>";
        html += "<button class='btn-green' onclick='editContent(" + String(i+1) + "," + String(t+1) + ")'>Edit</button>";
        html += "</div>";
      }
      html += "</div>";
    }
    
    html += "</div>";
    
    html += "<div id='addChapterModal' class='modal'><div class='modal-content'><h3>Add Chapter</h3><input id='chName' placeholder='Chapter name'><br><button class='btn-green' onclick='addChapter()'>Create</button><button onclick='closeModal()'>Cancel</button></div></div>";
    html += "<div id='addTopicModal' class='modal'><div class='modal-content'><h3>Add Topic</h3><input id='tpName' placeholder='Topic name'><br><button class='btn-green' onclick='addTopic()'>Create</button><button onclick='closeModal()'>Cancel</button></div></div>";
    html += "<div id='renameChapterModal' class='modal'><div class='modal-content'><h3>Rename Chapter</h3><input id='rcName'><br><button class='btn-orange' onclick='renameChapterSubmit()'>Rename</button><button onclick='closeModal()'>Cancel</button></div></div>";
    html += "<div id='renameTopicModal' class='modal'><div class='modal-content'><h3>Rename Topic</h3><input id='rtName'><br><button class='btn-orange' onclick='renameTopicSubmit()'>Rename</button><button onclick='closeModal()'>Cancel</button></div></div>";
    
    html += "<script>";
    html += "let chapterIdx, topicIdx;";
    html += "function showAddChapter(){document.getElementById('addChapterModal').style.display='flex';}";
    html += "function showAddTopic(idx){chapterIdx=idx;document.getElementById('addTopicModal').style.display='flex';}";
    html += "function renameChapter(idx,name){chapterIdx=idx;document.getElementById('rcName').value=name;document.getElementById('renameChapterModal').style.display='flex';}";
    html += "function renameTopic(idx,tidx,name){chapterIdx=idx;topicIdx=tidx;document.getElementById('rtName').value=name;document.getElementById('renameTopicModal').style.display='flex';}";
    html += "function addChapter(){let n=document.getElementById('chName').value;fetch('/api/add_chapter',{method:'POST',headers:{'Content-Type':'application/json'},body:JSON.stringify({name:n})}).then(()=>location.reload());}";
    html += "function addTopic(){let n=document.getElementById('tpName').value;fetch('/api/add_topic',{method:'POST',headers:{'Content-Type':'application/json'},body:JSON.stringify({chapter:chapterIdx,name:n})}).then(()=>location.reload());}";
    html += "function deleteChapter(idx,name){if(confirm('Delete '+name+'?')){fetch('/api/delete_chapter?chapter='+idx).then(()=>location.reload());}}";
    html += "function deleteTopic(idx,tidx,name){if(confirm('Delete '+name+'?')){fetch('/api/delete_topic?chapter='+idx+'&topic='+tidx).then(()=>location.reload());}}";
    html += "function renameChapterSubmit(){let n=document.getElementById('rcName').value;fetch('/api/rename_chapter',{method:'POST',headers:{'Content-Type':'application/json'},body:JSON.stringify({chapter:chapterIdx,name:n})}).then(()=>location.reload());}";
    html += "function renameTopicSubmit(){let n=document.getElementById('rtName').value;fetch('/api/rename_topic',{method:'POST',headers:{'Content-Type':'application/json'},body:JSON.stringify({chapter:chapterIdx,topic:topicIdx,name:n})}).then(()=>location.reload());}";
    html += "function editContent(ch,tp){window.location.href='/edit?c='+ch+'&t='+tp;}";
    html += "function closeModal(){document.querySelectorAll('.modal').forEach(m=>m.style.display='none');}";
    html += "window.onclick=function(e){if(e.target.classList.contains('modal'))e.target.style.display='none';}";
    html += "</script>";
    html += "</body></html>";
    server.send(200, "text/html", html);
  });
  
  server.on("/edit", []() {
    int ch = server.arg("c").toInt();
    int tp = server.arg("t").toInt();
    String content = "";
    String path = "/ch" + String(ch) + "_" + String(tp) + ".txt";
    File f = SPIFFS.open(path, FILE_READ);
    while (f.available()) content += (char)f.read();
    f.close();
    
    String html = "<!DOCTYPE html><html><head><meta name='viewport' content='width=device-width, initial-scale=1'>";
    html += "<title>Edit Content</title><style>";
    html += "body{font-family:Arial;background:#1a1a2e;padding:20px;color:#eee}";
    html += ".container{max-width:800px;margin:auto}";
    html += "textarea{width:100%;height:400px;padding:10px;background:#0f3460;border:1px solid #e94560;border-radius:5px;color:white;font-family:monospace}";
    html += "button{background:#4CAF50;color:white;padding:10px 20px;border:none;border-radius:5px;cursor:pointer;margin:10px 5px}";
    html += "</style></head><body>";
    html += "<div class='container'>";
    html += "<h2>Edit Content</h2>";
    html += "<form method='POST' action='/save'>";
    html += "<input type='hidden' name='c' value='" + String(ch) + "'>";
    html += "<input type='hidden' name='t' value='" + String(tp) + "'>";
    html += "<textarea name='content'>" + content + "</textarea><br>";
    html += "<button type='submit'>Save</button>";
    html += "<button type='button' onclick='location.href=\"/\"'>Back</button>";
    html += "</form></div></body></html>";
    server.send(200, "text/html", html);
  });
  
  server.on("/save", HTTP_POST, []() {
    int ch = server.arg("c").toInt();
    int tp = server.arg("t").toInt();
    String content = server.arg("content");
    String path = "/ch" + String(ch) + "_" + String(tp) + ".txt";
    File f = SPIFFS.open(path, FILE_WRITE);
    f.print(content);
    f.close();
    server.send(200, "text/html", "<html><body><script>location.href='/';</script></body></html>");
  });
  
  server.on("/api/add_chapter", HTTP_POST, []() {
    DynamicJsonDocument doc(256);
    deserializeJson(doc, server.arg("plain"));
    addChapter(doc["name"].as<String>());
    server.send(200, "application/json", "{\"success\":true}");
  });
  
  server.on("/api/delete_chapter", HTTP_GET, []() {
    deleteChapter(server.arg("chapter").toInt());
    server.send(200, "application/json", "{\"success\":true}");
  });
  
  server.on("/api/add_topic", HTTP_POST, []() {
    DynamicJsonDocument doc(256);
    deserializeJson(doc, server.arg("plain"));
    addTopic(doc["chapter"].as<int>(), doc["name"].as<String>());
    server.send(200, "application/json", "{\"success\":true}");
  });
  
  server.on("/api/delete_topic", HTTP_GET, []() {
    deleteTopic(server.arg("chapter").toInt(), server.arg("topic").toInt());
    server.send(200, "application/json", "{\"success\":true}");
  });
  
  server.on("/api/rename_chapter", HTTP_POST, []() {
    DynamicJsonDocument doc(256);
    deserializeJson(doc, server.arg("plain"));
    renameChapter(doc["chapter"].as<int>(), doc["name"].as<String>());
    server.send(200, "application/json", "{\"success\":true}");
  });
  
  server.on("/api/rename_topic", HTTP_POST, []() {
    DynamicJsonDocument doc(256);
    deserializeJson(doc, server.arg("plain"));
    renameTopic(doc["chapter"].as<int>(), doc["topic"].as<int>(), doc["name"].as<String>());
    server.send(200, "application/json", "{\"success\":true}");
  });
  
  server.begin();
}

// ================= DISPLAY FUNCTIONS =================
String getWeekday() {
  String days[7] = {"SUN", "MON", "TUE", "WED", "THU", "FRI", "SAT"};
  return days[weekDay];
}

void displayWatchFace() {
  display.clearDisplay();
  
  String timeStr = "";
  if (hours < 10) timeStr += "0";
  timeStr += String(hours) + ":";
  if (minutes < 10) timeStr += "0";
  timeStr += String(minutes);
  
  String dateStr = getWeekday() + " ";
  if (day < 10) dateStr += "0";
  dateStr += String(day) + "/";
  if (month < 10) dateStr += "0";
  dateStr += String(month);
  
  String secStr = "";
  if (seconds < 10) secStr += "0";
  secStr += String(seconds);
  
  display.setTextSize(3);
  int16_t x1, y1;
  uint16_t w, h;
  display.getTextBounds(timeStr, 0, 0, &x1, &y1, &w, &h);
  display.setCursor((128 - w) / 2, 15);
  display.print(timeStr);
  
  display.setTextSize(1);
  display.getTextBounds(secStr, 0, 0, &x1, &y1, &w, &h);
  display.setCursor((128 - w) / 2, 42);
  display.print(secStr);
  
  display.getTextBounds(dateStr, 0, 0, &x1, &y1, &w, &h);
  display.setCursor((128 - w) / 2, 54);
  display.print(dateStr);
  
  display.display();
}

void showMessage(String msg) {
  display.clearDisplay();
  display.setCursor(10, 30);
  display.print(msg);
  display.display();
  delay(1500);
}

void displayMenu() {
  display.clearDisplay();
  display.setCursor(0, 0);
  display.println("E-Reader");
  display.drawLine(0, 10, 128, 10, WHITE);
  
  const char* items[] = {"Start Reading", "AP Mode", "Update Time", "Exit"};
  for (int i = 0; i < 4; i++) {
    display.setCursor(0, 20 + i * 10);
    if (i == readerCursor) {
      display.print("> ");
      display.setTextColor(BLACK, WHITE);
      display.print(items[i]);
      display.setTextColor(WHITE);
    } else {
      display.print("  ");
      display.print(items[i]);
    }
  }
  display.display();
}

void displayChapters() {
  display.clearDisplay();
  display.setCursor(0, 0);
  display.println("Chapters");
  display.drawLine(0, 10, 128, 10, WHITE);
  
  if (selectedChapter < readerScrollTop) readerScrollTop = selectedChapter;
  if (selectedChapter >= readerScrollTop + VISIBLE_ITEMS) readerScrollTop = selectedChapter - VISIBLE_ITEMS + 1;
  
  for (int i = 0; i < VISIBLE_ITEMS; i++) {
    int idx = readerScrollTop + i;
    if (idx >= chapterCount) break;
    int y = 14 + i * 10;
    String txt = chapters[idx];
    if (txt.length() > 16) txt = txt.substring(0, 13) + "...";
    
    if (idx == selectedChapter) {
      display.fillRect(0, y-1, 128, 10, WHITE);
      display.setTextColor(BLACK);
      display.setCursor(2, y);
      display.print("> ");
      display.print(txt);
      display.setTextColor(WHITE);
    } else {
      display.setCursor(2, y);
      display.print("  ");
      display.print(txt);
    }
  }
  display.setCursor(0, 54);
  display.print(selectedChapter+1);
  display.print("/");
  display.print(chapterCount);
  display.display();
}

void displayTopics() {
  display.clearDisplay();
  display.setCursor(0, 0);
  String title = currentChapter;
  if (title.length() > 16) title = title.substring(0, 13) + "...";
  display.print(title);
  display.drawLine(0, 10, 128, 10, WHITE);
  
  if (selectedTopic < readerScrollTop) readerScrollTop = selectedTopic;
  if (selectedTopic >= readerScrollTop + VISIBLE_ITEMS) readerScrollTop = selectedTopic - VISIBLE_ITEMS + 1;
  
  for (int i = 0; i < VISIBLE_ITEMS; i++) {
    int idx = readerScrollTop + i;
    if (idx >= topicCount) break;
    int y = 14 + i * 10;
    String txt = topics[idx];
    if (txt.length() > 16) txt = txt.substring(0, 13) + "...";
    
    if (idx == selectedTopic) {
      display.fillRect(0, y-1, 128, 10, WHITE);
      display.setTextColor(BLACK);
      display.setCursor(2, y);
      display.print("> ");
      display.print(txt);
      display.setTextColor(WHITE);
    } else {
      display.setCursor(2, y);
      display.print("  ");
      display.print(txt);
    }
  }
  display.setCursor(0, 54);
  display.print(selectedTopic+1);
  display.print("/");
  display.print(topicCount);
  display.display();
}

void displayContent() {
  int startLine = scrollPosition;
  if (startLine < 0) startLine = 0;
  if (startLine >= lineCount) startLine = lineCount - LINES_PER_SCREEN;
  if (startLine < 0) startLine = 0;
  
  display.clearDisplay();
  for (int i = 0; i < LINES_PER_SCREEN; i++) {
    int idx = startLine + i;
    if (idx >= lineCount) break;
    display.setCursor(0, i * 8);
    String line = textLines[idx];
    if (line.length() > 21) line = line.substring(0, 18) + "..";
    display.println(line);
  }
  display.display();
}

void displayAPScreen() {
  display.clearDisplay();
  display.setCursor(0, 0);
  display.println("AP Mode");
  display.drawLine(0, 10, 128, 10, WHITE);
  display.setCursor(15, 25);
  display.println("FPV_CAM");
  display.setCursor(15, 40);
  display.println("Open 192.168.4.1");
  display.setCursor(15, 55);
  display.println("Press BACK to exit");
  display.display();
}

// ================= SLEEP FUNCTION WITH TIME TRACKING =================
void goToSleep() {
  if (currentMode == MODE_READER && readerScreen == SCR_READ) {
    rtc_lastChapter = selectedChapter;
    rtc_lastTopic = selectedTopic;
    rtc_lastScrollPos = scrollPosition;
    rtc_hasResumeData = true;
  }
  
  rtc_sleepEnterTime = esp_timer_get_time();
  rtc_unixTime = getCurrentUnixTime();
  rtc_timeSet = true;
  
  display.clearDisplay();
  display.display();
  digitalWrite(OLED_POWER_PIN, LOW);
  
  if (apEnabled) stopAPMode();
  WiFi.mode(WIFI_OFF);
  
  esp_sleep_enable_gpio_wakeup();
  gpio_wakeup_enable((gpio_num_t)BTN_SELECT, GPIO_INTR_LOW_LEVEL);
  gpio_wakeup_enable((gpio_num_t)BTN_UP, GPIO_INTR_LOW_LEVEL);
  gpio_wakeup_enable((gpio_num_t)BTN_DOWN, GPIO_INTR_LOW_LEVEL);
  gpio_wakeup_enable((gpio_num_t)BTN_BACK, GPIO_INTR_LOW_LEVEL);
  
  Serial.println("Entering light sleep...");
  esp_light_sleep_start();
  
  uint64_t wakeTime = esp_timer_get_time();
  uint64_t sleepDuration = (wakeTime - rtc_sleepEnterTime) / 1000000ULL;
  rtc_unixTime += sleepDuration;
  
  Serial.printf("Slept for %llu seconds\n", sleepDuration);
  
  updateTimeFromUnix();
  
  digitalWrite(OLED_POWER_PIN, HIGH);
  Wire.begin(SDA_PIN, SCL_PIN);
  display.begin(SSD1306_SWITCHCAPVCC, SCREEN_ADDRESS);
  display.setRotation(2);
  display.ssd1306_command(SSD1306_SETCONTRAST);
  display.ssd1306_command(brightness);
  
  lastActivityTime = millis();
  lastTick = millis() - 1000;
  
  if (currentMode == MODE_WATCH) {
    displayWatchFace();
  } else {
    displayMenu();
  }
}

void handleButtons() {
  unsigned long now = millis();
  if (now - lastButtonTime < DEBOUNCE_DELAY) return;
  
  bool up = digitalRead(BTN_UP) == LOW;
  bool down = digitalRead(BTN_DOWN) == LOW;
  bool sel = digitalRead(BTN_SELECT) == LOW;
  bool back = digitalRead(BTN_BACK) == LOW;
  
  if (up || down || sel || back) lastActivityTime = now;
  
  if (currentMode == MODE_WATCH && up && back) {
    lastButtonTime = now;
    goToSleep();
    return;
  }
  
  if (currentMode == MODE_WATCH && sel && up) {
    lastButtonTime = now;
    currentMode = MODE_READER;
    readerScreen = SCR_MAIN;
    readerCursor = 0;
    if (rtc_hasResumeData && rtc_lastChapter < chapterCount) {
      selectedChapter = rtc_lastChapter;
      selectedTopic = rtc_lastTopic;
      scrollPosition = rtc_lastScrollPos;
      currentChapter = chapters[selectedChapter];
      loadTopics(selectedChapter);
      loadContent(selectedChapter, selectedTopic);
      readerScreen = SCR_READ;
    }
    displayMenu();
    return;
  }
  
  if ((currentMode == MODE_READER || currentMode == MODE_AP) && sel && back) {
    lastButtonTime = now;
    if (currentMode == MODE_READER && readerScreen == SCR_READ) {
      rtc_lastChapter = selectedChapter;
      rtc_lastTopic = selectedTopic;
      rtc_lastScrollPos = scrollPosition;
      rtc_hasResumeData = true;
      rtc_unixTime = getCurrentUnixTime();
      rtc_timeSet = true;
    }
    if (apEnabled) stopAPMode();
    currentMode = MODE_WATCH;
    readerScreen = SCR_MAIN;
    displayWatchFace();
    return;
  }
  
  if (currentMode == MODE_READER) {
    if (up && !down && !sel && !back) {
      lastButtonTime = now;
      if (readerScreen == SCR_MAIN && readerCursor > 0) readerCursor--;
      else if (readerScreen == SCR_CHAPTERS && selectedChapter > 0) selectedChapter--;
      else if (readerScreen == SCR_TOPICS && selectedTopic > 0) selectedTopic--;
      else if (readerScreen == SCR_READ && scrollPosition > 0) scrollPosition--;
      
      if (readerScreen == SCR_READ) displayContent();
      else if (readerScreen == SCR_CHAPTERS) displayChapters();
      else if (readerScreen == SCR_TOPICS) displayTopics();
      else displayMenu();
    }
    else if (down && !up && !sel && !back) {
      lastButtonTime = now;
      if (readerScreen == SCR_MAIN && readerCursor < 3) readerCursor++;
      else if (readerScreen == SCR_CHAPTERS && selectedChapter < chapterCount-1) selectedChapter++;
      else if (readerScreen == SCR_TOPICS && selectedTopic < topicCount-1) selectedTopic++;
      else if (readerScreen == SCR_READ && scrollPosition + LINES_PER_SCREEN < lineCount) scrollPosition++;
      
      if (readerScreen == SCR_READ) displayContent();
      else if (readerScreen == SCR_CHAPTERS) displayChapters();
      else if (readerScreen == SCR_TOPICS) displayTopics();
      else displayMenu();
    }
    else if (sel && !up && !down && !back) {
      lastButtonTime = now;
      if (readerScreen == SCR_MAIN) {
        if (readerCursor == 0) {
          readerScreen = SCR_CHAPTERS;
          selectedChapter = 0;
          displayChapters();
        } else if (readerCursor == 1) {
          currentMode = MODE_AP;
          setupAPMode();
          displayAPScreen();
        } else if (readerCursor == 2) {
          syncTimeWithWiFi();
          displayMenu();
        } else if (readerCursor == 3) {
          if (apEnabled) stopAPMode();
          currentMode = MODE_WATCH;
          displayWatchFace();
        }
      }
      else if (readerScreen == SCR_CHAPTERS) {
        currentChapter = chapters[selectedChapter];
        loadTopics(selectedChapter);
        readerScreen = SCR_TOPICS;
        selectedTopic = 0;
        displayTopics();
      }
      else if (readerScreen == SCR_TOPICS) {
        loadContent(selectedChapter, selectedTopic);
        readerScreen = SCR_READ;
        scrollPosition = 0;
        displayContent();
      }
    }
    else if (back && !up && !down && !sel) {
      lastButtonTime = now;
      if (readerScreen == SCR_MAIN) {
        if (apEnabled) stopAPMode();
        currentMode = MODE_WATCH;
        displayWatchFace();
      }
      else if (readerScreen == SCR_CHAPTERS) {
        readerScreen = SCR_MAIN;
        readerCursor = 0;
        displayMenu();
      }
      else if (readerScreen == SCR_TOPICS) {
        readerScreen = SCR_CHAPTERS;
        displayChapters();
      }
      else if (readerScreen == SCR_READ) {
        readerScreen = SCR_TOPICS;
        displayTopics();
      }
    }
  }
  else if (currentMode == MODE_AP) {
    server.handleClient();
    if (back && !up && !down && !sel) {
      lastButtonTime = now;
      stopAPMode();
      currentMode = MODE_READER;
      readerScreen = SCR_MAIN;
      displayMenu();
    }
  }
}

// ================= MAIN LOOP =================
void loop() {
  updateTimeCounter();
  handleButtons();
  
  if (currentMode == MODE_WATCH) {
    if (millis() - lastActivityTime >= WATCH_SLEEP_TIMEOUT_MS) {
      goToSleep();
    }
  }
  
  delay(10);
}
