#include <Wire.h>
#include <Adafruit_GFX.h>
#include <Adafruit_SSD1306.h>
#include <U8g2lib.h>
#include <Fonts/FreeSans9pt7b.h>

#define FONT          u8g2_font_8x13_tr
#define LITTLE_FONT   u8g2_font_5x8_tr
#define LINE_SPACING  16

// OLED дисплей 128x64, I2C адрес по умолчанию — 0x3C
#define SCREEN_WIDTH 128
#define SCREEN_HEIGHT 64
#define OLED_RESET     -1
Adafruit_SSD1306 display(SCREEN_WIDTH, SCREEN_HEIGHT, &Wire, OLED_RESET);

// Инициализация дисплея SSD1306 128x64 по I2C
// Выбери нужный конструктор в зависимости от контроллера
// Этот - для SSD1306, I2C, с аппаратным reset (если нет reset - поставить U8G2_R0 и -1)
U8G2_SSD1306_128X64_NONAME_F_HW_I2C u8g2(U8G2_R0, /* reset=*/ U8X8_PIN_NONE);


// Кнопки
#define BTN_UP_PIN     3
#define BTN_DOWN_PIN   2
#define BTN_OK_PIN     1   // #
#define BTN_BACK_PIN   0   // *

extern uint16_t startAddr;
extern uint16_t versionAddr;
extern void writeStringToEEPROM(const char* str, uint16_t addr);
extern void readModuleFromEEPROM();
extern void readVersionFromEEPROM();
extern char versionStr[32];
extern char moduleStr[32];


enum MenuState {
  MENU_MAIN,
  MENU_SELECT_MODEL,
  MENU_SELECT_OPTION,
  MENU_SELECT_VERSION,
  MENU_SHOW_EEPROM
};


const char* modelList[] = { "DPO4", "DPO3", "DPO2", "DPO1", "MDO4", "MDO3", "MDO2", "MDO1"};
const uint8_t modelCount = sizeof(modelList) / sizeof(modelList[0]);

// Фильтрация опций по модели
uint8_t filteredOptionIndexes[50];
uint8_t filteredOptionCount = 0;

const char* optionList[] = {
  "BND",    // Bundle (all-in-one)
  "AERO",   // Aerospace serial buses
  "AFG",    // Arbitrary Function Generator
  "AUDIO",  // Audio serial buses
  "AUTO",   // Automotive serial buses
  "AUTOMAX",// Automotive serial buses
  "COMP",   // Computer serial buses
  "EMBD",   // Embedded systems
  "ENET",   // Ethernet
  "FLEX",   // FlexRay
  "LMT",    // Limit testing and mask testing
  "MSO",    // Mixed Signal Option
  "PWR",    // Power analysis
  "SA",     // Spectrum Analyzer
  "SA3",     // Spectrum Analyzer
  "SA3T6",     // Spectrum Analyzer
  "SA6",     // Spectrum Analyzer
  "SEC",    // Security extension
  "TRIG",   // RF Power Level Trigger
  "USB",    // USB analysis
  "VID"     // Video signal triggering
};
const uint8_t optionCount = sizeof(optionList) / sizeof(optionList[0]);

const char* optionDescriptions[] = {
  "Include all available options in one",
  "Aerospace buses: MIL-STD-1553",
  "Waveform generator for custom signals",
  "Audio buses: I2S, LJ, RJ, TDM",
  "Cars buses: CAN, CAN FD, LIN",
  "Cars buses: CAN, CAN FD, LIN, FlexRay",
  "Computer buses: RS-232/422/485, UART",
  "Embedded protocols: I2C, SPI",
  "Ethernet: 10BASE-T, 100BASE-TX",
  "Cars bus FlexRay",
  "Limit and mask testing: Go/No-Go",
  "Mixed Signal Oscill: analog + digital",
  "Power: efficiency, ripple, dI/dt, dV/dt",
  "Spectrum analysis: RF, high-frequency",
  "Spectr analysis (3 MHz): RF, high-frequency",
  "SA3 extention to 6 MHz for MDO4000C",
  "Spectr analysis (6 MHz): RF, high-frequency",
  "Security extension: access control",
  "RF power level triggering",
  "USB decoding: USB LS, FS, HS",
  "Video triggering: HDTV"
};


