#include <Button.h>
#include <Adafruit_NeoPixel.h>

constexpr bool KEEPALIVE_ENABLE = true;
constexpr int KEEPALIVE_TIMEOUT = 30000; // ms

constexpr int LEDS_TAIL_PIN = PA6;
constexpr int LEDS_TAIL_LENGTH = 10;
constexpr int LEDS_TAIL_BRIGHTNESS = 255;
constexpr int LEDS_TAIL_SIDE_LENGTH = 2;
constexpr int LEDS_TAIL_SIDE_MARGIN = 2;
constexpr int LEDS_TAIL_RIGHT_RANGE[] = {0, LEDS_TAIL_SIDE_LENGTH};
constexpr int LEDS_TAIL_LEFT_RANGE[] = {LEDS_TAIL_LENGTH - LEDS_TAIL_SIDE_LENGTH, LEDS_TAIL_SIDE_LENGTH};

constexpr int LEDS_LEFT_PIN = PA5;
constexpr int LEDS_RIGHT_PIN = PA7;

/* default configuration right lamp (left from the front)
RING						          EDGE
29 30 31 32 33 34 35 37		0  1  2  3  4  5  6  7  8  9  10
28 27 26 25 24 23 22 21		   20 19 18 17 16 15 14 13 12 11
*/
constexpr int LEDS_TOP_LENGTH = 12;
constexpr int LEDS_BOTTOM_LENGTH = 11;
constexpr int LEDS_RING_LENGTH = 16;
constexpr int LEDS_EDGE_LENGTH = LEDS_TOP_LENGTH + LEDS_BOTTOM_LENGTH;
constexpr int LEDS_HEAD_LENGTH = LEDS_TOP_LENGTH + LEDS_BOTTOM_LENGTH + LEDS_RING_LENGTH;
constexpr int LEDS_HEAD_BRIGHTNESS = 200;

constexpr int LEDS_DASH_PIN = PB0;
constexpr int LEDS_DASH_LENGTH = 10;
constexpr int LEDS_DASH_BRIGHTNESS = 255;

constexpr int PIN_ALIVE = PA2;
constexpr int PIN_BUCK = PA1;

constexpr int PIN_PEDAL = PA0;
constexpr int PIN_PEDAL_FORWARD = PA3;
constexpr int PIN_PEDAL_REVERSE = PA4;

constexpr int PIN_BUTTON_MODE = PA9;
constexpr int PIN_BUTTON_LEFT = PA8;
constexpr int PIN_BUTTON_RIGHT = PA10;

constexpr int LED_ON = LOW;
constexpr int LED_OFF = HIGH;

constexpr int PEDAL_ON = HIGH;
constexpr int PEDAL_OFF = LOW;

constexpr int BUCK_ON = HIGH;
constexpr int BUCK_OFF = LOW;

constexpr int ANIMATION_FPS = 30;
constexpr int ANIMATION_INTERVAL = lround(1000 / ANIMATION_FPS);

Button pedal(PIN_PEDAL);
Button pedalForward(PIN_PEDAL_FORWARD);
Button pedalReverse(PIN_PEDAL_REVERSE);

Button buttonMode(PIN_BUTTON_MODE);
Button buttonLeft(PIN_BUTTON_LEFT);
Button buttonRight(PIN_BUTTON_RIGHT);

Adafruit_NeoPixel ledsTail(LEDS_TAIL_LENGTH, LEDS_TAIL_PIN, NEO_GRBW + NEO_KHZ800);
Adafruit_NeoPixel ledsHeadLeft(LEDS_HEAD_LENGTH, LEDS_LEFT_PIN, NEO_GRBW + NEO_KHZ800);
Adafruit_NeoPixel ledsHeadRight(LEDS_HEAD_LENGTH, LEDS_RIGHT_PIN, NEO_GRBW + NEO_KHZ800);
Adafruit_NeoPixel ledsDash(LEDS_DASH_LENGTH, LEDS_DASH_PIN, NEO_GRBW + NEO_KHZ800);

enum Direction { LEFT, RIGHT };

enum Mode {
  MODE_HEAD,
  MODE_HAZARD,
  MODE_POLICE,
  MODE_RAINBOW,
  MODE_KITT,
  MODE_COUNT,
};

