#include <Arduino.h>
#include <SPI.h>
#include <Adafruit_GFX.h>
#include <Adafruit_ILI9341.h>
#include <XPT2046_Touchscreen.h>
#include <MFRC522.h>
#include <cstring>

/*
  SmartLockerESP32.ino
  - Fixes shared-SPI conflicts between TFT, touch, and RFID
  - Adds full UI with touch buttons:
      1) Assign Locker
      2) Retrieve
      3) Cancel
  - Integrates button flow with locker assignment/retrieval logic
  - Uses millis() (no delay)
*/

// ------------------------------------------------------------
// Shared SPI bus pins
// ------------------------------------------------------------
constexpr uint8_t PIN_SPI_SCK = 18;
constexpr uint8_t PIN_SPI_MOSI = 23;
constexpr uint8_t PIN_SPI_MISO = 19;

// ------------------------------------------------------------
// Device pins
// ------------------------------------------------------------
constexpr uint8_t PIN_TFT_CS = 15;
constexpr uint8_t PIN_TFT_DC = 21;
constexpr uint8_t PIN_TFT_RST = 4;

constexpr uint8_t PIN_TOUCH_CS = 22;
constexpr uint8_t PIN_TOUCH_IRQ = 34;

constexpr uint8_t PIN_RFID_SS = 5;
constexpr uint8_t PIN_RFID_RST = 27;

constexpr uint8_t PIN_STATUS_LED = 16;
constexpr uint8_t PIN_BUZZER = 25;

// ------------------------------------------------------------
// UI/timing configuration
// ------------------------------------------------------------
constexpr uint16_t SCREEN_W = 320;
constexpr uint16_t SCREEN_H = 240;

constexpr uint32_t TOUCH_DEBOUNCE_MS = 150;
constexpr uint32_t MESSAGE_SCREEN_MS = 2600;
constexpr uint32_t RFID_RESCAN_GUARD_MS = 1200;
constexpr uint32_t LED_PULSE_MS = 220;

constexpr size_t UID_STR_LEN = 24;
constexpr size_t LOCKER_COUNT = 8;

// ------------------------------------------------------------
// Touch calibration (adjust if orientation differs on your panel)
// ------------------------------------------------------------
constexpr bool TOUCH_SWAP_XY = true;
constexpr bool TOUCH_FLIP_X = false;
constexpr bool TOUCH_FLIP_Y = true;
constexpr int16_t TOUCH_RAW_MIN_X = 200;
constexpr int16_t TOUCH_RAW_MAX_X = 3800;
constexpr int16_t TOUCH_RAW_MIN_Y = 200;
constexpr int16_t TOUCH_RAW_MAX_Y = 3800;

// ------------------------------------------------------------
// Colors
// ------------------------------------------------------------
constexpr uint16_t COLOR_BG = ILI9341_BLACK;
constexpr uint16_t COLOR_TEXT = ILI9341_WHITE;
constexpr uint16_t COLOR_TITLE = ILI9341_CYAN;
constexpr uint16_t COLOR_ASSIGN = 0x03E0;    // Green
constexpr uint16_t COLOR_RETRIEVE = 0x001F;  // Blue
constexpr uint16_t COLOR_CANCEL = 0x7800;    // Maroon
constexpr uint16_t COLOR_OK = ILI9341_GREEN;
constexpr uint16_t COLOR_WARN = ILI9341_YELLOW;
constexpr uint16_t COLOR_ERROR = ILI9341_RED;

// ------------------------------------------------------------
// Types
// ------------------------------------------------------------
struct Rect {
  int16_t x;
  int16_t y;
  int16_t w;
  int16_t h;
};

enum class UiState : uint8_t {
  HOME = 0,
  WAITING_RFID,
  MESSAGE
};

enum class Action : uint8_t {
  NONE = 0,
  ASSIGN,
  RETRIEVE
};

struct Locker {
  bool occupied = false;
  char ownerUid[UID_STR_LEN] = "";
};

struct BuzzerPattern {
  uint8_t beepsRemaining = 0;
  bool toneOn = false;
  uint16_t frequency = 2200;
  uint16_t onMs = 90;
  uint16_t offMs = 70;
  uint32_t nextToggleMs = 0;
};

// ------------------------------------------------------------
// Globals
// ------------------------------------------------------------
Adafruit_ILI9341 tft(PIN_TFT_CS, PIN_TFT_DC, PIN_TFT_RST);
// IRQ can be noisy/floating on some wiring. Using CS-only read is often more reliable.
XPT2046_Touchscreen touch(PIN_TOUCH_CS);
MFRC522 rfid(PIN_RFID_SS, PIN_RFID_RST);

