#include <Wire.h>
#include <Adafruit_GFX.h>
#include <Adafruit_SSD1306.h>
#include <SPI.h>

// Pin Definitions
const int latchPin595 = 4;  // 74HC595 latch pin
const int clockPin595 = 3;  // 74HC595 clock pin
const int dataPin595 = 2;   // 74HC595 data pin

const int latchPin165 = 7;  // 74HC165 latch pin
const int clockPin165 = 6;  // 74HC165 clock pin
const int dataPin165 = 5;   // 74HC165 data pin

// Button Pins (Assuming buttons connected to 4x 74HC165)
const int buttonPins[32] = {0, 1, 2, 3, 4, 5, 6, 7,   // 1st 74HC165
                            8, 9, 10, 11, 12, 13, 14, 15, // 2nd 74HC165
                            16, 17, 18, 19, 20, 21, 22, 23, // 3rd 74HC165
                            24, 25, 26, 27, 28, 29, 30, 31}; // 4th 74HC165

// Mode Definitions
enum Mode {STANDBY, TEST, FIRING};
Mode currentMode = STANDBY;

// OLED display setup (using I2C)
#define SCREEN_WIDTH 128
#define SCREEN_HEIGHT 64
#define OLED_RESET    -1
Adafruit_SSD1306 display(SCREEN_WIDTH, SCREEN_HEIGHT, &Wire, OLED_RESET);

// LED Pins
const int greenLEDPin = 8;
const int amberLEDPin = 9;
const int redLEDPin = 10;

// Debounce parameters
const unsigned long debounceDelay = 50;
unsigned long lastDebounceTime[32] = {0};
bool lastButtonState[32] = {0};
bool buttonState[32] = {0};

// Button Mapping
const int rowButtons[8] = {0, 1, 2, 3, 4, 5, 6, 7}; // Buttons 1-8 for Rows 1-8
const int columnButtons[16] = {8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23}; // Buttons 9-24 for Columns A-P
const int nextModeButton = 24; // Button 25
const int previousModeButton = 25; // Button 26

void setup() {
  // Initialize serial communication
  Serial.begin(9600);
  
  // Initialize the OLED display
  if (!display.begin(SSD1306_I2C_ADDRESS, OLED_RESET)) {
    Serial.println(F("SSD1306 allocation failed"));
    for (;;);
  }
  display.clearDisplay();
  display.setTextSize(1);
  display.setTextColor(SSD1306_WHITE);

  // Initialize shift registers
  pinMode(latchPin595, OUTPUT);
  pinMode(clockPin595, OUTPUT);
  pinMode(dataPin595, OUTPUT);
  pinMode(latchPin165, OUTPUT);
  pinMode(clockPin165, OUTPUT);
  pinMode(dataPin165, INPUT);

  // Initialize LED pins as output
  pinMode(greenLEDPin, OUTPUT);
  pinMode(amberLEDPin, OUTPUT);
  pinMode(redLEDPin, OUTPUT);

  // Set all outputs to LOW initially
  for (int i = 0; i < 3; i++) {
    shiftOut(dataPin595, clockPin595, MSBFIRST, 0);
  }
  digitalWrite(latchPin595, HIGH);
  digitalWrite(latchPin595, LOW);
}

void loop() {
  // Read buttons
  readButtons();

  // Check mode change buttons
  if (buttonState[nextModeButton]) {
    changeMode(1);
  } else if (buttonState[previousModeButton]) {
    changeMode(-1);
  }

  // Handle current mode
  switch (currentMode) {
    case STANDBY:
      handleStandbyMode();
      break;
    case TEST:
      handleTestMode();
      break;
    case FIRING:
      handleFiringMode();
      break;
  }
}

void readButtons() {
  digitalWrite(latchPin165, LOW);
  delayMicroseconds(5);
  digitalWrite(latchPin165, HIGH);
  delayMicroseconds(5);

  for (int i = 0; i < 4; i++) {
    byte buttons = shiftIn(dataPin165, clockPin165, MSBFIRST);
    for (int j = 0; j < 8; j++) {
      int buttonIndex = i * 8 + j;
      bool reading = buttons & (1 << (7 - j));

      if (reading != lastButtonState[buttonIndex]) {
        lastDebounceTime[buttonIndex] = millis();
      }

      if ((millis() - lastDebounceTime[buttonIndex]) > debounceDelay) {
        if (reading != buttonState[buttonIndex]) {
          buttonState[buttonIndex] = reading;
          if (buttonState[buttonIndex] == HIGH) {
            // Button pressed
            buttonPressed(buttonIndex);
          }
        }
      }

      lastButtonState[buttonIndex] = reading;
    }
  }
}

