#include <Wire.h>
#include <RTClib.h>
#include <Adafruit_GFX.h>
#include <Adafruit_SSD1306.h>
#include <TOTP.h>  // Includes SHA1
#include <EEPROM.h>

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

//RTC_DS3231 rtc;  
// Or 
RTC_DS1307 rtc;

#define MAX_ACCOUNTS 10
#define NAME_LEN 16
#define SECRET_LEN 20

struct Account {
  char name[NAME_LEN];
  uint8_t secret[SECRET_LEN];
  uint8_t secretLen;
};

Account accounts[MAX_ACCOUNTS];
int numAccounts = 0;
int selected = 0;

enum State { CLOCK, MENU, OTP };
State state = CLOCK;

const int btnUp = 2, btnDown = 3, btnBack = 4, btnSelect = 5;

// Load accounts from EEPROM
void loadAccounts() {
  int addr = 0;
  EEPROM.get(addr, numAccounts);
  addr += sizeof(int);
  for (int i = 0; i < numAccounts; i++) {
    EEPROM.get(addr, accounts[i]);
    addr += sizeof(Account);
  }
}

// Save accounts to EEPROM
void saveAccounts() {
  int addr = 0;
  EEPROM.put(addr, numAccounts);
  addr += sizeof(int);
  for (int i = 0; i < numAccounts; i++) {
    EEPROM.put(addr, accounts[i]);
    addr += sizeof(Account);
  }
}

void setup() {
  Serial.begin(9600);
  pinMode(btnUp, INPUT_PULLUP);
  pinMode(btnDown, INPUT_PULLUP);
  pinMode(btnBack, INPUT_PULLUP);
  pinMode(btnSelect, INPUT_PULLUP);

  if (!rtc.begin()) {
    Serial.println("Couldn't find RTC");
    while (1);
  }
  /*if (rtc.lostPower()) {
    Serial.println("RTC lost power, set time!");
    // Set via serial or buttons here
  }*/

  if (!display.begin(SSD1306_SWITCHCAPVCC, 0x3C)) {
    Serial.println("SSD1306 allocation failed");
    while (1);
  }
  display.clearDisplay();
  display.display();

  loadAccounts();
}

void loop() {
  handleSerial();
  handleButtons();
  updateDisplay();
  delay(100);  // Debounce/simple refresh
}

void handleSerial() {
  if (Serial.available()) {
    String cmd = Serial.readStringUntil('\n');
    cmd.trim();
    if (cmd.startsWith("add ")) {
      // Parse: add name hex:0x12,0x34,...
      int spaceIdx = cmd.indexOf(' ', 4);
      String name = cmd.substring(4, spaceIdx);
      String hexStr = cmd.substring(spaceIdx + 1 + 4);  // Skip "hex:"
      // Find or add account
      int idx = -1;
      for (int i = 0; i < numAccounts; i++) {
        if (String(accounts[i].name) == name) { idx = i; break; }
      }
      if (idx == -1 && numAccounts < MAX_ACCOUNTS) { idx = numAccounts++; }
      if (idx != -1) {
        name.toCharArray(accounts[idx].name, NAME_LEN);
        // Parse HEX to bytes
        int len = 0;
        while (hexStr.length() > 0 && len < SECRET_LEN) {
          int commaIdx = hexStr.indexOf(',');
          String byteStr = (commaIdx == -1) ? hexStr : hexStr.substring(0, commaIdx);
          accounts[idx].secret[len++] = strtol(byteStr.c_str(), NULL, 0);
          hexStr = (commaIdx == -1) ? "" : hexStr.substring(commaIdx + 1);
        }
        accounts[idx].secretLen = len;
        saveAccounts();
        Serial.println("Added/updated: " + name);
      } else {
        Serial.println("Max accounts reached");
      }
    } else if (cmd.startsWith("remove ")) {
      String name = cmd.substring(7);
      for (int i = 0; i < numAccounts; i++) {
        if (String(accounts[i].name) == name) {
          for (int j = i; j < numAccounts - 1; j++) {
            accounts[j] = accounts[j + 1];
          }
          numAccounts--;
          saveAccounts();
          Serial.println("Removed: " + name);
          return;
        }
      }
      Serial.println("Not found");
    } else if (cmd == "list") {
      for (int i = 0; i < numAccounts; i++) {
        Serial.print(accounts[i].name); Serial.println();
      }
    }
  }
}

void handleButtons() {
  if (digitalRead(btnUp) == LOW && state != CLOCK) { selected = max(0, selected - 1); }
  if (digitalRead(btnDown) == LOW && state != CLOCK) { selected = min(numAccounts - 1, selected + 1); }
  if (digitalRead(btnSelect) == LOW) {
    if (state == CLOCK) state = MENU;
    else if (state == MENU) state = OTP;
  }
  if (digitalRead(btnBack) == LOW) {
    if (state == OTP) state = MENU;
    else if (state == MENU) state = CLOCK;
  }
}

void updateDisplay() {
  display.clearDisplay();
  DateTime now = rtc.now();
  char buffer[20];  // Buffer for formatted strings
  if (state == CLOCK) {
    display.setTextSize(2);
    display.setTextColor(SSD1306_WHITE);
    display.setCursor(10, 10);
    sprintf(buffer, "%02d:%02d:%02d", now.hour(), now.minute(), now.second());
    display.print(buffer);
    display.setTextSize(1);
    display.setCursor(10, 40);
    sprintf(buffer, "%02d/%02d/%04d", now.month(), now.day(), now.year());
    display.print(buffer);
  } else if (state == MENU) {
    display.setTextSize(1);
    display.setCursor(0, 0);
    display.println("Accounts:");
    for (int i = 0; i < numAccounts; i++) {
      display.setCursor(0, 10 + i * 10);
      if (i == selected) display.print("> ");
      display.println(accounts[i].name);
    }
  } else if (state == OTP) {
    TOTP totp = TOTP(accounts[selected].secret, accounts[selected].secretLen);
    char* code = totp.getCode(now.unixtime());
    display.setTextSize(2);
    display.setCursor(0, 0);
    display.println(accounts[selected].name);
    display.setCursor(0, 20);
    display.println(code);
    // Progress bar for 30s
    int remaining = 30 - (now.second() % 30);
    display.drawRect(0, 50, 128, 10, SSD1306_WHITE);
    display.fillRect(0, 50, map(remaining, 0, 30, 0, 128), 10, SSD1306_WHITE);
  }
  display.display();
}