Mode activeMode = MODE_HEAD;
// Mode activeMode = MODE_KITT;

uint32_t colorHSV(float hue, float sat = 1.0, float val = 1.0) {
  return Adafruit_NeoPixel::ColorHSV(
    (uint16_t)(hue / 360.0 * 65535),
    (uint8_t)(sat * 255),
    (uint8_t)(val * 255)
  );
}

const uint32_t COLOR_OFF = ledsTail.Color(0, 0, 0, 0);
const uint32_t COLOR_WHITE = ledsTail.Color(0, 0, 0, 255);
const uint32_t COLOR_WHITE_FULL = ledsTail.Color(255, 255, 255, 255);
const uint32_t COLOR_WHITE_DIM = ledsTail.Color(0, 0, 0, 100);
const uint32_t COLOR_RED_DIM = ledsTail.Color(150, 0, 0, 0);
const uint32_t COLOR_RED = ledsTail.Color(255, 0, 0, 0);
const uint32_t COLOR_AMBER = colorHSV(10); // easier value fade
const uint32_t COLOR_POLICE_BLUE = ledsTail.Color(0, 0, 255, 0);
const uint32_t COLOR_POLICE_RED = ledsTail.Color(255, 0, 0, 0);
const uint32_t COLOR_GREEN = ledsTail.Color(0, 255, 0, 0);

unsigned long animationFrame = 0;
unsigned long stableAnimationFrame = 0;
unsigned long lastFrame = 0;

unsigned long lastPedalPress = millis();
unsigned long lastPedalRelease = 0;

unsigned long lastButtonPress = 0;

bool isReversing = false;

unsigned long indicatorLeftStart = 0;
unsigned long indicatorRightStart = 0;
bool indicatorLeftActive = false;
bool indicatorRightActive = false;

bool policeActive = false;

void setup() {
  // initialize digital pin LED_BUILTIN as an output.
  pinMode(LED_BUILTIN, OUTPUT);
  pinMode(PIN_ALIVE, OUTPUT);  
  pinMode(PIN_BUCK, OUTPUT); 
  pinMode(PIN_PEDAL, INPUT_PULLUP);

  pinMode(PIN_PEDAL_FORWARD, INPUT_PULLUP);
  pinMode(PIN_PEDAL_REVERSE, INPUT_PULLUP);

  pinMode(PIN_BUTTON_MODE, INPUT_PULLUP);
  pinMode(PIN_BUTTON_LEFT, INPUT_PULLUP);
  pinMode(PIN_BUTTON_RIGHT, INPUT_PULLUP);

  // digitalWrite(LED_BUILTIN, LED_ON);
  digitalWrite(PIN_ALIVE, HIGH);
  digitalWrite(PIN_BUCK, BUCK_ON);
  digitalWrite(LED_BUILTIN, LED_OFF);

  ledsTail.setBrightness(LEDS_TAIL_BRIGHTNESS);
  ledsHeadLeft.setBrightness(LEDS_HEAD_BRIGHTNESS);
  ledsHeadRight.setBrightness(LEDS_HEAD_BRIGHTNESS);
  ledsDash.setBrightness(LEDS_DASH_BRIGHTNESS);

  ledsTail.begin();
  ledsHeadLeft.begin();
  ledsHeadRight.begin();
  ledsDash.begin();

  pedal.begin();
  pedalForward.begin();
  pedalReverse.begin();

  buttonMode.begin();
  buttonLeft.begin();
  buttonRight.begin();

  Serial.begin(9600);

  ledsTail.clear();
  ledsHeadLeft.clear();
  ledsHeadRight.clear();
  ledsDash.clear();
}

void animateHeadlights() {
  ledsHeadLeft.fill(COLOR_WHITE);
  ledsHeadRight.fill(COLOR_WHITE);
}

void animateTailLight() {
  ledsTail.fill(COLOR_OFF, 0, LEDS_TAIL_LENGTH);
  ledsTail.fill(COLOR_RED, LEDS_TAIL_SIDE_LENGTH, LEDS_TAIL_LENGTH - LEDS_TAIL_SIDE_LENGTH * 2);
}

