/*
  Six-party demonstration voting machine
  Board: Arduino Uno R3
  Current test mode: Serial Monitor acts as the complete display.
  Optional later display: 0.96-inch 128x64 I2C OLED (SSD1306), 0x3C.

  Buttons:
    Party 1..5: D2..D6
    Party 6: A0
    Green/confirm: A1
    Red/cancel: A2

  Every button connects between its Arduino pin and GND. The sketch uses
  INPUT_PULLUP, so no external button resistors are required.

  Hold GREEN + RED together for 2 seconds on the ready screen to enter the
  admin passcode. Party buttons enter digits 1..6, GREEN submits, and RED
  clears/cancels. Default passcode: 2614. On the results screen, hold GREEN
  for 3 seconds to reset every stored vote to zero.

  IMPORTANT: This is an educational demonstration. It is not suitable for a
  real public election; it has no voter authentication, certified audit trail,
  cryptographic security, or tamper resistance.
*/

// OLED is now enabled. Set this back to 0 only for Serial-Monitor-only testing.
#define USE_OLED 1

#if USE_OLED
  #include <Wire.h>
  #include <Adafruit_GFX.h>
  #include <Adafruit_SSD1306.h>
#endif
#include <EEPROM.h>
#include <string.h>

#if USE_OLED
constexpr uint8_t SCREEN_WIDTH = 128;
constexpr uint8_t SCREEN_HEIGHT = 64;
constexpr uint8_t OLED_ADDRESS = 0x3C;
constexpr int8_t OLED_RESET = -1;

Adafruit_SSD1306 display(SCREEN_WIDTH, SCREEN_HEIGHT, &Wire, OLED_RESET);
bool oledReady = false;
#endif

constexpr uint8_t PARTY_COUNT = 6;
const uint8_t PARTY_PINS[PARTY_COUNT] = {2, 3, 4, 5, 6, A0};
constexpr uint8_t GREEN_PIN = A1;
constexpr uint8_t RED_PIN = A2;

// Change this to any four digits from 1 through 6.
const char ADMIN_PASSCODE[] = "2614";

constexpr uint16_t EEPROM_MAGIC = 0x564D;
constexpr int EEPROM_MAGIC_ADDRESS = 0;
constexpr int EEPROM_VOTES_ADDRESS = 2;

constexpr uint16_t DEBOUNCE_MS = 30;
constexpr uint16_t MESSAGE_MS = 1600;
constexpr uint16_t CREDIT_MS = 3000;
constexpr uint16_t ADMIN_HOLD_MS = 2000;
constexpr uint16_t RESET_HOLD_MS = 3000;
#if USE_OLED
constexpr uint16_t OLED_KEEP_AWAKE_MS = 5000;
#endif

uint32_t votes[PARTY_COUNT] = {};

enum class Mode : uint8_t {
  READY,
  CONFIRM,
  MESSAGE,
  CREDITS,
  PASSCODE,
  RESULTS
};

Mode mode = Mode::READY;
uint8_t selectedParty = 0;
uint32_t messageStarted = 0;
uint32_t creditStarted = 0;
bool showCreditsAfterMessage = false;
uint32_t adminHoldStarted = 0;
bool adminHoldTriggered = false;
uint32_t resetHoldStarted = 0;
bool resetHoldTriggered = false;
#if USE_OLED
uint32_t lastOledKeepAwake = 0;
#endif

char enteredCode[5] = {};
uint8_t enteredDigits = 0;

struct ButtonState {
  bool stable = HIGH;
  bool previousReading = HIGH;
  uint32_t changedAt = 0;
  bool pressEvent = false;
};

ButtonState partyButtons[PARTY_COUNT];
ButtonState greenButton;
ButtonState redButton;

const __FlashStringHelper *partyName(uint8_t partyIndex) {
  switch (partyIndex) {
    case 0: return F("NATKHAT PARTY");
    case 1: return F("AURORAEON JANTA PARTY");
    case 2: return F("MACHHAR PARTY");
    case 3: return F("AVENGERS PARTY");
    case 4: return F("MADHUMAKHI PARTY");
    case 5: return F("BUCCHU PARTY");
    default: return F("UNKNOWN PARTY");
  }
}

void updateButton(uint8_t pin, ButtonState &button) {
  button.pressEvent = false;
  const bool reading = digitalRead(pin);
  const uint32_t now = millis();

  if (reading != button.previousReading) {
    button.previousReading = reading;
    button.changedAt = now;
  }

  if ((now - button.changedAt) >= DEBOUNCE_MS && reading != button.stable) {
    button.stable = reading;
    if (button.stable == LOW) button.pressEvent = true;
  }
}

