#include <WiFi.h>
#include <esp_now.h>
#include <esp_wifi.h>

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

#include <AccelStepper.h>
#include <math.h>
#include <string.h>

// ============================================================
// Pins
// ============================================================
#define STEP_PIN  D0
#define DIR_PIN   D1
#define LIMIT_PIN D2

#ifndef SDA_PIN
#define SDA_PIN D4
#endif

#ifndef SCL_PIN
#define SCL_PIN D5
#endif

// ============================================================
// OLED
// ============================================================
#define SCREEN_WIDTH 128
#define SCREEN_HEIGHT 64
static const uint8_t OLED_ADDR = 0x3C;

Adafruit_SSD1306 display(SCREEN_WIDTH, SCREEN_HEIGHT, &Wire, -1);

// ============================================================
// Stepper
// ============================================================
static const bool DIR_POLARITY = false;
AccelStepper stepper(AccelStepper::DRIVER, STEP_PIN, DIR_PIN);

static const long STEPS_PER_REV = 200;

static inline long degToSteps(float deg) {
  return lroundf((deg / 360.0f) * (float)STEPS_PER_REV);
}

// Motion tuning
static const float HOME_SPEED_SPS = 120.0f;
static const float MOVE_SPEED_SPS = 600.0f;
static const float ACCEL_SPS2     = 1500.0f;

// Tuned angles
static const long LOAD_OFFSET_STEPS = degToSteps(85.0f);
static const long SLAP_OFFSET_STEPS = degToSteps(10.0f);

// Debounce / ignore
static const uint32_t RELEASE_DEBOUNCE_MS = 40;
static const uint32_t SWITCH_IGNORE_MS    = 1000;
static uint32_t switchIgnoreUntilMs = 0;

// ============================================================
// ESP-NOW packet
// ============================================================
typedef struct __attribute__((packed)) {
  uint8_t card_id; // 1..52
} CardPacket;

// ============================================================
// Game rule toggles
// ============================================================
struct RuleToggles {
  bool doubles   = true;
  bool sandwich  = true;
  bool marriage  = false;
  bool divorce   = false;
  bool addTo10   = false;
};

static RuleToggles rules;

// ============================================================
// Card / pile memory
// ============================================================

// Rule-evaluation pile memory
static const uint8_t MAX_PILE = 12;
static uint8_t pileRanks[MAX_PILE];
static uint8_t pileCount = 0;

// OLED card history memory
static const uint8_t CARD_HISTORY_LEN = 6;
static uint8_t cardHistory[CARD_HISTORY_LEN];
static uint8_t cardHistoryCount = 0;

// Mailbox from ESP-NOW callback -> main loop
static volatile bool rxCardPending = false;
static volatile uint8_t rxCardId = 0;

// Slap request generated by rules
static volatile bool slapRequested = false;

// Status text for OLED / serial
static char latestCardText[4] = "--";
static char latestReason[16]  = "Waiting";
static bool hasAnyCard = false;

// ============================================================
// State machine
// ============================================================
enum State {
  WAIT_FOR_PRESS,               // move CW while waiting for switch press
  HOME_MOVE_UNTIL_RELEASE_STABLE,
  MOVE_TO_LOAD,
  READY,
  MOVE_TO_SLAP_STRIKE,
  SLAP_SEEK_PRESS,
  POST_SLAP_WAIT_1S,
  POST_SLAP_SEEK_RELEASE_500
};

static State state = WAIT_FOR_PRESS;

static uint32_t releaseCandidateSinceMs = 0;
static uint32_t settleUntilMs = 0;

// ============================================================
// Limit switch helpers
// ============================================================
static inline bool limitPressedRaw() {
  return digitalRead(LIMIT_PIN) == LOW; // INPUT_PULLUP => pressed = LOW
}

static inline bool limitPressedFiltered() {
  if ((int32_t)(millis() - switchIgnoreUntilMs) < 0) return false;
  return limitPressedRaw();
}

static inline void beginSwitchIgnore(uint32_t durationMs = SWITCH_IGNORE_MS) {
  switchIgnoreUntilMs = millis() + durationMs;
}

// ============================================================
// Stepper helpers
// ============================================================
static void setDirectionPolarity() {
  stepper.setPinsInverted(false, DIR_POLARITY, false);
}

static void startLoadFromZero() {
  stepper.setMaxSpeed(MOVE_SPEED_SPS);
  stepper.setAcceleration(ACCEL_SPS2);
  stepper.moveTo(LOAD_OFFSET_STEPS);
  state = MOVE_TO_LOAD;
}