void animateRightTailSide(uint32_t color, bool withMargin = true) {
  ledsTail.fill(color, 0, LEDS_TAIL_SIDE_LENGTH);

  if (withMargin) {
    ledsTail.fill(COLOR_OFF, LEDS_TAIL_SIDE_LENGTH, LEDS_TAIL_SIDE_MARGIN);
  }
}

void animateLeftTailSide(uint32_t color, bool withMargin = true) {
  ledsTail.fill(color, LEDS_TAIL_LENGTH - LEDS_TAIL_SIDE_LENGTH, LEDS_TAIL_SIDE_LENGTH);

  if (withMargin) {
    ledsTail.fill(COLOR_OFF, LEDS_TAIL_LENGTH - LEDS_TAIL_SIDE_LENGTH -  LEDS_TAIL_SIDE_MARGIN, LEDS_TAIL_SIDE_MARGIN);
  }
}

void animateDash() {
  ledsDash.fill(COLOR_GREEN);
}

void animateReverse() {
  ledsTail.fill(COLOR_WHITE);
}

void animateBrake() {
    animateLeftTailSide(COLOR_RED, false);
    animateRightTailSide(COLOR_RED, false);
}

constexpr int INDICATOR_DURATION = 3200; // total ms
constexpr int INDICATOR_HOLD_FRAMES = 15;
constexpr int INDICATOR_OFF_FRAMES = 20;
constexpr float INDICATOR_SPEED = 2;

void animateIndicator(Direction direction) {
  Adafruit_NeoPixel& ledsTarget = direction == LEFT
    ? ledsHeadLeft
    : ledsHeadRight;

  const long indicatorFrame = lround(animationFrame * INDICATOR_SPEED) % (LEDS_EDGE_LENGTH + INDICATOR_HOLD_FRAMES + INDICATOR_OFF_FRAMES);

  const long sweepEnd = LEDS_TOP_LENGTH;
  const long holdEnd = LEDS_TOP_LENGTH + INDICATOR_HOLD_FRAMES;

  for (int ledIndex = LEDS_TOP_LENGTH - 1; ledIndex >= 0; ledIndex--) {
    if (indicatorFrame < sweepEnd) {
      // SWEEP
      bool isLit = ledIndex >= (LEDS_TOP_LENGTH - 1 - indicatorFrame);

      ledsTarget.setPixelColor(ledIndex, isLit ? COLOR_AMBER : COLOR_OFF);

      if (ledIndex > 0) {
        ledsTarget.setPixelColor(LEDS_EDGE_LENGTH - ledIndex, isLit ? COLOR_AMBER : COLOR_OFF);
      }
    } else if (indicatorFrame < holdEnd) {
      // HOLD
      ledsTarget.setPixelColor(ledIndex, COLOR_AMBER);

      if (ledIndex > 0) {
        ledsTarget.setPixelColor(LEDS_EDGE_LENGTH - ledIndex, COLOR_AMBER);
      }
    } else {
      // OFF
      ledsTarget.fill(COLOR_OFF, 0, LEDS_EDGE_LENGTH);
    }
  }

  void (*animateTailSide)(uint32_t, bool) = direction == LEFT
    ? animateLeftTailSide
    : animateRightTailSide;

  const int dashIndex = direction == LEFT
    ? 1
    : 0;

  if (indicatorFrame < sweepEnd) {
    const int colorAmberFaded = colorHSV(10, 1, (float)indicatorFrame / sweepEnd);

    animateTailSide(colorAmberFaded, true);
    ledsDash.setPixelColor(dashIndex, colorAmberFaded);
  } else if (indicatorFrame < holdEnd) {
    animateTailSide(COLOR_AMBER, true);
    ledsDash.setPixelColor(dashIndex, COLOR_AMBER);
  } else {
    animateTailSide(COLOR_OFF, true);
    ledsDash.setPixelColor(dashIndex, COLOR_OFF);
  }
}

constexpr float POLICE_SPEED = 0.75;

enum PoliceStates { POLICE_LEFT, POLICE_RIGHT, POLICE_OFF };

const PoliceStates policeBeaconFrames[12] = {
  POLICE_LEFT,
  POLICE_OFF,
  POLICE_LEFT,
  POLICE_OFF,
  POLICE_OFF,
  POLICE_OFF,
  POLICE_RIGHT,
  POLICE_OFF,
  POLICE_RIGHT,
  POLICE_OFF,
  POLICE_OFF,
  POLICE_OFF,
};