bool isOptionForDPO[] = {
  true,   // BND     → DPO4BND
  true,   // AERO    → DPO4AERO
  false,  // AFG     → (не используется в DPO)
  true,   // AUDIO   → DPO4AUDIO
  true,   // AUTO    → DPO4AUTO
  true,   // AUTOMAX → нет в DPO
  true,   // COMP    → DPO4COMP
  true,   // EMBD    → DPO4EMBD
  true,   // ENET    → DPO4ENET
  true,   // FLEX    → нет упоминания, но оставим
  true,   // LMT     → DPO4LMT
  false,  // MSO     → только для MDO (MDO4MSO)
  true,   // PWR     → DPO4PWR
  false,  // SA      → только MDO
  false,  // SA3     → только MDO
  false,  // SA3T6   → только MDO
  false,  // SA6     → только MDO
  false,  // SEC     → только MDO (MDO4SEC)
  false,  // TRIG    → только MDO (MDO4TRIG)
  true,   // USB     → DPO4USB
  true    // VID     → DPO4VID
};

bool isOptionForMDO[] = {
  false,  // BND     → нет прямого MDO4BND
  false,  // AERO
  true,   // AFG     → MDO4AFG
  false,  // AUDIO
  false,  // AUTO
  false,  // AUTOMAX → пусть будет для MDO
  false,  // COMP
  false,  // EMBD
  false,  // ENET
  false,  // FLEX
  false,  // LMT
  true,   // MSO     → MDO4MSO
  false,  // PWR
  true,   // SA      → MDO4SA3/MDO4SA6
  true,   // SA3     → MDO4SA3
  true,   // SA3     → MDO4SA3T6
  true,   // SA6     → MDO4SA6
  true,   // SEC     → MDO4SEC
  true,   // TRIG    → MDO4TRIG
  false,  // USB
  false   // VID
};

void filterOptionsForModel(const char* modelName) {
  filteredOptionCount = 0;

  bool isMDO = strstr(modelName, "MDO") != nullptr;
  bool isDPO = strstr(modelName, "DPO") != nullptr;

  for (uint8_t i = 0; i < optionCount; ++i) {
    if ((isMDO && isOptionForMDO[i]) || (isDPO && isOptionForDPO[i])) {
      filteredOptionIndexes[filteredOptionCount++] = i;
    }
  }

  // Подстраховка
  if (filteredOptionCount == 0) {
    for (uint8_t i = 0; i < sizeof(optionList); ++i)
      filteredOptionIndexes[filteredOptionCount++] = i;
  }
}


uint8_t selectedModel = 0;
uint8_t selectedOption = 0;

MenuState currentMenu = MENU_MAIN;
uint8_t cursorPos = 0;
char editableVersion[16] = "v1.00";
uint8_t editIndex = 0;
bool editingVersion = false;


void setupUI() {
  pinMode(BTN_UP_PIN, INPUT_PULLUP);
  pinMode(BTN_DOWN_PIN, INPUT_PULLUP);
  pinMode(BTN_OK_PIN, INPUT_PULLUP);
  pinMode(BTN_BACK_PIN, INPUT_PULLUP);

  if (!display.begin(SSD1306_SWITCHCAPVCC, 0x3C)) {
    Serial.println(F("SSD1306 not found"));
    while (true);
  }
  //display.clearDisplay();
  //display.setFont(&FONT);
  //display.setTextSize(1);
  //display.setTextColor(SSD1306_WHITE);
  //display.display();

  u8g2.begin();

  // Выбираем шрифт - очень маленький, около 6pt (пример: u8g2_font_6x10_tr)
  u8g2.setFont(FONT);


  showMainMenu();
}

void loopUI() {
  static uint32_t lastPress = 0;
  if (millis() - lastPress < 200) return;

  if (currentMenu == MENU_SELECT_VERSION && editingVersion) {
    if (!digitalRead(BTN_UP_PIN)) {
      if (editableVersion[editIndex] >= '0' && editableVersion[editIndex] <= '9') {
        if (editableVersion[editIndex] == '9') editableVersion[editIndex] = '0';
        else editableVersion[editIndex]++;
      }
      updateMenuDisplay();
      lastPress = millis();
      return;
    }

    if (!digitalRead(BTN_DOWN_PIN)) {
      if (editableVersion[editIndex] >= '0' && editableVersion[editIndex] <= '9') {
        if (editableVersion[editIndex] == '0') editableVersion[editIndex] = '9';
        else editableVersion[editIndex]--;
      }
      updateMenuDisplay();
      lastPress = millis();
      return;
    }

    if (!digitalRead(BTN_OK_PIN)) {
      // Перейти к следующему редактируемому символу, пропуская точки и буквы
      do {
        editIndex++;
      } while (editIndex < strlen(editableVersion) && (editableVersion[editIndex] < '0' || editableVersion[editIndex] > '9'));

      if (editIndex >= strlen(editableVersion)) {
        // Закончили редактировать
        writeStringToEEPROM(editableVersion, versionAddr);
        strncpy(versionStr, editableVersion, sizeof(versionStr));
        editingVersion = false;
        currentMenu = MENU_MAIN;
        cursorPos = 0;
      }

      updateMenuDisplay();
      lastPress = millis();
      return;
    }

    if (!digitalRead(BTN_BACK_PIN)) {
      // Назад по цифрам
      do {
        if (editIndex == 0) {
          editingVersion = false;
          currentMenu = MENU_MAIN;
          cursorPos = 0;
          break;
        } else {
          editIndex--;
        }
      } while (editableVersion[editIndex] < '0' || editableVersion[editIndex] > '9');

      updateMenuDisplay();
      lastPress = millis();
      return;
    }

    return;
  }


  if (!digitalRead(BTN_UP_PIN)) {
    if (cursorPos > 0) cursorPos--;
    updateMenuDisplay();
    lastPress = millis();
  } else if (!digitalRead(BTN_DOWN_PIN)) {
    cursorPos++;
    updateMenuDisplay();
    lastPress = millis();
  } else if (!digitalRead(BTN_OK_PIN)) {
    handleOk();
    lastPress = millis();
  } else if (!digitalRead(BTN_BACK_PIN)) {
    handleBack();
    lastPress = millis();
  }
}