Locker lockers[LOCKER_COUNT];
BuzzerPattern buzzer;

UiState uiState = UiState::HOME;
Action pendingAction = Action::NONE;

bool touchReady = false;
bool touchLatched = false;

uint32_t lastTouchAtMs = 0;
uint32_t messageShownAtMs = 0;
uint32_t ledPulseUntilMs = 0;
uint32_t lastUidSeenAtMs = 0;
char lastUidSeen[UID_STR_LEN] = "";

char messageLine1[48] = "";
char messageLine2[48] = "";
uint16_t messageColor = COLOR_OK;

const Rect BTN_ASSIGN = {20, 60, 280, 50};
const Rect BTN_RETRIEVE = {20, 124, 280, 50};
const Rect BTN_CANCEL = {20, 188, 280, 34};
const Rect BTN_BACK = {20, 188, 280, 34};

// ------------------------------------------------------------
// Declarations
// ------------------------------------------------------------
void initPins();
void initDisplay();
void initTouch();
void initRfid();
void hardResetTft();
void waitMs(uint32_t ms);

void deselectAllSpiDevices();
void prepareForTft();
void prepareForTouch();
void prepareForRfid();

void drawHomeScreen();
void drawRfidPromptScreen();
void drawMessageScreen();
void drawButton(const Rect &r, uint16_t color, const char *label);
bool pointInRect(const Rect &r, int16_t x, int16_t y);
void setMessage(const char *line1, const char *line2, uint16_t color);

bool readTouchTap(int16_t &x, int16_t &y);
void handleTouchInput();
void updateUiTimers();

void pollRfid();
bool readRfidUid(char *uidOut, size_t outSize);
void uidToString(const MFRC522::Uid &uid, char *out, size_t outSize);
void handleScannedCard(const char *uid);

int8_t findLockerByUid(const char *uid);
int8_t findFirstFreeLocker();
int8_t assignLockerToUid(const char *uid);
int8_t retrieveLockerForUid(const char *uid);
uint8_t countFreeLockers();

void startBuzzerPattern(uint8_t beeps, uint16_t frequency, uint16_t onMs, uint16_t offMs);
void updateBuzzer();
void pulseLed(uint32_t ms);
void updateLed();

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

  initPins();
  initDisplay();
  initTouch();
  initRfid();

  drawHomeScreen();
  startBuzzerPattern(1, 2400, 90, 70);

  Serial.println();
  Serial.println("Smart Locker ready");
}

void loop() {
  handleTouchInput();
  pollRfid();
  updateUiTimers();
  updateBuzzer();
  updateLed();
}

// ------------------------------------------------------------
// Initialization
// ------------------------------------------------------------
void initPins() {
  pinMode(PIN_TFT_CS, OUTPUT);
  pinMode(PIN_TOUCH_CS, OUTPUT);
  pinMode(PIN_RFID_SS, OUTPUT);
  pinMode(PIN_TFT_RST, OUTPUT);
  pinMode(PIN_TOUCH_IRQ, INPUT);
  deselectAllSpiDevices();

  pinMode(PIN_STATUS_LED, OUTPUT);
  digitalWrite(PIN_STATUS_LED, LOW);

  pinMode(PIN_BUZZER, OUTPUT);
  noTone(PIN_BUZZER);

  SPI.begin(PIN_SPI_SCK, PIN_SPI_MISO, PIN_SPI_MOSI);
}

void initDisplay() {
  hardResetTft();
  prepareForTft();
  // Lower init SPI frequency improves stability on shared-bus setups.
  tft.begin(10000000);
  tft.setRotation(1);  // 320x240 landscape
  tft.fillScreen(COLOR_BG);
}

void initTouch() {
  prepareForTouch();
  touchReady = touch.begin();
  touch.setRotation(0);  // keep raw; mapping is handled manually

  Serial.print("Touch init: ");
  Serial.println(touchReady ? "OK" : "FAILED");
}

void initRfid() {
  prepareForRfid();
  rfid.PCD_Init();
  Serial.println("RFID init: OK");
}

void hardResetTft() {
  digitalWrite(PIN_TFT_RST, HIGH);
  waitMs(5);
  digitalWrite(PIN_TFT_RST, LOW);
  waitMs(20);
  digitalWrite(PIN_TFT_RST, HIGH);
  waitMs(150);
}