void animatePolice() {
  const long policeFrame = lround(animationFrame * POLICE_SPEED) % 12;

  if (policeBeaconFrames[policeFrame] == POLICE_LEFT) {
    ledsHeadLeft.fill(COLOR_POLICE_BLUE, 0, LEDS_EDGE_LENGTH);
    ledsHeadRight.fill(COLOR_OFF, 0, LEDS_EDGE_LENGTH);

    // deliberately inverted from edge lights for a more alternating effect
    animateLeftTailSide(COLOR_OFF);
    animateRightTailSide(COLOR_POLICE_BLUE);

    ledsDash.setPixelColor(1, COLOR_POLICE_BLUE);
    ledsDash.setPixelColor(0, COLOR_OFF);
  } else if (policeBeaconFrames[policeFrame] == POLICE_RIGHT) {
    ledsHeadLeft.fill(COLOR_OFF, 0, LEDS_EDGE_LENGTH);
    ledsHeadRight.fill(COLOR_POLICE_BLUE, 0, LEDS_EDGE_LENGTH);

    animateLeftTailSide(COLOR_POLICE_BLUE);
    animateRightTailSide(COLOR_OFF);

    ledsDash.setPixelColor(1, COLOR_OFF);
    ledsDash.setPixelColor(0, COLOR_POLICE_BLUE);
  } else {
    ledsHeadLeft.fill(COLOR_OFF, 0, LEDS_EDGE_LENGTH);
    ledsHeadRight.fill(COLOR_OFF, 0, LEDS_EDGE_LENGTH);

    animateLeftTailSide(COLOR_OFF);
    animateRightTailSide(COLOR_OFF);

    ledsDash.setPixelColor(1, COLOR_OFF);
    ledsDash.setPixelColor(0, COLOR_OFF);
  }

  if (policeFrame < 6) {
    ledsHeadLeft.fill(COLOR_WHITE_DIM, LEDS_EDGE_LENGTH, LEDS_EDGE_LENGTH + LEDS_RING_LENGTH);
    ledsHeadRight.fill(COLOR_WHITE, LEDS_EDGE_LENGTH, LEDS_EDGE_LENGTH + LEDS_RING_LENGTH);
  } else {
    ledsHeadLeft.fill(COLOR_WHITE, LEDS_EDGE_LENGTH, LEDS_EDGE_LENGTH + LEDS_RING_LENGTH);
    ledsHeadRight.fill(COLOR_WHITE_DIM, LEDS_EDGE_LENGTH, LEDS_EDGE_LENGTH + LEDS_RING_LENGTH);
  }
}

constexpr float RAINBOW_SPEED = 2.0;
constexpr float RAINBOW_DENSITY = 10.0;

void animateRainbow() {
  // EDGES, flowing color
  for (int ledIndex = LEDS_TOP_LENGTH - 1; ledIndex >= 0; ledIndex--) {
    const uint32_t pixelColor = colorHSV((stableAnimationFrame * RAINBOW_SPEED) + (ledIndex * RAINBOW_DENSITY));

    ledsHeadLeft.setPixelColor(ledIndex, pixelColor);
    ledsHeadLeft.setPixelColor(LEDS_EDGE_LENGTH - ledIndex, pixelColor);

    ledsHeadRight.setPixelColor(ledIndex, pixelColor);
    ledsHeadRight.setPixelColor(LEDS_EDGE_LENGTH - ledIndex, pixelColor);
  }

  // RINGS, single color
  ledsHeadLeft.fill(colorHSV(stableAnimationFrame * RAINBOW_SPEED), LEDS_EDGE_LENGTH, LEDS_RING_LENGTH);
  ledsHeadRight.fill(colorHSV(stableAnimationFrame * RAINBOW_SPEED), LEDS_EDGE_LENGTH, LEDS_RING_LENGTH);

  for (int ledIndex = 0; ledIndex < LEDS_TAIL_LENGTH; ledIndex++) {
    const int distFromCenter = abs(ledIndex * 2 - (LEDS_TAIL_LENGTH - 1));
    const uint32_t pixelColor = colorHSV((stableAnimationFrame * RAINBOW_SPEED) - (distFromCenter * RAINBOW_DENSITY * 2));

    ledsTail.setPixelColor(ledIndex, pixelColor);
  }

  ledsDash.setPixelColor(1, colorHSV(stableAnimationFrame * RAINBOW_SPEED));
  ledsDash.setPixelColor(0, colorHSV(stableAnimationFrame * RAINBOW_SPEED + 90));
}

