#include <FastLED.h>
#include <EEPROM.h>

#define NUM_LEDS 21
#define LED_PIN 3

CRGB leds[NUM_LEDS];

// Button pins
const int buttonPins[10] = {8, 9, 10, 11, 12, 13, A0, A1, A2, A3};
const int colorCyclePin = A4;
const int brightnessPotPin = A5;

// LED zones per button
const int button0Leds[] = {4, 3, 2, 1, 0};
const int button1Leds[] = {16, 17, 18, 19, 20};
const int button2Leds[] = {5, 6, 7};
const int button3Leds[] = {8, 9};
const int button4Leds[] = {10};
const int button5Leds[] = {11};
const int button6Leds[] = {12};
const int button7Leds[] = {13};
const int button8Leds[] = {14};
const int button9Leds[] = {15};

const int* buttonLeds[] = {
  button0Leds, button1Leds, button2Leds, button3Leds, button4Leds,
  button5Leds, button6Leds, button7Leds, button8Leds, button9Leds
};

const int buttonLedCounts[] = {5, 5, 3, 2, 1, 1, 1, 1, 1, 1};

// Color palette
const CRGB colorPalette[] = {
  CRGB::Black, CRGB::White, CRGB(173, 216, 230), CRGB::Yellow,
  CRGB(255, 165, 0), CRGB(255, 80, 0), CRGB::Green, CRGB(0, 255, 127),
  CRGB::Cyan, CRGB(255, 105, 180), CRGB::Purple, CRGB::Blue, CRGB::Red
};
const int numColors = sizeof(colorPalette) / sizeof(CRGB);

// EEPROM storage
int colorIndex[10];
int brightness = 128;
int kittColorIndex = 1;

// KITT animation control
unsigned long lastActivityTime = 0;
bool kitForward = true;
int kitPos = 0;
unsigned long lastKitUpdate = 0;
const unsigned long kitInterval = 25;
const unsigned long idleTimeout = 1000;

// Sweep logic
const unsigned long sweepDelay = 50;
unsigned long lastSweepUpdate[10] = {0};
int activeLedCount[10] = {0};

// Flash burst logic
const int burstPulses = 3;
const unsigned long flashOnTime = 100;
const unsigned long flashOffTime = 100;
int currentPulse = 0;
bool flashState = false;
unsigned long lastFlashTime = 0;
bool burstActive = false;

void setup() {
  FastLED.addLeds<WS2812B, LED_PIN, GRB>(leds, NUM_LEDS);
  Serial.begin(9600);

  for (int i = 0; i < 10; i++) {
    pinMode(buttonPins[i], INPUT_PULLUP);
    colorIndex[i] = EEPROM.read(i);
    if (colorIndex[i] >= numColors) colorIndex[i] = 1;
  }

  pinMode(colorCyclePin, INPUT_PULLUP);
  pinMode(brightnessPotPin, INPUT);

  brightness = EEPROM.read(20);
  if (brightness > 255) brightness = 128;

  kittColorIndex = EEPROM.read(21);
  if (kittColorIndex >= numColors) kittColorIndex = 1;

  FastLED.setBrightness(brightness);
  FastLED.clear();
  FastLED.show();
}

void loop() {
  bool anyButtonPressed = false;
  bool colorCyclePressed = digitalRead(colorCyclePin) == LOW;

  brightness = analogRead(brightnessPotPin) / 4;
  FastLED.setBrightness(brightness);
  EEPROM.update(20, brightness);

  // KITT color change via A4 alone
  if (colorCyclePressed) {
    bool otherPressed = false;
    for (int i = 0; i < 10; i++) {
      if (digitalRead(buttonPins[i]) == LOW) {
        otherPressed = true;
        break;
      }
    }
    if (!otherPressed) {
      kittColorIndex = (kittColorIndex + 1) % numColors;
      EEPROM.update(21, kittColorIndex);
      delay(300);
    }
  }

  bool button0 = digitalRead(buttonPins[0]) == LOW;
  bool button1 = digitalRead(buttonPins[1]) == LOW;
  bool burstTrigger = button0 && button1;

  if (burstTrigger) {
    anyButtonPressed = true;
    unsigned long now = millis();

    if (!burstActive) {
      currentPulse = 0;
      flashState = false;
      lastFlashTime = now;
      burstActive = true;
    }

    if (currentPulse < burstPulses) {
      if ((flashState && now - lastFlashTime >= flashOnTime) ||
          (!flashState && now - lastFlashTime >= flashOffTime)) {
        flashState = !flashState;
        lastFlashTime = now;
        if (!flashState) currentPulse++;
      }

      CRGB c0 = colorPalette[colorIndex[0]];
      CRGB c1 = colorPalette[colorIndex[1]];
      for (int j = 0; j < buttonLedCounts[0]; j++) {
        leds[buttonLeds[0][j]] = flashState ? c0 : CRGB::Black;
      }
      for (int j = 0; j < buttonLedCounts[1]; j++) {
        leds[buttonLeds[1][j]] = flashState ? c1 : CRGB::Black;
      }

    } else {
      for (int j = 0; j < buttonLedCounts[0]; j++) {
        leds[buttonLeds[0][j]] = colorPalette[colorIndex[0]];
      }
      for (int j = 0; j < buttonLedCounts[1]; j++) {
        leds[buttonLeds[1][j]] = colorPalette[colorIndex[1]];
      }
    }
  } else {
    burstActive = false;

    // Standard button logic
    for (int i = 0; i < 10; i++) {
      bool pressed = digitalRead(buttonPins[i]) == LOW;

      if (pressed) {
        anyButtonPressed = true;
        lastActivityTime = millis();

        if (colorCyclePressed) {
          colorIndex[i] = (colorIndex[i] + 1) % numColors;
          EEPROM.update(i, colorIndex[i]);
          delay(300);
        }

        if ((i == 0 || i == 1) && burstActive) continue;

        if (buttonLedCounts[i] > 1 && !colorCyclePressed) {
          if (millis() - lastSweepUpdate[i] >= sweepDelay && activeLedCount[i] < buttonLedCounts[i]) {
            activeLedCount[i]++;
            lastSweepUpdate[i] = millis();
          }

          for (int j = 0; j < buttonLedCounts[i]; j++) {
            leds[buttonLeds[i][j]] = (j < activeLedCount[i]) ? colorPalette[colorIndex[i]] : CRGB::Black;
          }

        } else {
          for (int j = 0; j < buttonLedCounts[i]; j++) {
            leds[buttonLeds[i][j]] = colorPalette[colorIndex[i]];
          }
          activeLedCount[i] = buttonLedCounts[i];
        }

      } else {
        if ((i == 0 || i == 1) && burstActive) continue;

        for (int j = 0; j < buttonLedCounts[i]; j++) {
          leds[buttonLeds[i][j]] = CRGB::Black;
        }
        activeLedCount[i] = 0;
      }
    }
  }

  // KITT animation
  if (!anyButtonPressed && (millis() - lastActivityTime > idleTimeout)) {
    if (millis() - lastKitUpdate > kitInterval) {
      lastKitUpdate = millis();

      FastLED.clear();
      leds[kitPos] = colorPalette[kittColorIndex];
      FastLED.show();

      if (kitForward) {
        kitPos++;
        if (kitPos >= NUM_LEDS - 1) {
          kitPos = NUM_LEDS - 1;
          kitForward = false;
        }
      } else {
        kitPos--;
        if (kitPos <= 0) {
          kitPos = 0;
          kitForward = true;
        }
      }
    }
  } else {
    FastLED.show();
  }
}