#include <FastLED.h>
#include <Wire.h>
#include <Adafruit_GFX.h>
#include <Adafruit_SSD1306.h>
#include <Fonts/FreeSansBold12pt7b.h>

// --- Definitions ---
#define LED_PIN          10
#define BUZZER_PIN       9    // Buzzer back to Pin 9
#define BATTERY_PIN      A6
#define NUM_LEDS         36
#define BRIGHTNESS       255
#define BATT_UPDATE_MS   2000 

Adafruit_SSD1306 display(128, 64, &Wire, -1);

// Global Variables
CRGB leds[NUM_LEDS];
CRGB currentColor = CRGB::Red;
const char* currentColorName = "RED";
unsigned long lastBattUpdate = 0;

const int rowPins[7] = {0, 2, 3, 4, 5, 6, 7};
const int colPins[6] = {A3, A2, A1, A0, 13, 12};

CRGB row7Colors[6] = {CRGB::Red, 0xFFA500, CRGB::Yellow, CRGB::Green, CRGB::Blue, CRGB::Magenta};
const char* colorNames[6] = {"RED", "ORANGE", "YELLOW", "GREEN", "BLUE", "MAGENTA"};
const int colorTones[6] = {440, 494, 523, 587, 659, 698};

void drawBattery(int percent) {
    // Clear top-right corner
    display.fillRect(80, 0, 48, 15, SSD1306_BLACK);
    
    // Battery Box
    display.drawRect(105, 2, 20, 10, SSD1306_WHITE);
    display.fillRect(125, 5, 2, 4, SSD1306_WHITE);
    
    // Level fill
    int fill = (percent * 16) / 100;
    display.fillRect(107, 4, fill, 6, SSD1306_WHITE);
    
    // Percentage text
    display.setFont(); 
    display.setCursor(82, 4);
    display.print(percent); display.print(F("%"));
}

void updateMainDisplay() {
    display.clearDisplay();
    
    // Draw Current Color Name
    display.setFont(&FreeSansBold12pt7b);
    display.setTextColor(SSD1306_WHITE);
    int16_t x1, y1;
    uint16_t w, h;
    display.getTextBounds(currentColorName, 0, 0, &x1, &y1, &w, &h);
    display.setCursor((128 - w) / 2, (64 + h) / 2 + 5);
    display.print(currentColorName);
    
    // Battery Math (1.1V internal ref is already set in setup)
    int raw = analogRead(BATTERY_PIN);
    float v = raw * (1.1 / 1023.0) * (122.0 / 22.0);
    int percent = map(constrain(v * 10, 33, 45), 33, 45, 0, 100);
    drawBattery(percent);
    
    display.display();
}

void setup() {
    // 1. Set Analog Reference FIRST
    analogReference(INTERNAL);
    
    // 2. Initialize Buzzer
    pinMode(BUZZER_PIN, OUTPUT);
    digitalWrite(BUZZER_PIN, LOW);

    // 3. Initialize OLED
    if(display.begin(SSD1306_SWITCHCAPVCC, 0x3C)) {
        display.clearDisplay();
        updateMainDisplay();
    }

    // 4. Initialize LEDs
    FastLED.addLeds<WS2812B, LED_PIN, GRB>(leds, NUM_LEDS);
    FastLED.setBrightness(BRIGHTNESS);
    
    // 5. Initialize Matrix Pins
    for (int i = 0; i < 7; i++) pinMode(rowPins[i], INPUT);
    for (int i = 0; i < 6; i++) pinMode(colPins[i], INPUT_PULLUP);
    
    // Startup Tone to confirm Pin 9 is alive
    tone(BUZZER_PIN, 880, 150);
}

void loop() {
    // Non-blocking battery update
    if (millis() - lastBattUpdate >= BATT_UPDATE_MS) {
        updateMainDisplay();
        lastBattUpdate = millis();
    }

    bool ledChanged = false;
    for (int r = 0; r < 7; r++) {
        pinMode(rowPins[r], OUTPUT);
        digitalWrite(rowPins[r], LOW);
        delayMicroseconds(40);

        for (int c = 0; c < 6; c++) {
            if (digitalRead(colPins[c]) == LOW) {
                if (r < 6) { 
                    int ledIndex = (r * 6) + c;
                    if (ledIndex < NUM_LEDS) {
                        leds[ledIndex] = currentColor;
                        ledChanged = true;
                        tone(BUZZER_PIN, 1800, 10); // High pitch short click
                    }
                } 
                else {
                    if (strcmp(currentColorName, colorNames[c]) != 0) {
                        currentColor = row7Colors[c];
                        currentColorName = colorNames[c];
                        tone(BUZZER_PIN, colorTones[c], 150); // Melodic beep
                        updateMainDisplay();
                    }
                }
                // Debounce / wait for release
                delay(150); 
            }
        }
        pinMode(rowPins[r], INPUT);
    }

    if (ledChanged) {
        FastLED.show();
    }
}