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

// Define the OLED reset pin
#define SCREEN_WIDTH 128
#define SCREEN_HEIGHT 64
#define OLED_RESET    -1 // Reset pin # (or -1 if sharing Arduino reset pin)
Adafruit_SSD1306 display(SCREEN_WIDTH, SCREEN_HEIGHT, &Wire, OLED_RESET);

// Define the relay and button pins
int DWELL = 20; // How long to activate the relay (in ms)
int RELAY = 2; // Relay control pin
int SEMI = A1; // Semi-auto switch pin
int FULLAUTO = A2; // Full-auto switch pin
int BUTTON = A3; // Fire button pin
int RESETBUTTON = 4; // Button to reset ammo left and ammo used
int INCREASEBUTTON = 12; // Button to increase ammo count
int DECREASEBUTTON = 11; // Button to decrease ammo count

// Initialize states
int SEMISTATE = LOW; // Start state of semi switch
int BUTTONSTATE = HIGH; // Start state of button
int LASTBUTTONSTATE = HIGH; // Previous state of button
unsigned long lastDebounceTime = 0; // The last time the output pin was toggled
unsigned long debounceDelay = 50; // The debounce time; increase if the output flickers

// Initialize the counters and timers
volatile int relayActivationCount = 0;
volatile bool relayActivated = false;
unsigned long totalActiveTime = 0; // Total time relay was active
unsigned long lastUpdateTime = 0; // Last time the display was updated
const unsigned long DISPLAY_UPDATE_INTERVAL = 100; // Display update interval in milliseconds

int ammoCount = 90; // Ammo count starting at 90 rounds
int initialAmmoCount = 90; // Store the initial ammo count
int ammoUsed = 0; // Total ammo used

// Timer variables for auto-reset
unsigned long zeroAmmoTime = 0; // Time when ammo count reached zero
bool waitingForReset = false; // Flag to indicate if waiting for reset

void setup() {
  // Initialize serial communication for debugging purposes
  Serial.begin(9600);

  // Initialize the OLED display
  if (!display.begin(SSD1306_SWITCHCAPVCC, 0x3C)) { // Address 0x3C for 128x64
    Serial.println(F("SSD1306 allocation failed"));
    for (;;);
  }
  delay(2000); // Pause for 2 seconds

  // Clear the buffer
  display.clearDisplay();

  // Set text size and color
  display.setTextSize(1.5); // Increase text size
  display.setTextColor(SSD1306_WHITE);
  
  // Initialize the button pins and relay
  pinMode(RELAY, OUTPUT); // Relay
  pinMode(SEMI, INPUT_PULLUP); // Semi Auto
  pinMode(FULLAUTO, INPUT_PULLUP); // Full Auto
  pinMode(BUTTON, INPUT_PULLUP); // Fire Button
  pinMode(RESETBUTTON, INPUT_PULLUP); // Reset Button
  pinMode(INCREASEBUTTON, INPUT_PULLUP); // Increase ammo button
  pinMode(DECREASEBUTTON, INPUT_PULLUP); // Decrease ammo button
  digitalWrite(RELAY, LOW); // Ensure relay is initially off
  
  lastUpdateTime = millis(); // Initialize last update time
}