static void startSlapStrikeCW() {
  stepper.setMaxSpeed(MOVE_SPEED_SPS);
  stepper.setAcceleration(ACCEL_SPS2);
  stepper.move(SLAP_OFFSET_STEPS);
  state = MOVE_TO_SLAP_STRIKE;
}

static void startSeekPressCW() {
  stepper.setMaxSpeed(HOME_SPEED_SPS);
  stepper.setSpeed(fabs(HOME_SPEED_SPS));
  state = SLAP_SEEK_PRESS;
}

static void stopMotorNow() {
  stepper.setSpeed(0);
  stepper.stop();
  stepper.moveTo(stepper.currentPosition());
}

// ============================================================
// Card helpers
// ID order:
// 1..13  = AC..KC
// 14..26 = AH..KH
// 27..39 = AS..KS
// 40..52 = AD..KD
// ============================================================
static inline uint8_t cardIdToRank(uint8_t id) {
  return (uint8_t)(((id - 1) % 13) + 1); // 1..13
}

static void cardIdToText(uint8_t id, char out[4]) {
  out[0] = '\0';

  if (id < 1 || id > 52) {
    out[0] = '?';
    out[1] = '\0';
    return;
  }

  uint8_t suitIndex = (id - 1) / 13;      // 0=C,1=H,2=S,3=D
  uint8_t rankIndex = (id - 1) % 13 + 1;  // 1..13

  if (rankIndex == 1) {
    out[0] = 'A'; out[1] = '\0';
  } else if (rankIndex >= 2 && rankIndex <= 9) {
    out[0] = char('0' + rankIndex); out[1] = '\0';
  } else if (rankIndex == 10) {
    out[0] = '1'; out[1] = '0'; out[2] = '\0';
  } else if (rankIndex == 11) {
    out[0] = 'J'; out[1] = '\0';
  } else if (rankIndex == 12) {
    out[0] = 'Q'; out[1] = '\0';
  } else {
    out[0] = 'K'; out[1] = '\0';
  }

  char suitChar = 'C';
  if (suitIndex == 0) suitChar = 'C';
  else if (suitIndex == 1) suitChar = 'H';
  else if (suitIndex == 2) suitChar = 'S';
  else suitChar = 'D';

  size_t n = strlen(out);
  out[n] = suitChar;
  out[n + 1] = '\0';
}

// ============================================================
// Pile / history helpers
// ============================================================
static inline void clearPile() {
  pileCount = 0;
}

static inline void pushRank(uint8_t r) {
  if (pileCount < MAX_PILE) {
    pileRanks[pileCount++] = r;
  } else {
    for (uint8_t i = 1; i < MAX_PILE; i++) pileRanks[i - 1] = pileRanks[i];
    pileRanks[MAX_PILE - 1] = r;
  }
}

static inline uint8_t lastRank(uint8_t back) {
  return pileRanks[pileCount - back];
}

static void pushCardHistory(uint8_t id) {
  if (cardHistoryCount < CARD_HISTORY_LEN) {
    cardHistory[cardHistoryCount++] = id;
  } else {
    for (uint8_t i = 1; i < CARD_HISTORY_LEN; i++) cardHistory[i - 1] = cardHistory[i];
    cardHistory[CARD_HISTORY_LEN - 1] = id;
  }
}

// ============================================================
// Rules
// ============================================================
static inline bool checkDoubles() {
  return (pileCount >= 2) && (lastRank(1) == lastRank(2));
}

static inline bool checkSandwich() {
  return (pileCount >= 3) && (lastRank(1) == lastRank(3));
}

static inline bool checkMarriage() {
  if (pileCount < 2) return false;
  uint8_t a = lastRank(1), b = lastRank(2);
  return (a == 13 && b == 12) || (a == 12 && b == 13);
}

static inline bool checkDivorce() {
  if (pileCount < 2) return false;
  uint8_t a = lastRank(1), b = lastRank(2);
  return (a == 13 && b == 11) || (a == 11 && b == 13);
}

static inline bool checkAddTo10() {
  if (pileCount < 2) return false;
  uint8_t a = lastRank(1), b = lastRank(2);
  return (a + b) == 10;
}

static inline void triggerSlapAndClear(const char* reason) {
  if (!slapRequested) {
    slapRequested = true;
    strncpy(latestReason, reason, sizeof(latestReason) - 1);
    latestReason[sizeof(latestReason) - 1] = '\0';

    Serial.print("SLAP! Reason: ");
    Serial.println(reason);
  }
  clearPile();
}