void buttonPressed(int buttonIndex) {
  // Implement button press logic if needed
  // For now, we'll handle mode change buttons directly in loop()
}

void changeMode(int direction) {
  if (direction == 1) {
    if (currentMode == STANDBY) currentMode = TEST;
    else if (currentMode == TEST) currentMode = FIRING;
  } else if (direction == -1) {
    if (currentMode == FIRING) currentMode = TEST;
    else if (currentMode == TEST) currentMode = STANDBY;
  }

  // Clear button states to prevent multiple triggers
  for (int i = 0; i < 32; i++) {
    buttonState[i] = LOW;
  }

  // Update LEDs based on the new mode
  if (currentMode == STANDBY) {
    digitalWrite(greenLEDPin, HIGH);
    digitalWrite(amberLEDPin, LOW);
    digitalWrite(redLEDPin, LOW);
  } else if (currentMode == TEST) {
    digitalWrite(greenLEDPin, LOW);
    digitalWrite(amberLEDPin, HIGH);
    digitalWrite(redLEDPin, LOW);
  } else if (currentMode == FIRING) {
    digitalWrite(greenLEDPin, LOW);
    digitalWrite(amberLEDPin, LOW);
    // Red LED will blink in handleFiringMode()
  }

  // Update display based on new mode
  display.clearDisplay();
  if (currentMode == STANDBY) {
    display.setCursor(0, 0);
    display.print("Standby");
  } else if (currentMode == TEST) {
    display.setCursor(0, 0);
    display.print("Test Mode");
  } else if (currentMode == FIRING) {
    display.setCursor(0, 0);
    display.print("ARMED!");
  }
  display.display();
}

void handleStandbyMode() {
  // No current sent to outputs in standby mode
  // Just display "Standby" on OLED
  display.setCursor(0, 0);
  display.print("Standby");
  display.display();
}

void handleTestMode() {
  // Disable buttons 1 to 24
  for (int i = 0; i < 24; i++) {
    buttonState[i] = LOW;
  }

  // Send 5mA to all 128 outputs simultaneously
  // This is simulated here; you would adjust the actual circuit for low current
  for (int i = 0; i < 3; i++) {
    shiftOut(dataPin595, clockPin595, MSBFIRST, 0xFF);
  }
  digitalWrite(latchPin595, HIGH);
  digitalWrite(latchPin595, LOW);

  // Read status of all 128 outputs
  char outputStatus[129];
  outputStatus[128] = '\0';

  digitalWrite(latchPin165, LOW);
  delayMicroseconds(5);
  digitalWrite(latchPin165, HIGH);
  delayMicroseconds(5);

  for (int i = 0; i < 16; i++) {
    byte status = shiftIn(dataPin165, clockPin165, MSBFIRST);
    for (int j = 0; j < 8; j++) {
      outputStatus[i * 8 + j] = (status & (1 << (7 - j))) ? 'Y' : 'N';
    }
  }

  // Display status on OLED
  display.clearDisplay();
  display.setCursor(0, 0);
  for (int i = 0; i < 128; i++) {
    display.print(outputStatus[i]);
    if ((i + 1) % 16 == 0) {
      display.print("\n");
    }
  }
  display.display();
}

void handleFiringMode() {
  // Send 18V full power to selected outputs
  // Use buttons 1-8 to select rows, 9-24 to select columns
  int selectedRow = -1;
  int selectedColumn = -1;

  for (int i = 0; i < 8; i++) {
    if (buttonState[rowButtons[i]]) {
      selectedRow = i;
    }
  }

  for (int i = 0; i < 16; i++) {
    if (buttonState[columnButtons[i]]) {
      selectedColumn = i;
    }
  }

  if (selectedRow != -1 && selectedColumn != -1) {
    int rowBits = 1 << selectedRow;
    int columnBits = 1 << selectedColumn;

    shiftOut(dataPin595, clockPin595, MSBFIRST, columnBits);
    shiftOut(dataPin595, clockPin595, MSBFIRST, columnBits >> 8);
    shiftOut(dataPin595, clockPin595, MSBFIRST, rowBits);

    digitalWrite(latchPin595, HIGH);
    digitalWrite(latchPin595, LOW);

    // Update display
    display.clearDisplay();
    display.setCursor(0, 0);
    display.print("ARMED!");
    display.display();
  }

  // Blink the red LED
  static unsigned long lastBlinkTime = 0;
  static bool ledState = LOW;
  if (millis() - lastBlinkTime >= 500) { // 0.5 seconds
    ledState = !ledState;
    digitalWrite(redLEDPin, ledState);
    lastBlinkTime = millis();
  }
}