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

// ----------------- OLED Setup -----------------
// OLED display dimensions
#define SCREEN_WIDTH 128
#define SCREEN_HEIGHT 64

// Initialize OLED display using I2C
Adafruit_SSD1306 display(SCREEN_WIDTH, SCREEN_HEIGHT, &Wire, -1);

// State tracking variables
String currentLayer = "Base Layer";  // Default layer
bool numpadState = false;            // Default: Numpad OFF
int encoderPosition = 0;             // Rotary encoder position
bool isMuted = false;                // Mute status
String lastKeyPressed = "";          // Last key pressed for display

// ----------------- Encoder Setup -----------------
// Rotary encoder pins
#define CLK 8
#define DT 9
#define SW 10 // Rotary Encoder button
volatile int lastCLKState;
volatile bool encoderTurned = false;
volatile bool buttonPressed = false;

// Interrupt service routine for rotary encoder
void readEncoder() {
  int currentCLKState = digitalRead(CLK);
  if (currentCLKState != lastCLKState) {
    if (digitalRead(DT) != currentCLKState) {
      encoderPosition++;  // Clockwise rotation
    } else {
      encoderPosition--;  // Counterclockwise rotation
    }
    encoderTurned = true;  // Flag to update screen
  }
  lastCLKState = currentCLKState;
}

// Interrupt for the encoder button
void buttonPress() {
  buttonPressed = !buttonPressed;  // Toggle mute/unmute on button press
  isMuted = buttonPressed;
}

// ----------------- Keypad Setup -----------------
#define ROWS 4
#define COLS 4

// Keypad pin definitions
const int rowPins[ROWS] = {0, 1, 2, 3};
const int colPins[COLS] = {4, 5, 6, 7};

// Key map for the 4x4 keypad
const char keys[ROWS][COLS] = {
    {'1', '2', '3', 'A'}, // 'A' toggles layer
    {'4', '5', '6', 'B'}, // 'B' toggles numpad
    {'7', '8', '9', 'C'},
    {'*', '0', '#', 'D'}}; // '*' toggles numpad state

// Function to toggle layers
void toggleLayer() {
  currentLayer = (currentLayer == "Base Layer") ? "Function Layer" : "Base Layer";
  updateScreen();  // Ensure the OLED is updated
}

// Function to toggle numpad state
void toggleNumpad() {
  numpadState = !numpadState;
  updateScreen();  // Update OLED
}

// Function to scan the keypad
void scanKeypad() {
  for (int r = 0; r < ROWS; r++) {
    pinMode(rowPins[r], OUTPUT);
    digitalWrite(rowPins[r], LOW); // Set the row to LOW

    for (int c = 0; c < COLS; c++) {
      pinMode(colPins[c], INPUT_PULLUP);
      if (digitalRead(colPins[c]) == LOW) { // Key is pressed
        char key = keys[r][c];
        delay(50); // Debounce delay

        // Update the last key pressed
        lastKeyPressed = String(key);

        // Handle key actions
        if (key == 'A') toggleLayer();
        if (key == 'B') toggleNumpad();
        if (key == '*') toggleNumpad(); // Toggle numpad on/off with '*'

        // Debug output
        Serial.print("Key Pressed: ");
        Serial.println(key);

        // Update OLED after key press
        updateScreen();
      }
    }
    digitalWrite(rowPins[r], HIGH); // Reset row to HIGH
  }
}

// ----------------- OLED Update -----------------
void updateScreen() {
  // Debugging to check if updateScreen is being called
  Serial.println("Updating OLED...");

  display.clearDisplay();
  display.setTextSize(1);
  display.setTextColor(WHITE);
  display.setCursor(0, 0);

  // Display layer and numpad state
  display.print("Key Layer: ");
  display.println(currentLayer);
  display.print("Numpad: ");
  display.println(numpadState ? "ON" : "OFF");

  // Display volume and mute status
  display.print("Volume: ");
  if (isMuted) {
    display.println("Muted");
  } else {
    display.print(encoderPosition);
    display.println(" %");
  }

  // Display last key pressed
  if (lastKeyPressed != "") {
    display.print("Last Key: ");
    display.println(lastKeyPressed);
  }

  display.display(); // Make sure display updates
}

// ----------------- Main Program -----------------
void setup() {
  // Initialize Serial for debugging
  Serial.begin(9600);

  // Initialize OLED display
  if (!display.begin(SSD1306_SWITCHCAPVCC, 0x3C)) {
    Serial.println(F("SSD1306 allocation failed"));
    for (;;)
      ; // Halt execution
  }

  // Clear the buffer
  display.clearDisplay();
  display.display();

  // Initialize rotary encoder pins
  pinMode(CLK, INPUT_PULLUP);
  pinMode(DT, INPUT_PULLUP);
  pinMode(SW, INPUT_PULLUP); // Encoder button pin
  lastCLKState = digitalRead(CLK);
  attachInterrupt(digitalPinToInterrupt(CLK), readEncoder, CHANGE);
  attachInterrupt(digitalPinToInterrupt(SW), buttonPress, FALLING);

  // Initial screen update
  updateScreen();
}

void loop() {
  // Scan the keypad
  scanKeypad();

  // Update the screen if the encoder position changed
  if (encoderTurned) {
    updateScreen();
    encoderTurned = false;
  }

  delay(100); // Adjust for smooth operation
}