static inline void evaluateRules() {
  if (rules.doubles  && checkDoubles())  { triggerSlapAndClear("Doubles");  return; }
  if (rules.sandwich && checkSandwich()) { triggerSlapAndClear("Sandwich"); return; }
  if (rules.marriage && checkMarriage()) { triggerSlapAndClear("Marriage"); return; }
  if (rules.divorce  && checkDivorce())  { triggerSlapAndClear("Divorce");  return; }
  if (rules.addTo10  && checkAddTo10())  { triggerSlapAndClear("AddTo10");  return; }

  strncpy(latestReason, "No slap", sizeof(latestReason) - 1);
  latestReason[sizeof(latestReason) - 1] = '\0';
}

// ============================================================
// OLED drawing
// ============================================================
static void drawScreen() {
  display.clearDisplay();
  display.setTextColor(SSD1306_WHITE);

  // Header
  display.setTextSize(1);
  display.setCursor(0, 0);
  display.print("State: ");
  switch (state) {
    case WAIT_FOR_PRESS:               display.println("HOME_WAIT"); break;
    case HOME_MOVE_UNTIL_RELEASE_STABLE: display.println("HOME_REL"); break;
    case MOVE_TO_LOAD:                 display.println("LOAD"); break;
    case READY:                        display.println("READY"); break;
    case MOVE_TO_SLAP_STRIKE:          display.println("SLAP"); break;
    case SLAP_SEEK_PRESS:              display.println("SEEK_SW"); break;
    case POST_SLAP_WAIT_1S:            display.println("SETTLE"); break;
    case POST_SLAP_SEEK_RELEASE_500:   display.println("RELOAD"); break;
    default:                           display.println("?"); break;
  }

  // Latest card
  display.setTextSize(2);
  display.setCursor(0, 14);
  display.print("Card:");
  display.setCursor(60, 14);
  display.print(latestCardText);

  // Reason
  display.setTextSize(1);
  display.setCursor(0, 36);
  display.print("Rule: ");
  display.println(latestReason);

  // History
  display.setCursor(0, 48);
  display.print("Last: ");
  for (uint8_t i = 0; i < cardHistoryCount; i++) {
    char txt[4];
    cardIdToText(cardHistory[i], txt);
    display.print(txt);
    if (i + 1 < cardHistoryCount) display.print(" ");
  }

  display.display();
}

// ============================================================
// ESP-NOW callback
// ============================================================
static void onEspNowRecv(const esp_now_recv_info_t *info, const uint8_t *data, int len) {
  (void)info;

  if (len != (int)sizeof(CardPacket)) return;

  CardPacket pkt;
  memcpy(&pkt, data, sizeof(pkt));

  if (pkt.card_id < 1 || pkt.card_id > 52) return;

  rxCardId = pkt.card_id;
  rxCardPending = true;
}

// ============================================================
// ESP-NOW init
// ============================================================
static bool initEspNowReceiver(uint8_t channel = 1) {
  WiFi.mode(WIFI_STA);
  WiFi.setSleep(false);

  esp_wifi_set_promiscuous(true);
  esp_wifi_set_channel(channel, WIFI_SECOND_CHAN_NONE);
  esp_wifi_set_promiscuous(false);

  if (esp_now_init() != ESP_OK) return false;
  esp_now_register_recv_cb(onEspNowRecv);
  return true;
}

// ============================================================
// Process received card in main loop
// ============================================================
static void processPendingCard() {
  if (!rxCardPending) return;

  noInterrupts();
  uint8_t id = rxCardId;
  rxCardPending = false;
  interrupts();

  uint8_t rank = cardIdToRank(id);

  pushRank(rank);
  pushCardHistory(id);
  cardIdToText(id, latestCardText);
  hasAnyCard = true;

  Serial.print("Received: ");
  Serial.print(latestCardText);
  Serial.print("  ID=");
  Serial.print(id);
  Serial.print("  Rank=");
  Serial.print(rank);
  Serial.print("  pileCount=");
  Serial.println(pileCount);

  evaluateRules();
  drawScreen();
}