void showMainMenu() {
  currentMenu = MENU_MAIN;
  cursorPos = 0;
  updateMenuDisplay();
}

void drawStringWrapped(int x, int y, int maxWidth, const char* text, int lineHeight) {
  int lineStart = 0;
  int textLen = strlen(text);

  while (lineStart < textLen) {
    int lineEnd = lineStart;
    int lastSpace = -1;
    int width = 0;

    while (lineEnd < textLen) {
      char c = text[lineEnd];
      // Вычисляем ширину подстроки с lineStart до lineEnd + 1
      char buffer[64];  // должен быть достаточного размера
      int length = lineEnd - lineStart + 1;
      if (length >= sizeof(buffer)) length = sizeof(buffer) - 1;
      memcpy(buffer, text + lineStart, length);
      buffer[length] = '\0';

      int w = u8g2.getStrWidth(buffer);

      if (w > maxWidth) {
        break;
      }

      if (c == ' ' || c == '-' || c == '\t') {
        lastSpace = lineEnd;
      }

      lineEnd++;
    }

    // Если мы вышли за предел maxWidth
    if (lineEnd == lineStart) {
      // Если один символ не помещается — выводим его все равно (чтобы не зацикливаться)
      lineEnd = lineStart + 1;
    } else if (lineEnd < textLen && lastSpace != -1) {
      // Откатываемся до последнего пробела, чтобы не разрывать слово
      lineEnd = lastSpace;
    }

    // Подстрока для вывода длиной (lineEnd - lineStart)
    char buf[128];
    int len = lineEnd - lineStart;
    if (len >= sizeof(buf)) len = sizeof(buf) - 1;
    strncpy(buf, text + lineStart, len);
    buf[len] = 0;

    u8g2.drawStr(x, y, buf);

    y += lineHeight;
    lineStart = (lineEnd == lastSpace) ? lineEnd + 1 : lineEnd;
  }
}


