#include <Wire.h>
#include "config.h"
#include <Adafruit_GFX.h>
#include <Adafruit_SSD1306.h>
#include <ESP8266WiFi.h>

void playTune(int tuneNumber);
void beep(int buzz);

unsigned long middle_long_press_started = 0;
unsigned long right_long_press_started = 0; 
bool is_display_off = false;

Adafruit_SSD1306 display(SCREEN_WIDTH, SCREEN_HEIGHT, &Wire, OLED_RESET);

// Initialize display, buttons and buzzer
void initSystems() {

  pinMode(PIN_BTN_L, INPUT);
  pinMode(PIN_BTN_M, INPUT);
  pinMode(PIN_BTN_R, INPUT);
  pinMode(PIN_BUZZER, OUTPUT);
  pinMode(LED_BUILTIN, OUTPUT);
  digitalWrite(LED_BUILTIN, HIGH);
  // Initialize the SSD1306 display
  if (!display.begin(SSD1306_SWITCHCAPVCC, 0x3C)) { // Adjust I2C address if needed
    for(;;); // Don't proceed, loop forever
  }

  // Rotate the display 180 degrees
  display.setRotation(2);

  // Clear the display buffer
  display.clearDisplay();
  display.setTextSize(1);      
  display.setTextColor(WHITE);
  display.setCursor(0, 0);
  display.display();
}

// Connect to WiFi
void connectToWifi() {

  display.clearDisplay();
  display.setCursor(0, 0);
  display.println("Connecting to WiFi");
  display.display();

  WiFi.begin(WIFI_SSID, WIFI_PASSWORD);

  int attemptCount = 0;
  while (WiFi.status() != WL_CONNECTED && attemptCount < 20) { // Timeout after 20 attempts
    delay(500);
    display.print(".");
    display.display();
    attemptCount++;
  }

  // Check if connected
  if (WiFi.status() == WL_CONNECTED) {
    Serial.println("\nConnected to WiFi!");
    Serial.print("IP Address: ");
    Serial.println(WiFi.localIP());

    // Display connection success and IP address
    display.clearDisplay();
    display.setCursor(0, 0);
    display.print("WiFi Connected!");
    display.setCursor(0, 10);
    display.print("IP: ");
    display.print(WiFi.localIP());
    display.display();
    beep(800);
  } else {
    Serial.println("\nFailed to connect to WiFi.");
    display.clearDisplay();
    display.setCursor(0, 0);
    display.print("WiFi Connection");
    display.setCursor(0, 10);
    display.print("Failed");
    display.display();
  }
  delay(1000);
}

void commonButtonHandler() {
    unsigned long currentMillis = millis();
    static unsigned long leftPressStart = 0;
    static unsigned long middlePressStart = 0;
    static unsigned long rightPressStart = 0;

    static bool leftBeeped = false;
    static bool middleBeeped = false;
    static bool rightBeeped = false;

    bool leftPressed = (digitalRead(PIN_BTN_L) == HIGH);
    bool middlePressed = (digitalRead(PIN_BTN_M) == HIGH);
    bool rightPressed = (digitalRead(PIN_BTN_R) == HIGH);

    // Short press detection
    if (leftPressed) {
        if ((currentMillis - leftPressStart > 50)) { // Debounce delay
            if (!leftBeeped) {
                beep(1000); // Play beep sound
                leftBeeped = true; // Set beeped state
                // Handle left short press action here
            }
        }
    } else {
        leftPressStart = currentMillis; // Reset the timer if button is not pressed
        leftBeeped = false; // Reset beeped state
    }

    if (middlePressed) {
        if ((currentMillis - middlePressStart > 50)) { // Debounce delay
            if (!middleBeeped) {
                beep(1000); // Play beep sound
                middleBeeped = true; // Set beeped state
                // Handle middle short press action here
            }
        }
    } else {
        middlePressStart = currentMillis; // Reset the timer if button is not pressed
        middleBeeped = false; // Reset beeped state
    }

    if (rightPressed) {
        if ((currentMillis - rightPressStart > 50)) { // Debounce delay
            if (!rightBeeped) {
                beep(1000); // Play beep sound
                rightBeeped = true; // Set beeped state
                if (is_display_off) {
                    display.ssd1306_command(SSD1306_DISPLAYON); // Turn on display
                    is_display_off = false;
                    beep(1300);
                }            
            }
        }
    } else {
        rightPressStart = currentMillis; // Reset the timer if button is not pressed
        rightBeeped = false; // Reset beeped state
    }

    // Long press detection
    if (leftPressed && (currentMillis - leftPressStart > 2000)) {
        if (!leftBeeped) {
            beep(1000); // Play beep sound
            leftBeeped = true; // Set beeped state
            // Handle left long press action here
        }
    }

    if (middlePressed && (currentMillis - middlePressStart > 2000)) {
        if (!middleBeeped) {
            beep(1000); // Play beep sound
            middleBeeped = true; // Set beeped state
            // Handle middle long press action here
        }
    }

    if (rightPressed && (currentMillis - rightPressStart > 2000)) {
            if (!is_display_off) {
                display.ssd1306_command(SSD1306_DISPLAYOFF); // Turn off display
                is_display_off = true;
                playTune(3); // Play beep sound
            }
    }

    // Combination of Left and Middle long press
    if (leftPressed && middlePressed && 
        (currentMillis - leftPressStart > 2000) && (currentMillis - middlePressStart > 2000)) {
            beep(1000); // Play beep sound
            ESP.restart();
    }
}