// ============================================================
// Setup
// ============================================================
void setup() {
  Serial.begin(115200);

  pinMode(LIMIT_PIN, INPUT_PULLUP);

  // I2C / OLED
  Wire.begin(SDA_PIN, SCL_PIN);
  if (!display.begin(SSD1306_SWITCHCAPVCC, OLED_ADDR)) {
    Serial.println("OLED init failed");
    while (true) delay(1000);
  }

  // Stepper
  stepper.setMinPulseWidth(2);
  setDirectionPolarity();

  // ESP-NOW
  if (!initEspNowReceiver(1)) {
    Serial.println("ESP-NOW init failed");
    display.clearDisplay();
    display.setTextSize(1);
    display.setCursor(0, 0);
    display.println("ESP-NOW FAILED");
    display.display();
    while (true) delay(1000);
  }

  strncpy(latestReason, "Waiting", sizeof(latestReason) - 1);
  latestReason[sizeof(latestReason) - 1] = '\0';

  // Start homing motion immediately:
  // wait for switch press while moving CW
  stepper.setMaxSpeed(HOME_SPEED_SPS);
  stepper.setSpeed(fabs(HOME_SPEED_SPS));
  state = WAIT_FOR_PRESS;

  drawScreen();

  Serial.println("Merged receiver ready.");
  Serial.println("Rules: Doubles + Sandwich enabled.");
  Serial.println("Homing: moving CW until switch press, then CW until release, then load.");
}

// ============================================================
// Main loop
// ============================================================
void loop() {
  // 1) Service motor continuously
  switch (state) {
    case WAIT_FOR_PRESS:
    case HOME_MOVE_UNTIL_RELEASE_STABLE:
    case SLAP_SEEK_PRESS:
    case POST_SLAP_SEEK_RELEASE_500:
      stepper.runSpeed();
      break;

    case MOVE_TO_LOAD:
    case MOVE_TO_SLAP_STRIKE:
      stepper.run();
      break;

    default:
      break;
  }

  // 2) Process any newly received card
  processPendingCard();

  // 3) Inputs / timing
  bool lp_raw  = limitPressedRaw();
  bool lp_filt = limitPressedFiltered();
  uint32_t now = millis();

  // 4) Motor state machine
  switch (state) {
    case WAIT_FOR_PRESS:
      // Moving CW already. Wait for switch to become pressed.
      if (lp_raw) {
        Serial.println("Homing: switch pressed. Continue CW until release...");
        releaseCandidateSinceMs = 0;
        state = HOME_MOVE_UNTIL_RELEASE_STABLE;
        drawScreen();
      }
      break;

    case HOME_MOVE_UNTIL_RELEASE_STABLE:
      if (!lp_raw) {
        if (releaseCandidateSinceMs == 0) releaseCandidateSinceMs = now;
        if ((now - releaseCandidateSinceMs) >= RELEASE_DEBOUNCE_MS) {
          stepper.setCurrentPosition(0);
          Serial.println("Homing: released stable. Home=0. Loading...");
          startLoadFromZero();
          drawScreen();
        }
      } else {
        releaseCandidateSinceMs = 0;
      }
      break;

    case MOVE_TO_LOAD:
      if (stepper.distanceToGo() == 0) {
        Serial.println("Loaded. READY.");
        state = READY;
        drawScreen();
      }
      break;

    case READY:
      if (slapRequested) {
        slapRequested = false;
        Serial.println("Executing slap...");
        startSlapStrikeCW();
        drawScreen();
      }
      break;

    case MOVE_TO_SLAP_STRIKE:
      if (stepper.distanceToGo() == 0) {
        Serial.println("Strike done. Seeking switch press...");
        startSeekPressCW();
        drawScreen();
      }
      break;

    case SLAP_SEEK_PRESS:
      if (lp_filt) {
        Serial.println("Limit pressed. Settling...");
        beginSwitchIgnore(SWITCH_IGNORE_MS);
        stopMotorNow();
        settleUntilMs = now + SWITCH_IGNORE_MS;
        state = POST_SLAP_WAIT_1S;
        drawScreen();
      }
      break;

    case POST_SLAP_WAIT_1S:
      if ((int32_t)(now - settleUntilMs) >= 0) {
        if (limitPressedRaw()) {
          Serial.println("Switch still pressed. Seeking release at move speed...");
          stepper.setMaxSpeed(MOVE_SPEED_SPS);
          stepper.setSpeed(fabs(MOVE_SPEED_SPS));
          state = POST_SLAP_SEEK_RELEASE_500;
          drawScreen();
        } else {
          Serial.println("Switch not pressed after settle. Seeking press again...");
          startSeekPressCW();
          drawScreen();
        }
      }
      break;

    case POST_SLAP_SEEK_RELEASE_500:
      if (!lp_raw) {
        stepper.setSpeed(0);
        stepper.setCurrentPosition(0);
        Serial.println("Released -> set 0. Loading...");
        startLoadFromZero();
        drawScreen();
      }
      break;
  }

  delay(2);
}