void loop() {
  SEMISTATE = digitalRead(SEMI); // Read current status of semi switch
  int reading = digitalRead(BUTTON); // Read current status of button
  int resetReading = digitalRead(RESETBUTTON); // Read current status of reset button
  int increaseReading = digitalRead(INCREASEBUTTON); // Read current status of increase ammo button
  int decreaseReading = digitalRead(DECREASEBUTTON); // Read current status of decrease ammo button

  // Handle button debounce
  if (reading != LASTBUTTONSTATE) {
    lastDebounceTime = millis();
  }

  // Check if the button state has been stable for longer than the debounce delay
  if ((millis() - lastDebounceTime) > debounceDelay) {
    // If the button state has changed
    if (reading != BUTTONSTATE) {
      BUTTONSTATE = reading;
      // Detect button press (transition from HIGH to LOW)
      if (BUTTONSTATE == LOW && SEMISTATE == LOW) {
        semiAuto();
      }
    }
  }

  // Update the button state
  LASTBUTTONSTATE = reading;

  if (digitalRead(FULLAUTO) == LOW) { // Full auto switch is closed
    fullAuto(); // Only call fullAuto() if the full-auto switch is on
  } else {
    digitalWrite(RELAY, LOW); // Ensure relay is off if full-auto switch is off
  }

  // Check if the reset button is pressed
  if (resetReading == LOW) {
    ammoCount = initialAmmoCount; // Restore the ammo count to the stored initial value
    ammoUsed = 0; // Reset the ammo used count
    waitingForReset = false; // Stop waiting for reset
    lastUpdateTime = millis(); // Force display update immediately
    delay(200); // Debounce delay for the reset button
  }

  // Check if the increase or decrease ammo buttons are pressed
  if (increaseReading == LOW) {
    if (ammoCount < 90) {
      ammoCount++;
      initialAmmoCount = ammoCount; // Update the initial ammo count
    }
    delay(200); // Debounce delay for the increase button
  }

  if (decreaseReading == LOW) {
    if (ammoCount > 0) {
      ammoCount--;
      initialAmmoCount = ammoCount; // Update the initial ammo count
    }
    delay(200); // Debounce delay for the decrease button
  }

  // Handle automatic reset when ammo count reaches zero
  if (ammoCount == 0 && !waitingForReset) {
    zeroAmmoTime = millis(); // Record the time when ammo count reached zero
    waitingForReset = true; // Start waiting for reset
  }

  if (waitingForReset && millis() - zeroAmmoTime >= 5000) { // 5 seconds have passed
    ammoCount = initialAmmoCount; // Restore the ammo count to the stored initial value
    ammoUsed = 0; // Reset the ammo used count
    waitingForReset = false; // Stop waiting for reset
    lastUpdateTime = millis(); // Force display update immediately
  }

  // Update the display at regular intervals
  if (millis() - lastUpdateTime >= DISPLAY_UPDATE_INTERVAL) {
    display.clearDisplay();
    
    // Display the current relay activation count
    centerText("Count: " + String(relayActivationCount), 10);
    
    // Display the ammo count and ammo used
    centerText("Ammo Left: " + String(ammoCount), 30);
    centerText("Ammo Used: " + String(ammoUsed), 50);

    display.display();
    lastUpdateTime = millis(); // Update the last update time
  }
}

void semiAuto() {
  if (ammoCount > 0) {
    digitalWrite(RELAY, HIGH); // Activate relay
    delay(DWELL); // Wait for the specified dwell time
    digitalWrite(RELAY, LOW); // Deactivate relay

    // Increment the counter and decrement ammo count
    relayActivationCount++;
    ammoCount--;
    ammoUsed++;
    totalActiveTime += DWELL; // Accumulate total active time
    relayActivated = true;
  }
}

void fullAuto() {
  static unsigned long lastFireTime = 0;
  const unsigned long FIRE_INTERVAL = 100; // Time between each shot in milliseconds
  int buttonState = digitalRead(BUTTON);

  if (buttonState == LOW && ammoCount > 0) {
    unsigned long currentTime = millis();

    if (currentTime - lastFireTime >= FIRE_INTERVAL) {
      digitalWrite(RELAY, HIGH); // Activate relay
      delay(DWELL); // Wait for the specified dwell time
      digitalWrite(RELAY, LOW); // Deactivate relay

      // Increment the counter and decrement ammo count
      relayActivationCount++;
      ammoCount--;
      ammoUsed++;
      totalActiveTime += DWELL; // Accumulate total active time
      relayActivated = true;

      // Update last fire time
      lastFireTime = currentTime;
    }
  } else {
    digitalWrite(RELAY, LOW); // Ensure relay is off when button is not pressed or ammo is empty
  }
}

// Function to center text on the OLED display
void centerText(const String &text, int y) {
  int16_t textWidth;
  int16_t textHeight;
  int16_t textX1;
  int16_t textY1;
  
  display.getTextBounds(text.c_str(), 0, 0, &textX1, &textY1, &textWidth, &textHeight);
  int16_t textCursorX = (SCREEN_WIDTH - textWidth) / 2;
  
  display.setCursor(textCursorX, y);
  display.print(text);
}