void waitMs(uint32_t ms) {
  const uint32_t start = millis();
  while (millis() - start < ms) {
    yield();
  }
}

// ------------------------------------------------------------
// Shared SPI arbitration (CS control)
// ------------------------------------------------------------
void deselectAllSpiDevices() {
  digitalWrite(PIN_TFT_CS, HIGH);
  digitalWrite(PIN_TOUCH_CS, HIGH);
  digitalWrite(PIN_RFID_SS, HIGH);
}

void prepareForTft() {
  digitalWrite(PIN_TOUCH_CS, HIGH);
  digitalWrite(PIN_RFID_SS, HIGH);
}

void prepareForTouch() {
  digitalWrite(PIN_TFT_CS, HIGH);
  digitalWrite(PIN_RFID_SS, HIGH);
}

void prepareForRfid() {
  digitalWrite(PIN_TFT_CS, HIGH);
  digitalWrite(PIN_TOUCH_CS, HIGH);
}

// ------------------------------------------------------------
// UI rendering
// ------------------------------------------------------------
void drawButton(const Rect &r, uint16_t color, const char *label) {
  prepareForTft();
  tft.fillRoundRect(r.x, r.y, r.w, r.h, 8, color);
  tft.drawRoundRect(r.x, r.y, r.w, r.h, 8, ILI9341_WHITE);
  tft.setTextSize(2);
  tft.setTextColor(ILI9341_WHITE);
  tft.setCursor(r.x + 10, r.y + 16);
  tft.print(label);
}

void drawHomeScreen() {
  uiState = UiState::HOME;
  pendingAction = Action::NONE;

  prepareForTft();
  tft.fillScreen(COLOR_BG);

  tft.setTextSize(2);
  tft.setTextColor(COLOR_TITLE);
  tft.setCursor(16, 16);
  tft.print("Smart Locker");

  tft.setTextSize(1);
  tft.setTextColor(COLOR_TEXT);
  tft.setCursor(196, 22);
  tft.print("Free: ");
  tft.print(countFreeLockers());
  tft.print("/");
  tft.print(LOCKER_COUNT);

  drawButton(BTN_ASSIGN, COLOR_ASSIGN, "Assign Locker");
  drawButton(BTN_RETRIEVE, COLOR_RETRIEVE, "Retrieve");
  drawButton(BTN_CANCEL, COLOR_CANCEL, "Cancel");
}

void drawRfidPromptScreen() {
  uiState = UiState::WAITING_RFID;

  prepareForTft();
  tft.fillScreen(COLOR_BG);

  tft.setTextSize(2);
  tft.setTextColor(COLOR_TITLE);
  tft.setCursor(20, 24);
  tft.print("Scan RFID Card");

  tft.setTextSize(2);
  tft.setTextColor(COLOR_TEXT);
  tft.setCursor(20, 66);
  if (pendingAction == Action::ASSIGN) {
    tft.print("Mode: Assign");
  } else if (pendingAction == Action::RETRIEVE) {
    tft.print("Mode: Retrieve");
  } else {
    tft.print("Mode: None");
  }

  tft.setTextSize(1);
  tft.setTextColor(ILI9341_LIGHTGREY);
  tft.setCursor(20, 98);
  tft.print("Tap card on reader to continue.");

  drawButton(BTN_BACK, COLOR_CANCEL, "Cancel");
}

void drawMessageScreen() {
  uiState = UiState::MESSAGE;
  messageShownAtMs = millis();

  prepareForTft();
  tft.fillScreen(COLOR_BG);

  tft.setTextSize(2);
  tft.setTextColor(messageColor);
  tft.setCursor(20, 44);
  tft.print(messageLine1);

  tft.setTextSize(2);
  tft.setTextColor(COLOR_TEXT);
  tft.setCursor(20, 86);
  tft.print(messageLine2);

  tft.setTextSize(1);
  tft.setTextColor(ILI9341_LIGHTGREY);
  tft.setCursor(20, 132);
  tft.print("Returning to menu...");
}

void setMessage(const char *line1, const char *line2, uint16_t color) {
  snprintf(messageLine1, sizeof(messageLine1), "%s", line1);
  snprintf(messageLine2, sizeof(messageLine2), "%s", line2);
  messageColor = color;
}

bool pointInRect(const Rect &r, int16_t x, int16_t y) {
  return (x >= r.x && x < (r.x + r.w) && y >= r.y && y < (r.y + r.h));
}