void updateAllButtons() {
  for (uint8_t i = 0; i < PARTY_COUNT; ++i) {
    updateButton(PARTY_PINS[i], partyButtons[i]);
    if (partyButtons[i].pressEvent) {
      Serial.print(F("BUTTON: PARTY "));
      Serial.print(i + 1);
      Serial.print(F(" - "));
      Serial.println(partyName(i));
    }
  }
  updateButton(GREEN_PIN, greenButton);
  updateButton(RED_PIN, redButton);

  if (greenButton.pressEvent) Serial.println(F("BUTTON: GREEN"));
  if (redButton.pressEvent) Serial.println(F("BUTTON: RED"));
}

void saveVotes() {
  for (uint8_t i = 0; i < PARTY_COUNT; ++i) {
    EEPROM.put(EEPROM_VOTES_ADDRESS + i * (int)sizeof(uint32_t), votes[i]);
  }
}

void loadVotes() {
  uint16_t magic = 0;
  EEPROM.get(EEPROM_MAGIC_ADDRESS, magic);

  if (magic != EEPROM_MAGIC) {
    EEPROM.put(EEPROM_MAGIC_ADDRESS, EEPROM_MAGIC);
    saveVotes();
    return;
  }

  for (uint8_t i = 0; i < PARTY_COUNT; ++i) {
    EEPROM.get(EEPROM_VOTES_ADDRESS + i * (int)sizeof(uint32_t), votes[i]);
  }
}

void beginScreen(const __FlashStringHelper *title) {
  Serial.println();
  Serial.println(F("================================"));
  Serial.println(title);
  Serial.println(F("================================"));

#if USE_OLED
  if (oledReady) {
    display.clearDisplay();
    display.setTextColor(SSD1306_WHITE);
    display.setTextWrap(false);
    display.setTextSize(1);
    display.setCursor(0, 0);
    display.println(title);
    display.drawFastHLine(0, 10, SCREEN_WIDTH, SSD1306_WHITE);
  }
#endif
}

#if USE_OLED
void keepOledAwake() {
  if (!oledReady || millis() - lastOledKeepAwake < OLED_KEEP_AWAKE_MS) return;

  // Some generic SSD1306 modules unexpectedly enter display-off mode.
  // Reassert normal display mode and resend the current screen buffer.
  display.ssd1306_command(SSD1306_DISPLAYALLON_RESUME);
  display.ssd1306_command(SSD1306_NORMALDISPLAY);
  display.ssd1306_command(SSD1306_DISPLAYON);
  display.display();
  lastOledKeepAwake = millis();
}
#endif

void drawReady() {
  mode = Mode::READY;
  adminHoldStarted = 0;
  adminHoldTriggered = false;

  beginScreen(F("VOTING READY"));
  Serial.println(F("Press a party button:"));
  for (uint8_t i = 0; i < PARTY_COUNT; ++i) {
    Serial.print(i + 1);
    Serial.print(F(" - "));
    Serial.println(partyName(i));
  }
  Serial.println(F("After selection: GREEN confirms, RED cancels."));
  Serial.println(F("Admin results: hold GREEN + RED for 2 seconds."));

#if USE_OLED
  if (oledReady) {
    display.setTextSize(1);
    display.setCursor(0, 16);
    display.println(F("Press party key 1-6"));
    display.setCursor(0, 29);
    display.println(F("GREEN = confirm"));
    display.setCursor(0, 40);
    display.println(F("RED   = cancel"));
    display.setCursor(0, 54);
    display.print(F("Admin: hold G+R"));
    display.display();
  }
#endif
}

void drawConfirmation(uint8_t partyIndex) {
  mode = Mode::CONFIRM;
  beginScreen(F("CONFIRM VOTE"));
  Serial.print(F("Selected: Party "));
  Serial.print(partyIndex + 1);
  Serial.print(F(" - "));
  Serial.println(partyName(partyIndex));
  Serial.println(F("Press GREEN to confirm this vote."));
  Serial.println(F("Press RED to cancel this vote."));

#if USE_OLED
  if (oledReady) {
    display.setTextSize(1);
    display.setCursor(0, 17);
    display.print(F("P"));
    display.print(partyIndex + 1);
    display.print(F(": "));
    display.setTextWrap(true);
    display.println(partyName(partyIndex));
    display.setTextWrap(false);
    display.setCursor(0, 52);
    display.print(F("GREEN yes  RED no"));
    display.display();
  }
#endif
}

