#include "USB.h"
#include "USBHIDKeyboard.h"
#include "esp_system.h"
#include "soc/rtc_cntl_reg.h"

#define FW_VERSION "1.3.0"

USBHIDKeyboard Keyboard;

const int BUTTON_PIN = 2;
const int EXT_LED_PIN = 3;
const int HEARTBEAT_LED_PIN = 21;

const unsigned long DEBOUNCE_MS = 25;
const unsigned long BLINK_ON_MS = 60;
const unsigned long BLINK_OFF_MS = 60;
const unsigned long HEARTBEAT_INTERVAL = 500;
const unsigned long LONG_PRESS_MS = 15000;
const unsigned long HOLD_INDICATOR_START = 10000;
const unsigned long HOLD_FLASH_INTERVAL = 200;

bool lastReading = HIGH;
bool buttonState = HIGH;
unsigned long lastDebounceTime = 0;

bool blinking = false;
int blinkStep = 0;
unsigned long blinkTimer = 0;

bool heartbeatState = false;
unsigned long heartbeatTimer = 0;

unsigned long pressStartTime = 0;
bool longPressHandled = false;

bool holdIndicatorActive = false;
bool holdFlashState = false;
unsigned long holdFlashTimer = 0;

void enterBootloader() {
  for (int i = 0; i < 8; i++) {
    digitalWrite(EXT_LED_PIN, (i % 2 == 0) ? LOW : HIGH);
    digitalWrite(HEARTBEAT_LED_PIN, (i % 2 == 0) ? HIGH : LOW);
    delay(60);
  }
  REG_WRITE(RTC_CNTL_OPTION1_REG, RTC_CNTL_FORCE_DOWNLOAD_BOOT);
  esp_restart();
}

void startBlink() {
  blinking = true;
  blinkStep = 0;
  blinkTimer = millis();
  digitalWrite(EXT_LED_PIN, LOW);
}

void updateBlink() {
  if (!blinking) return;
  unsigned long elapsed = millis() - blinkTimer;
  unsigned long stepDuration = (blinkStep % 2 == 0) ? BLINK_OFF_MS : BLINK_ON_MS;

  if (elapsed >= stepDuration) {
    blinkStep++;
    blinkTimer = millis();
    if (blinkStep >= 4) {
      blinking = false;
      digitalWrite(EXT_LED_PIN, HIGH);
    } else {
      digitalWrite(EXT_LED_PIN, blinkStep % 2 == 0 ? LOW : HIGH);
    }
  }
}

void updateHeartbeat() {
  if (millis() - heartbeatTimer >= HEARTBEAT_INTERVAL) {
    heartbeatTimer = millis();
    heartbeatState = !heartbeatState;
    digitalWrite(HEARTBEAT_LED_PIN, heartbeatState ? LOW : HIGH);
  }
}

void updateHoldIndicator() {
  if (!holdIndicatorActive) return;
  if (millis() - holdFlashTimer >= HOLD_FLASH_INTERVAL) {
    holdFlashTimer = millis();
    holdFlashState = !holdFlashState;
    digitalWrite(EXT_LED_PIN, holdFlashState ? LOW : HIGH);
    digitalWrite(HEARTBEAT_LED_PIN, holdFlashState ? HIGH : LOW);
  }
}

void setup() {
  pinMode(BUTTON_PIN, INPUT_PULLUP);
  pinMode(EXT_LED_PIN, OUTPUT);
  pinMode(HEARTBEAT_LED_PIN, OUTPUT);
  digitalWrite(EXT_LED_PIN, HIGH);
  digitalWrite(HEARTBEAT_LED_PIN, HIGH);

  USB.productName("SpaceButton FW " FW_VERSION);
  USB.manufacturerName("GJ");

  Keyboard.begin();
  USB.begin();
}

void loop() {
  bool reading = digitalRead(BUTTON_PIN);

  if (reading != lastReading) {
    lastDebounceTime = millis();
  }

  if ((millis() - lastDebounceTime) > DEBOUNCE_MS) {
    if (reading != buttonState) {
      buttonState = reading;

      if (buttonState == LOW) {
        Keyboard.press(' ');
        Keyboard.release(' ');
        startBlink();
        pressStartTime = millis();
        longPressHandled = false;
        holdIndicatorActive = false;
      } else {
        if (holdIndicatorActive) {
          holdIndicatorActive = false;
          digitalWrite(EXT_LED_PIN, HIGH);
          digitalWrite(HEARTBEAT_LED_PIN, HIGH);
        }
      }
    }
  }

  if (buttonState == LOW && !longPressHandled) {
    unsigned long held = millis() - pressStartTime;

    if (held >= HOLD_INDICATOR_START && !holdIndicatorActive) {
      holdIndicatorActive = true;
      holdFlashTimer = millis();
    }

    if (held >= LONG_PRESS_MS) {
      longPressHandled = true;
      enterBootloader();
    }
  }

  if (holdIndicatorActive) {
    updateHoldIndicator();
  } else {
    updateBlink();
    updateHeartbeat();
  }

  lastReading = reading;
}