// ------------------------------------------------------------
// Touch handling
// ------------------------------------------------------------
bool readTouchTap(int16_t &x, int16_t &y) {
  if (!touchReady) {
    return false;
  }

  prepareForTouch();
  const bool pressed = touch.touched();

  if (!pressed) {
    touchLatched = false;
    return false;
  }

  if (touchLatched) {
    return false;
  }

  const uint32_t now = millis();
  if (now - lastTouchAtMs < TOUCH_DEBOUNCE_MS) {
    return false;
  }

  TS_Point p = touch.getPoint();
  if (p.z < 200) {
    return false;
  }
  int16_t rawX = p.x;
  int16_t rawY = p.y;

  if (TOUCH_SWAP_XY) {
    int16_t temp = rawX;
    rawX = rawY;
    rawY = temp;
  }

  x = map(rawX, TOUCH_RAW_MIN_X, TOUCH_RAW_MAX_X, 0, SCREEN_W - 1);
  y = map(rawY, TOUCH_RAW_MIN_Y, TOUCH_RAW_MAX_Y, 0, SCREEN_H - 1);

  if (TOUCH_FLIP_X) {
    x = (SCREEN_W - 1) - x;
  }
  if (TOUCH_FLIP_Y) {
    y = (SCREEN_H - 1) - y;
  }

  x = constrain(x, 0, SCREEN_W - 1);
  y = constrain(y, 0, SCREEN_H - 1);

  touchLatched = true;
  lastTouchAtMs = now;
  return true;
}

void handleTouchInput() {
  int16_t x = 0;
  int16_t y = 0;
  if (!readTouchTap(x, y)) {
    return;
  }

  Serial.printf("[Touch] x=%d y=%d\n", x, y);

  if (uiState == UiState::HOME) {
    if (pointInRect(BTN_ASSIGN, x, y)) {
      pendingAction = Action::ASSIGN;
      drawRfidPromptScreen();
      return;
    }
    if (pointInRect(BTN_RETRIEVE, x, y)) {
      pendingAction = Action::RETRIEVE;
      drawRfidPromptScreen();
      return;
    }
    if (pointInRect(BTN_CANCEL, x, y)) {
      setMessage("Canceled", "No action selected", COLOR_WARN);
      drawMessageScreen();
      return;
    }
  }

  if (uiState == UiState::WAITING_RFID && pointInRect(BTN_BACK, x, y)) {
    setMessage("Canceled", "RFID flow aborted", COLOR_WARN);
    drawMessageScreen();
    return;
  }
}

void updateUiTimers() {
  if (uiState == UiState::MESSAGE && millis() - messageShownAtMs >= MESSAGE_SCREEN_MS) {
    drawHomeScreen();
  }
}

// ------------------------------------------------------------
// RFID + locker logic integration
// ------------------------------------------------------------
void pollRfid() {
  if (uiState != UiState::WAITING_RFID) {
    return;
  }

  char uid[UID_STR_LEN] = "";
  if (!readRfidUid(uid, sizeof(uid))) {
    return;
  }

  const uint32_t now = millis();
  if ((strncmp(uid, lastUidSeen, UID_STR_LEN) == 0) &&
      (now - lastUidSeenAtMs < RFID_RESCAN_GUARD_MS)) {
    return;
  }

  snprintf(lastUidSeen, sizeof(lastUidSeen), "%s", uid);
  lastUidSeenAtMs = now;

  pulseLed(LED_PULSE_MS);
  startBuzzerPattern(1, 2600, 85, 70);

  Serial.printf("[RFID] UID=%s\n", uid);
  handleScannedCard(uid);
  drawMessageScreen();
}

bool readRfidUid(char *uidOut, size_t outSize) {
  // Explicitly disable TFT + Touch before RFID read (shared SPI safety).
  prepareForRfid();

  if (!rfid.PICC_IsNewCardPresent()) {
    return false;
  }
  if (!rfid.PICC_ReadCardSerial()) {
    return false;
  }

  uidToString(rfid.uid, uidOut, outSize);
  rfid.PICC_HaltA();
  rfid.PCD_StopCrypto1();
  return true;
}

void uidToString(const MFRC522::Uid &uid, char *out, size_t outSize) {
  if (outSize == 0) {
    return;
  }

  size_t idx = 0;
  for (byte i = 0; i < uid.size; i++) {
    if (i > 0) {
      if (idx + 1 >= outSize) break;
      out[idx++] = ':';
    }

    if (idx + 2 >= outSize) break;
    const uint8_t b = uid.uidByte[i];
    const char hex[] = "0123456789ABCDEF";
    out[idx++] = hex[(b >> 4) & 0x0F];
    out[idx++] = hex[b & 0x0F];
  }
  out[idx] = '\0';
}