void updateMenuDisplay() {
  u8g2.clearBuffer();
  //display.clearDisplay();
  //display.setCursor(0, LINE_SPACING);

  if (currentMenu == MENU_MAIN) {
    // Выводим три строки меню с курсором >
    if(cursorPos > 2)
    {
      cursorPos = 2;
    }
    u8g2.drawStr(0, LINE_SPACING * 1, cursorPos == 0 ? ">Set Options" : " Set Options");
    u8g2.drawStr(0, LINE_SPACING * 2, cursorPos == 1 ? ">Set Version" : " Set Version");
    u8g2.drawStr(0, LINE_SPACING * 3, cursorPos == 2 ? ">Read Module" : " Read Module");
    /*display.print(cursorPos == 0 ? ">Set Options" : " Set Options");
    display.setCursor(0, LINE_SPACING *2);
    display.print(cursorPos == 1 ? ">Set Version" : " Set Version");
    display.setCursor(0, LINE_SPACING *3);
    display.print(cursorPos == 2 ? ">Read" : " Read");*/

  } else if (currentMenu == MENU_SELECT_MODEL) {
    u8g2.drawStr(0, LINE_SPACING * 1, "Model:");

    // Формируем строку ">ModelName"
    char line[20];
    snprintf(line, sizeof(line), ">%s", modelList[cursorPos % modelCount]);
    u8g2.drawStr(0, LINE_SPACING * 2, line);
    /*display.print("Model:");
    display.setCursor(0, LINE_SPACING *2);
    display.print(">");
    display.print(modelList[cursorPos % modelCount]);*/

  } else if (currentMenu == MENU_SELECT_OPTION) {
    u8g2.setFont(FONT);  // основной шрифт
    u8g2.drawStr(0, LINE_SPACING * 1, "Option:");

    if (cursorPos >= filteredOptionCount) cursorPos = filteredOptionCount - 1;

    char line[20];
    uint8_t optIndex = filteredOptionIndexes[cursorPos];
    snprintf(line, sizeof(line), ">%s", optionList[optIndex]);
    u8g2.drawStr(0, LINE_SPACING * 2, line);

    // Мелкий шрифт — описание
    u8g2.setFont(LITTLE_FONT);
    drawStringWrapped(0, LINE_SPACING * 3 + 4, SCREEN_WIDTH, optionDescriptions[optIndex], 10);
    u8g2.setFont(FONT);
    /*display.print("Option:");
    display.setCursor(0, LINE_SPACING *2);
    display.print(">");
    display.print(optionList[cursorPos % optionCount]);*/

  } else if (currentMenu == MENU_SELECT_VERSION) {
    u8g2.drawStr(0, LINE_SPACING * 1, "UART set ver:");

    // Для версии выводим символы с квадратными скобками вокруг редактируемого
    int x = 0;
    for (uint8_t i = 0; i < strlen(editableVersion); i++) {
      char buf[4];
      if (i == editIndex && editingVersion) {
        snprintf(buf, sizeof(buf), "[%c]", editableVersion[i]);
      } else {
        snprintf(buf, sizeof(buf), " %c ", editableVersion[i]);
      }
      u8g2.drawStr(x, LINE_SPACING * 2, buf);
      x += u8g2.getStrWidth(buf); // смещение курсора вправо на ширину предыдущей части
    }


    /*display.print("UART set ver:");
    display.setCursor(0, LINE_SPACING *2);

    for (uint8_t i = 0; i < strlen(editableVersion); i++) {
      if (i == editIndex && editingVersion) {
        display.print("[");
        display.print(editableVersion[i]);
        display.print("]");
      } else {
        display.print(" ");
        display.print(editableVersion[i]);
        display.print(" ");
      }
    }*/
  } else if (currentMenu == MENU_SHOW_EEPROM) {
    char buf[40];

    snprintf(buf, sizeof(buf), "Model Option:");
    u8g2.drawStr(0, LINE_SPACING * 1, buf);

    snprintf(buf, sizeof(buf), "  %s", moduleStr);
    u8g2.drawStr(0, LINE_SPACING * 2, buf);

    snprintf(buf, sizeof(buf), "Version:");
    u8g2.drawStr(0, LINE_SPACING * 3, buf);

    snprintf(buf, sizeof(buf), "  %s", versionStr);
    u8g2.drawStr(0, LINE_SPACING * 4, buf);
  
    /*display.print("Mod:");
    display.println(moduleStr);
    display.setCursor(0, LINE_SPACING *3);
    display.print("Ver:");
    display.println(versionStr);*/
  }

  //display.display();
  u8g2.sendBuffer();
}


void handleOk() {
  if (currentMenu == MENU_MAIN) {
    if (cursorPos == 0) {
      currentMenu = MENU_SELECT_MODEL;
      cursorPos = 0;
    } else if (cursorPos == 1) {
      currentMenu = MENU_SELECT_VERSION;
    } else if (cursorPos == 2) {
      readVersionFromEEPROM();
      readModuleFromEEPROM();
      currentMenu = MENU_SHOW_EEPROM;
    }

  } else if (currentMenu == MENU_SELECT_MODEL) {
    selectedModel = cursorPos % modelCount;
    filterOptionsForModel(modelList[selectedModel]);  // ← ВАЖНО!
    cursorPos = 0;
    currentMenu = MENU_SELECT_OPTION;

  } else if (currentMenu == MENU_SELECT_OPTION) {
    selectedOption = filteredOptionIndexes[cursorPos];;

    char fullModule[32];
    snprintf(fullModule, sizeof(fullModule), "%s%s", modelList[selectedModel], optionList[selectedOption]);
    writeStringToEEPROM(fullModule, startAddr);
    readAndDumpEEPROM();
    strncpy(moduleStr, fullModule, sizeof(moduleStr) - 1);
    moduleStr[sizeof(moduleStr) - 1] = 0;

    Serial.print("[UI] Options written: ");
    Serial.println(fullModule);

    currentMenu = MENU_MAIN;
    cursorPos = 0;

  } else if (currentMenu == MENU_SELECT_VERSION) {
    editingVersion = true;
    editIndex = 1;
  }


  updateMenuDisplay();
}




void handleBack() {
  if (currentMenu == MENU_SELECT_OPTION) {
    currentMenu = MENU_SELECT_MODEL;
  } else {
    currentMenu = MENU_MAIN;
  }
  cursorPos = 0;
  updateMenuDisplay();
}