constexpr float KITT_SPEED = 0.5;
constexpr int KITT_TRAIL = 4;
constexpr float KITT_TAIL_SPEED = 3.0f; 

int kittScanPos(int vPos) {
  if (vPos < LEDS_RING_LENGTH) {
    // first ring — reversed so it flows naturally into the edge
    return LEDS_EDGE_LENGTH + (LEDS_RING_LENGTH - 1 - vPos);
  } else if (vPos < LEDS_RING_LENGTH + LEDS_EDGE_LENGTH) {
    // edge forward
    return vPos - LEDS_RING_LENGTH;
  } else if (vPos < 2 * LEDS_RING_LENGTH + LEDS_EDGE_LENGTH) {
    // second ring
    return LEDS_EDGE_LENGTH + (vPos - LEDS_RING_LENGTH - LEDS_EDGE_LENGTH);
  } else {
    // edge backward
    return LEDS_EDGE_LENGTH - 1 - (vPos - 2 * LEDS_RING_LENGTH - LEDS_EDGE_LENGTH);
  }
}

void animateKitt() {
  const long kittFrame = lround(stableAnimationFrame * KITT_SPEED);
  const int cycleLength = 2 * LEDS_EDGE_LENGTH + 2 * LEDS_RING_LENGTH;
  const int vPos = kittFrame % cycleLength;

  ledsHeadLeft.fill(COLOR_OFF);
  ledsHeadRight.fill(COLOR_OFF);

  // trail behind
  for (int t = 0; t < KITT_TRAIL; t++) {
    const int pastVPos = ((vPos - t) % cycleLength + cycleLength) % cycleLength;
    const int scanPos = kittScanPos(pastVPos);

    const uint32_t color = colorHSV(0, 1, 1.0f - (float)t / KITT_TRAIL);
    ledsHeadLeft.setPixelColor(scanPos, color);
    ledsHeadRight.setPixelColor(scanPos, color);
  }

  // lead ahead — shorter and dimmer than trail
  constexpr int KITT_LEAD = 2;
  for (int t = 1; t <= KITT_LEAD; t++) {
    const int futureVPos = ((vPos + t) % cycleLength + cycleLength) % cycleLength;
    const int scanPos = kittScanPos(futureVPos);

    const uint32_t color = colorHSV(0, 1, 0.3f * (1.0f - (float)t / KITT_LEAD));
    ledsHeadLeft.setPixelColor(scanPos, color);
    ledsHeadRight.setPixelColor(scanPos, color);
  }

  const int scanPos = kittScanPos(vPos);
  

  // normalized cycle (0–1), sped up
  float tailCycle = (float)vPos * KITT_TAIL_SPEED / cycleLength;
  tailCycle -= (int)tailCycle;  // wrap to 0–1

  // triangle wave: 0 → 1 → 0 (KITT bounce)
  float tailT = tailCycle * 2.0f;
  if (tailT > 1.0f) tailT = 2.0f - tailT;

  // optional smoothing (gives that analog feel)
  tailT = tailT * tailT * (3.0f - 2.0f * tailT); // smoothstep

  // map to LED index
  const int tailPos = lround(tailT * (LEDS_TAIL_LENGTH - 1));

  // draw tail with falloff
  for (int i = 0; i < LEDS_TAIL_LENGTH; i++) {
    const int dist = abs(i - tailPos);
    float brightness = 0.0f;

    if (dist < KITT_TRAIL) {
      brightness = 1.0f - (float)dist / KITT_TRAIL;
    }

    ledsTail.setPixelColor(i, colorHSV(0, 1, brightness));
  }

  ledsDash.fill(COLOR_OFF);

  float speed = 3.0;

  float dashCycle = (float)vPos * speed / cycleLength;
  dashCycle -= (int)dashCycle;  // keep it in 0–1

  // triangle wave: 0→1→0 over the cycle
  float t = dashCycle * 2.0;
  if (t > 1.0) t = 2.0 - t;

  ledsDash.setPixelColor(0, colorHSV(0, 1, t));
  ledsDash.setPixelColor(1, colorHSV(0, 1, 1.0 - t));
}