void showMessage(const __FlashStringHelper *title,
                 const __FlashStringHelper *message) {
  mode = Mode::MESSAGE;
  beginScreen(title);
  Serial.println(message);

#if USE_OLED
  if (oledReady) {
    display.setTextSize(1);
    display.setCursor(0, 24);
    display.println(message);
    display.display();
  }
#endif
  messageStarted = millis();
}

void drawCredits() {
  mode = Mode::CREDITS;
  Serial.println();
  Serial.println(F("made by Chitraksh AND Ansh did not make it"));

#if USE_OLED
  if (oledReady) {
    display.clearDisplay();
    display.setTextColor(SSD1306_WHITE);
    display.setTextWrap(false);
    display.setTextSize(1);
    display.setCursor(0, 16);
    display.println(F("made by Chitraksh"));
    display.println(F("AND Ansh did not"));
    display.println(F("make it"));
    display.display();
  }
#endif

  creditStarted = millis();
}

void redrawPasscodeField() {
  Serial.print(F("CODE: "));
  for (uint8_t i = 0; i < 4; ++i) {
    Serial.print(i < enteredDigits ? '*' : '_');
  }
  Serial.println();

#if USE_OLED
  if (oledReady) {
    display.fillRect(0, 27, SCREEN_WIDTH, 18, SSD1306_BLACK);
    display.setTextColor(SSD1306_WHITE);
    display.setTextSize(2);
    display.setCursor(36, 28);
    for (uint8_t i = 0; i < 4; ++i) {
      display.print(i < enteredDigits ? '*' : '_');
    }
    display.display();
  }
#endif
}

void drawPasscode() {
  mode = Mode::PASSCODE;
  enteredDigits = 0;
  enteredCode[0] = '\0';

  beginScreen(F("ADMIN PASSCODE"));
  Serial.println(F("Use party buttons as digits 1 through 6."));
  Serial.println(F("Press GREEN to submit."));
  Serial.println(F("Press RED to delete; RED on an empty code exits."));

#if USE_OLED
  if (oledReady) {
    display.setTextSize(1);
    display.setCursor(0, 14);
    display.println(F("Use party keys 1-6"));
    display.setCursor(0, 49);
    display.println(F("GREEN OK  RED cancel"));
  }
#endif
  redrawPasscodeField();
}

void drawResults() {
  mode = Mode::RESULTS;
  resetHoldStarted = 0;
  resetHoldTriggered = false;
  beginScreen(F("RESULTS"));
  for (uint8_t i = 0; i < PARTY_COUNT; ++i) {
    Serial.print(i + 1);
    Serial.print(F(" - "));
    Serial.print(partyName(i));
    Serial.print(F(": "));
    Serial.println(votes[i]);
  }

#if USE_OLED
  if (oledReady) {
    display.setTextSize(1);
    for (uint8_t i = 0; i < PARTY_COUNT; ++i) {
      const uint8_t column = i / 3;
      const uint8_t row = i % 3;
      display.setCursor(column * 66, 14 + row * 10);
      display.print('P');
      display.print(i + 1);
      display.print(':');
      display.print(votes[i]);
    }
  }
#endif

  uint32_t highest = 0;
  for (uint8_t i = 0; i < PARTY_COUNT; ++i) {
    if (votes[i] > highest) highest = votes[i];
  }

  if (highest == 0) {
    Serial.println(F("No votes yet."));
#if USE_OLED
    if (oledReady) {
      display.setCursor(0, 47);
      display.print(F("No votes yet"));
    }
#endif
  } else {
    Serial.println(F("Leading party/parties:"));
#if USE_OLED
    if (oledReady) {
      display.setCursor(0, 47);
      display.print(F("Leader:"));
    }
#endif
    for (uint8_t i = 0; i < PARTY_COUNT; ++i) {
      if (votes[i] == highest) {
        Serial.print(F("- "));
        Serial.println(partyName(i));
#if USE_OLED
        if (oledReady) {
          display.print(' ');
          display.print('P');
          display.print(i + 1);
        }
#endif
      }
    }
  }

  Serial.println(F("Press RED to return to voting."));
  Serial.println(F("Hold GREEN for 3 seconds to RESET ALL VOTES."));
#if USE_OLED
  if (oledReady) {
    display.setCursor(0, 57);
    display.print(F("RED back; hold G reset"));
    display.display();
  }
#endif
}