void handleScannedCard(const char *uid) {
  if (pendingAction == Action::ASSIGN) {
    const int8_t already = findLockerByUid(uid);
    if (already >= 0) {
      char line2[48];
      snprintf(line2, sizeof(line2), "Already Locker #%d", already + 1);
      setMessage("Assign Skipped", line2, COLOR_WARN);
      return;
    }

    const int8_t assigned = assignLockerToUid(uid);
    if (assigned >= 0) {
      char line2[48];
      snprintf(line2, sizeof(line2), "Assigned Locker #%d", assigned + 1);
      setMessage("Assignment Done", line2, COLOR_OK);
    } else {
      setMessage("Assignment Failed", "No free locker", COLOR_ERROR);
      startBuzzerPattern(3, 1300, 75, 75);
    }
    return;
  }

  if (pendingAction == Action::RETRIEVE) {
    const int8_t released = retrieveLockerForUid(uid);
    if (released >= 0) {
      char line2[48];
      snprintf(line2, sizeof(line2), "Freed Locker #%d", released + 1);
      setMessage("Retrieve Done", line2, COLOR_OK);
    } else {
      setMessage("Retrieve Failed", "Card has no locker", COLOR_ERROR);
      startBuzzerPattern(2, 1400, 75, 75);
    }
    return;
  }

  setMessage("No Action", "Select a button first", COLOR_WARN);
}

int8_t findLockerByUid(const char *uid) {
  for (size_t i = 0; i < LOCKER_COUNT; i++) {
    if (!lockers[i].occupied) {
      continue;
    }
    if (strncmp(lockers[i].ownerUid, uid, UID_STR_LEN) == 0) {
      return static_cast<int8_t>(i);
    }
  }
  return -1;
}

int8_t findFirstFreeLocker() {
  for (size_t i = 0; i < LOCKER_COUNT; i++) {
    if (!lockers[i].occupied) {
      return static_cast<int8_t>(i);
    }
  }
  return -1;
}

int8_t assignLockerToUid(const char *uid) {
  const int8_t freeIdx = findFirstFreeLocker();
  if (freeIdx < 0) {
    return -1;
  }

  Locker &locker = lockers[freeIdx];
  locker.occupied = true;
  snprintf(locker.ownerUid, sizeof(locker.ownerUid), "%s", uid);
  return freeIdx;
}

int8_t retrieveLockerForUid(const char *uid) {
  const int8_t idx = findLockerByUid(uid);
  if (idx < 0) {
    return -1;
  }

  Locker &locker = lockers[idx];
  locker.occupied = false;
  locker.ownerUid[0] = '\0';
  return idx;
}

uint8_t countFreeLockers() {
  uint8_t freeCount = 0;
  for (size_t i = 0; i < LOCKER_COUNT; i++) {
    if (!lockers[i].occupied) {
      freeCount++;
    }
  }
  return freeCount;
}

// ------------------------------------------------------------
// Feedback helpers
// ------------------------------------------------------------
void startBuzzerPattern(uint8_t beeps, uint16_t frequency, uint16_t onMs, uint16_t offMs) {
  buzzer.beepsRemaining = beeps;
  buzzer.toneOn = false;
  buzzer.frequency = frequency;
  buzzer.onMs = onMs;
  buzzer.offMs = offMs;
  buzzer.nextToggleMs = 0;
}

void updateBuzzer() {
  if (buzzer.beepsRemaining == 0) {
    if (buzzer.toneOn) {
      noTone(PIN_BUZZER);
      buzzer.toneOn = false;
    }
    return;
  }

  const uint32_t now = millis();
  if (now < buzzer.nextToggleMs) {
    return;
  }

  if (!buzzer.toneOn) {
    tone(PIN_BUZZER, buzzer.frequency);
    buzzer.toneOn = true;
    buzzer.nextToggleMs = now + buzzer.onMs;
  } else {
    noTone(PIN_BUZZER);
    buzzer.toneOn = false;
    buzzer.beepsRemaining--;
    buzzer.nextToggleMs = now + buzzer.offMs;
  }
}

void pulseLed(uint32_t ms) {
  ledPulseUntilMs = millis() + ms;
}

void updateLed() {
  digitalWrite(PIN_STATUS_LED, (millis() < ledPulseUntilMs) ? HIGH : LOW);
}