void animate(long now) {
  // always animate headlights, let other modes override
  animateHeadlights();
  animateTailLight();
  animateDash();

  if (isReversing) {
    animateReverse();
  }

  if (activeMode == MODE_POLICE) {
    animatePolice();
  }

  if (activeMode == MODE_HAZARD) {
    animateIndicator(LEFT);
    animateIndicator(RIGHT);
  }

  if (activeMode == MODE_KITT) {
    animateKitt();
  }

  if (now - lastPedalRelease < 1500) {
    animateBrake();
  }

  if (activeMode == MODE_RAINBOW) {
    animateRainbow();
  }

  // we need the boolean to prevent the indicators from turning on right after boot
  if (indicatorLeftActive && now - indicatorLeftStart < INDICATOR_DURATION) {
    animateIndicator(LEFT);
  }

  if (indicatorRightActive && now - indicatorRightStart < INDICATOR_DURATION) {
    animateIndicator(RIGHT);
  }
}

void render() {
    ledsTail.show();
    ledsHeadLeft.show();
    ledsHeadRight.show();
    ledsDash.show();
}

bool isOff = false;

void shutdown() {
    // digitalWrite(PIN_BUCK, BUCK_OFF);
  Serial.println("SHUTTING DOWN");

  Serial.println("BUCK OFF");
  digitalWrite(PIN_BUCK, LOW);

  delay(100);

  Serial.println("SELF OFF, GOODBYE");
  digitalWrite(PIN_ALIVE, LOW);

  isOff = true;
}

// only used when STLink or USB prevents full power cut
void restart() {
  lastFrame = millis();
  animationFrame = 0;
  stableAnimationFrame = 0;
  isOff = false;

  digitalWrite(PIN_BUCK, BUCK_ON);
}

unsigned long modePressStart = 0;
bool modePressed = false;

// the loop function runs over and over again forever
void loop() {
  const unsigned long now = millis();

  if (buttonMode.pressed()) {
    modePressed = true;
    modePressStart = now;

    activeMode = (Mode)((activeMode + 1) % MODE_COUNT);
    animationFrame = 0; // ensures animations start cleanly instead of halfway

    lastButtonPress = now;

    Serial.print("BUTTON MODE ");
    Serial.println(activeMode);
  }

  if (buttonMode.released()) {
    modePressed = false;
  }

  if (modePressed && now - modePressStart > 5000) {
    shutdown();
  }

  if (buttonLeft.pressed()) {
    animationFrame = 0;
    
    indicatorLeftActive = true;
    indicatorLeftStart = now;

    indicatorRightActive = false;
    indicatorRightStart = 0;

    lastButtonPress = now;

    Serial.println("BUTTON LEFT");
  }

  if (buttonRight.pressed()) {
    animationFrame = 0;

    indicatorRightActive = true;
    indicatorRightStart = now;

    indicatorLeftActive = false;
    indicatorLeftStart = 0;

    lastButtonPress = now;

    Serial.println("BUTTON RIGHT");
  }

  if (pedalReverse.read() == Button::PRESSED && pedalForward.read() != Button::PRESSED) {
    lastPedalPress = now;
    isReversing = true;

    if (isOff) {
      restart();
    }

    Serial.println("PEDAL REVERSE");
  } else {
    isReversing = false;
  }

  if (pedalForward.read() == Button::PRESSED) {
    lastPedalPress = now;

    if (isOff) {
      restart();
    }

    Serial.println("PEDAL FORWARD");
  }`3 

  if (pedal.released()) {
    lastPedalRelease = now;
  }

  if (!isOff && now - lastFrame >= ANIMATION_INTERVAL) {
    lastFrame += ANIMATION_INTERVAL;
    animationFrame += 1;
    stableAnimationFrame += 1;

    animate(now);
    render();
  }

  if (KEEPALIVE_ENABLE && now - lastPedalPress > KEEPALIVE_TIMEOUT && now - lastButtonPress > KEEPALIVE_TIMEOUT) {
    shutdown();
  }
}