void resetAllVotes() {
  for (uint8_t i = 0; i < PARTY_COUNT; ++i) {
    votes[i] = 0;
  }
  saveVotes();
  showMessage(F("VOTES RESET"), F("All vote totals are now zero"));
}

void handleResults() {
  if (greenButton.stable == LOW) {
    if (resetHoldStarted == 0) resetHoldStarted = millis();
    if (!resetHoldTriggered && millis() - resetHoldStarted >= RESET_HOLD_MS) {
      resetHoldTriggered = true;
      resetAllVotes();
    }
    return;
  }

  resetHoldStarted = 0;
  resetHoldTriggered = false;

  if (redButton.pressEvent) drawReady();
}

void handleReady() {
  const bool bothHeld = greenButton.stable == LOW && redButton.stable == LOW;

  if (bothHeld) {
    if (adminHoldStarted == 0) adminHoldStarted = millis();
    if (!adminHoldTriggered && millis() - adminHoldStarted >= ADMIN_HOLD_MS) {
      adminHoldTriggered = true;
      drawPasscode();
    }
    return;
  }

  adminHoldStarted = 0;
  adminHoldTriggered = false;

  for (uint8_t i = 0; i < PARTY_COUNT; ++i) {
    if (partyButtons[i].pressEvent) {
      selectedParty = i;
      drawConfirmation(i);
      return;
    }
  }
}

void handleConfirmation() {
  if (greenButton.pressEvent) {
    ++votes[selectedParty];
    saveVotes();
    showCreditsAfterMessage = true;
    showMessage(F("VOTE SAVED"), F("Thank you"));
  } else if (redButton.pressEvent) {
    showCreditsAfterMessage = false;
    showMessage(F("VOTE CANCELLED"), F("No vote was added"));
  }
}

void handlePasscode() {
  for (uint8_t i = 0; i < PARTY_COUNT; ++i) {
    if (partyButtons[i].pressEvent && enteredDigits < 4) {
      enteredCode[enteredDigits++] = '1' + i;
      enteredCode[enteredDigits] = '\0';
      redrawPasscodeField();
      return;
    }
  }

  if (redButton.pressEvent) {
    if (enteredDigits > 0) {
      --enteredDigits;
      enteredCode[enteredDigits] = '\0';
      redrawPasscodeField();
    } else {
      drawReady();
    }
  } else if (greenButton.pressEvent) {
    if (strcmp(enteredCode, ADMIN_PASSCODE) == 0) {
      drawResults();
    } else {
      showMessage(F("WRONG CODE"), F("Access denied"));
    }
  }
}

void setup() {
  for (uint8_t i = 0; i < PARTY_COUNT; ++i) {
    pinMode(PARTY_PINS[i], INPUT_PULLUP);
  }
  pinMode(GREEN_PIN, INPUT_PULLUP);
  pinMode(RED_PIN, INPUT_PULLUP);

  Serial.begin(9600);
  Serial.println();
  Serial.println(F("SIX-PARTY VOTING MACHINE"));
  Serial.println(F("Serial Monitor display ready at 9600 baud"));

#if USE_OLED
  Wire.begin();
  oledReady = display.begin(SSD1306_SWITCHCAPVCC, OLED_ADDRESS);
  if (!oledReady) Serial.println(F("OLED not detected; continuing with Serial Monitor."));
#else
  Serial.println(F("OLED disabled: no display or OLED libraries required."));
#endif

  loadVotes();
  drawReady();
}

void loop() {
  updateAllButtons();

#if USE_OLED
  keepOledAwake();
#endif

  if (mode == Mode::MESSAGE) {
    if (millis() - messageStarted >= MESSAGE_MS) {
      if (showCreditsAfterMessage) {
        showCreditsAfterMessage = false;
        drawCredits();
      } else {
        drawReady();
      }
    }
    return;
  }

  if (mode == Mode::CREDITS) {
    if (millis() - creditStarted >= CREDIT_MS) drawReady();
    return;
  }

  switch (mode) {
    case Mode::READY:
      handleReady();
      break;
    case Mode::CONFIRM:
      handleConfirmation();
      break;
    case Mode::PASSCODE:
      handlePasscode();
      break;
    case Mode::RESULTS:
      handleResults();
      break;
    case Mode::MESSAGE:
    case Mode::CREDITS:
      break;
  }
}