// CYD Plug Dashboard — 4-button UI for zigbee2mqtt smart plugs
//
// Hardware: Hosyond ESP32-32E 4" (320x480, ST7796, XPT2046 touch)
//   TFT pins: User_Setup.h (vendored from klubhaus-doorbell)
//   Backlight: GPIO 27
//   Touch CS:  GPIO 33, IRQ: GPIO 36
//
// MQTT broker: 192.168.81.147:1883 (set in secrets.h)
//   Subscribe: zigbee2mqtt/<plug>      (state JSON {"state":"ON"})
//   Publish:   zigbee2mqtt/<plug>/set  (TOGGLE)
//
// Layout: 2x2 grid of buttons, status bar at top
//   ┌─────────────────────────────┐
//   │ WiFi✓ MQTT✓ Squiggle● ON    │
//   ├──────────────┬──────────────┤
//   │  Squiggle    │  Sideboard   │
//   │      ON      │      OFF     │
//   ├──────────────┼──────────────┤
//   │ Bamboo Lights│  DJ Booth    │
//   │      ON      │      OFF     │
//   └──────────────┴──────────────┘

#include <Arduino.h>
#include <TFT_eSPI.h>
#include <WiFi.h>
#include <PubSubClient.h>
#include <ArduinoJson.h>
#include <esp_random.h>
#include "secrets.h"

// ── Plugs ──────────────────────────────────────────────────────
struct Plug {
  const char* name;
  const char* topic;
  bool state;
};

static Plug plugs[] = {
  {"Squiggle",  "zigbee2mqtt/Squiggle"},
  {"Sideboard", "zigbee2mqtt/Sideboard Lamp"},
  {"Bamboo",    "zigbee2mqtt/Bamboo Lights"},
  {"DJ Booth",  "zigbee2mqtt/DJ Booth"},
};
static const int PLUG_COUNT = sizeof(plugs) / sizeof(plugs[0]);

// ── Party Lock ─────────────────────────────────────────────────
#define PARTY_LOCK_TOPIC "party-lock/set"
bool partyLocked = false;
bool partyLockPending = false;

// ── Double-tap ────────────────────────────────────────────────
#define DOUBLE_TAP_MS 2000
static int doubleTapIdx = -1;
static uint32_t doubleTapTime = 0;

// ── State query ───────────────────────────────────────────────
static int queryIdx = -1;
static uint32_t queryAt = 0;

// ── Display ────────────────────────────────────────────────────
#define PIN_LCD_BL 27
#define TFT_ROTATION 1   // landscape (480x320)
TFT_eSPI tft = TFT_eSPI();

// ── MQTT ───────────────────────────────────────────────────────
WiFiClient wifi;
PubSubClient mqtt(wifi);

String clientId;
bool mqttConnected = false;
bool wifiConnected = false;

// ── Touch ──────────────────────────────────────────────────────
// TFT_eSPI handles touch via TOUCH_CS / TOUCH_IRQ from User_Setup.h.
// Calibration from cyd-calibrate sketch:
//   raw_x  → screen_x  (linear, both axes increase rightward)
//   raw_y  → screen_y  (inverted — high raw_y = top of screen)
// Range: raw_x 31..454, raw_y 32..281
#define TOUCH_MIN_RAW_X 31
#define TOUCH_MAX_RAW_X 454
#define TOUCH_MIN_RAW_Y 32
#define TOUCH_MAX_RAW_Y 281

// ── Layout ─────────────────────────────────────────────────────
#define STATUS_BAR_H 28
#define BUTTON_GAP 6
#define BUTTON_PAD 6

struct Button {
  int x, y, w, h;
  int plug_index;
};

Button buttons[PLUG_COUNT];

// Screen-off button in the upper-right corner of the status bar.
// "OFF" at textSize(2) needs ~48px; make button 64px wide with padding.
#define SCREEN_OFF_W 64
#define SCREEN_OFF_H STATUS_BAR_H
#define SCREEN_OFF_X (tft.width() - SCREEN_OFF_W - 4)
#define SCREEN_OFF_Y 0

// Party lock toggle in the status bar (between MQTT and OFF).
#define LOCK_BTN_W 64
#define LOCK_BTN_H STATUS_BAR_H
#define LOCK_BTN_X (SCREEN_OFF_X - LOCK_BTN_W - 4)
#define LOCK_BTN_Y 0

bool displayOn = true;

// ── Colors ─────────────────────────────────────────────────────
#define COL_BG        0x0000
#define COL_STATUS_BG 0x18E3   // dark blue
#define COL_BTN_OFF   0x2945   // dark gray
#define COL_BTN_ON    0x07E0   // green
#define COL_BORDER    0x4208
#define COL_TEXT      0xFFFF
#define COL_TEXT_DIM  0xAD55
#define COL_OK        0x07E0   // green dot
#define COL_BAD       0xF800   // red dot
#define COL_TRACE     0x5D5D   // gray trace dot

// ── Forward decls ──────────────────────────────────────────────
void connectWifi();
void connectMqtt();
void mqttCallback(char* topic, byte* payload, unsigned int length);
void publishToggle(int idx);
void drawStatusBar();
void drawButton(int idx);
void drawAll();
int  buttonAt(int tx, int ty);
void layoutButtons();
uint32_t buttonColor(const Plug& p);

// ── Forward decls ──────────────────────────────────────────────
void drawButton(int idx);
void drawAll();
uint32_t buttonColor(const Plug& p);
int  buttonAt(int tx, int ty);
void pushTrace(int x, int y);
void drawTraces();
void pruneTraces();

// ── Trace particles ────────────────────────────────────────────
#define MAX_TRACES 100
#define TRACE_FADE_MS 1000

struct Trace {
  int x, y;
  uint32_t time;
};

static Trace traces[MAX_TRACES];
static int traceHead = 0;
static int traceCount = 0;

static uint32_t lastTraceDraw = 0;

// ── Setup ──────────────────────────────────────────────────────
void setup() {
  Serial.begin(115200);
  delay(200);
  Serial.println();
  Serial.println("=== cyd-dashboard ===");

  // Backlight on
  pinMode(PIN_LCD_BL, OUTPUT);
  digitalWrite(PIN_LCD_BL, HIGH);

  // TFT init
  tft.init();
  tft.setRotation(TFT_ROTATION);
  tft.fillScreen(COL_BG);

  // Show "booting" message
  tft.setTextColor(COL_TEXT, COL_BG);
  tft.setTextSize(2);
  tft.setCursor(10, 10);
  tft.println("booting...");

  layoutButtons();

  // MQTT client id with random suffix so reconnects don't collide
  uint32_t r = esp_random();
  char buf[32];
  snprintf(buf, sizeof(buf), "cyd-%08X", r);
  clientId = buf;
  mqtt.setServer(MQTT_HOST, MQTT_PORT);
  mqtt.setCallback(mqttCallback);
  mqtt.setBufferSize(512);

  connectWifi();

  // Touch self-test: print raw Z to confirm XPT2046 is responding
  uint16_t tz = tft.getTouchRawZ();
  Serial.printf("touch self-test: raw Z=%u (non-zero = controller detected)\n", tz);

  drawAll();
}

// ── Loop ───────────────────────────────────────────────────────
void loop() {
  if (WiFi.status() != WL_CONNECTED) {
    wifiConnected = false;
  // (Power note: WiFi TX on marginal hubs (unpowered USB) can cause brownouts.
  //  If disconnects recur every ~5s, move board to direct USB or powered hub.)
  connectWifi();
  } else {
    wifiConnected = true;
  }

  if (wifiConnected && !mqtt.connected()) {
    connectMqtt();
  }
  mqtt.loop();

  // Touch polling — simple debounced tap.
  // From cyd-calibrate: raw_x → screen_x (linear), raw_y → screen_y (inverted).
  // Clamp because the OFF button sits at sy=0 and user taps above the
  // calibrated max produce negative sy without clamping.
  // touchWasDown: only act on finger-DOWN edge, ignore drag/repeat.
  // lastTraceX/Y: last position where a trace dot was drawn for drag feedback.
  static bool touchWasDown = false;
  static int lastTraceX = -1, lastTraceY = -1;
  static uint32_t last_touch = 0;
  uint16_t tx, ty;
  bool nowTouching = tft.getTouch(&tx, &ty, 100);
  if (nowTouching && !touchWasDown && millis() - last_touch > 250) {
      int sx = constrain(map(tx, TOUCH_MIN_RAW_X, TOUCH_MAX_RAW_X, 0, tft.width()),
                         0, tft.width());
      int sy = constrain(map(ty, TOUCH_MIN_RAW_Y, TOUCH_MAX_RAW_Y, tft.height(), 0),
                         0, tft.height());

      // Zone classification (for debug + action)
      const char* zone = "unknown";
      if (!displayOn) {
        zone = "WAKE";
      } else if (sy < STATUS_BAR_H) {
        if (sx >= SCREEN_OFF_X) zone = "OFF-btn";
        else if (sx >= LOCK_BTN_X) zone = "lock-btn";
        else zone = "status-bar";
      } else {
        int idx = buttonAt(sx, sy);
        if (idx >= 0) zone = plugs[idx].name;
        else zone = "between-buttons";
      }
      Serial.printf("touch raw=(%u,%u) screen=(%d,%d) zone=%s (off_btn=%d..%d)\n",
                    tx, ty, sx, sy, zone, SCREEN_OFF_X, SCREEN_OFF_X + SCREEN_OFF_W);

      if (!displayOn) {
        digitalWrite(PIN_LCD_BL, HIGH);
        displayOn = true;
        Serial.println("display: on (wake)");
        last_touch = millis();
        return;
      }

      // Party lock toggle: in the status bar between MQTT and OFF.
      if (sy < STATUS_BAR_H &&
          sx >= LOCK_BTN_X && sx < LOCK_BTN_X + LOCK_BTN_W) {
        doubleTapIdx = -1;
        partyLocked = !partyLocked;
        mqtt.publish(PARTY_LOCK_TOPIC, partyLocked ? "ON" : "OFF");
        drawStatusBar();
        Serial.printf("party-lock: %s\n", partyLocked ? "ON" : "OFF");
        last_touch = millis();
        return;
      }

      // Screen-off button: any tap in the upper-right of the status bar counts,
      // even if sy would extrapolate slightly negative without clamping.
      if (sy < STATUS_BAR_H &&
          sx >= SCREEN_OFF_X && sx < SCREEN_OFF_X + SCREEN_OFF_W) {
        digitalWrite(PIN_LCD_BL, LOW);
        displayOn = false;
        Serial.println("display: off");
        last_touch = millis();
        return;
      }

      int idx = buttonAt(sx, sy);
      if (idx >= 0) {
        char dbg[64];
        snprintf(dbg, sizeof(dbg), "tap %s at (%d,%d) locked=%d",
                 plugs[idx].name, sx, sy, partyLocked);
        mqtt.publish("cyd-debug/touch", dbg);
        if (partyLocked) {
          tft.fillRoundRect(LOCK_BTN_X, LOCK_BTN_Y, LOCK_BTN_W, LOCK_BTN_H, 4, COL_BAD);
          tft.setTextColor(COL_TEXT, COL_BAD);
          tft.setTextSize(2);
          tft.setCursor(LOCK_BTN_X + 6, 4);
          tft.print("LOCK");
          delay(150);
          drawStatusBar();
          } else if (idx == 3) {
          // DJ Booth requires double-tap
          uint32_t now = millis();
          if (doubleTapIdx == 3 && now - doubleTapTime < DOUBLE_TAP_MS) {
            doubleTapIdx = -1;
            publishToggle(idx);
          } else {
            doubleTapIdx = 3;
            doubleTapTime = now;
            drawButton(idx);
            tft.setTextColor(COL_TEXT_DIM, COL_BTN_OFF);
            tft.setTextSize(1);
            tft.setCursor(buttons[idx].x + 4, buttons[idx].y + buttons[idx].h - 14);
            tft.print("tap again");
          }
        } else {
          doubleTapIdx = -1;
          publishToggle(idx);
        }
        last_touch = millis();
      }
      touchWasDown = true;
  } else if (nowTouching && touchWasDown) {
    int sx = constrain(map(tx, TOUCH_MIN_RAW_X, TOUCH_MAX_RAW_X, 0, tft.width()), 0, tft.width());
    int sy = constrain(map(ty, TOUCH_MIN_RAW_Y, TOUCH_MAX_RAW_Y, tft.height(), 0), 0, tft.height());
    pushTrace(sx, sy);
    touchWasDown = true;
  } else if (!nowTouching) {
    touchWasDown = false;
  }
  // Trace particle upkeep: prune + draw every 30ms
  if (millis() - lastTraceDraw > 30) {
    pruneTraces();
    drawTraces();
    lastTraceDraw = millis();
  }

  // Double-tap timeout: reset if user didn't tap again in time
  if (doubleTapIdx >= 0 && millis() - doubleTapTime >= DOUBLE_TAP_MS) {
    doubleTapIdx = -1;
    if (PLUG_COUNT > 3) drawButton(3);
  }

  // Delayed state query via /get to bypass party-lock ACL
  if (queryIdx >= 0 && millis() >= queryAt) {
    char get_topic[64];
    snprintf(get_topic, sizeof(get_topic), "%s/get", plugs[queryIdx].topic);
    mqtt.publish(get_topic, "{\"state\":\"\"}");
    Serial.printf("queried state for %s via /get\n", plugs[queryIdx].name);
    queryIdx = -1;
  }
}

// ── Trace particles ────────────────────────────────────────────
int buttonAt(int tx, int ty);

void pushTrace(int x, int y) {
  if (buttonAt(x, y) >= 0) return;
  traces[traceHead] = {x, y, millis()};
  traceHead = (traceHead + 1) % MAX_TRACES;
  if (traceCount < MAX_TRACES) traceCount++;
}

void drawTraces() {
  uint32_t now = millis();
  for (int i = 0; i < traceCount; i++) {
    int idx = (traceHead - traceCount + i + MAX_TRACES) % MAX_TRACES;
    uint32_t age = now - traces[idx].time;
    if (age > TRACE_FADE_MS) continue;
    int bright = 31 - (31 * age / TRACE_FADE_MS);
    if (bright < 0) bright = 0;
    uint16_t color = tft.color565(bright, bright, bright);
    tft.fillCircle(traces[idx].x, traces[idx].y, 3, color);
  }
}

void pruneTraces() {
  uint32_t now = millis();
  int prune = 0;
  while (prune < traceCount) {
    int idx = (traceHead - traceCount + prune + MAX_TRACES) % MAX_TRACES;
    if (now - traces[idx].time <= TRACE_FADE_MS) break;
    tft.fillCircle(traces[idx].x, traces[idx].y, 3, COL_BG);
    prune++;
  }
  if (prune > 0) {
    traceCount -= prune;
  }
}

// ── WiFi ───────────────────────────────────────────────────────
void connectWifi() {
  if (WiFi.status() == WL_CONNECTED) {
    if (!wifiConnected) {
      Serial.printf("WiFi: connected to %s, IP=%s\n",
                    WiFi.SSID().c_str(), WiFi.localIP().toString().c_str());
      wifiConnected = true;
      drawStatusBar();
    }
    return;
  }

  static int ssid_idx = 0;
  static uint32_t last_attempt = 0;
  static uint32_t first_try_at = 0;
  if (millis() - last_attempt < 5000) return;  // back off between attempts

  // If we've been trying the same SSID for >15s without success, give up and move on
  if (first_try_at == 0) first_try_at = millis();
  if (millis() - first_try_at > 15000) {
    Serial.printf("WiFi: giving up on %s after 15s\n", WIFI_SSIDS[ssid_idx]);
    WiFi.disconnect();
    ssid_idx = (ssid_idx + 1) % WIFI_NETWORK_COUNT;
    first_try_at = millis();
    last_attempt = millis();
    return;
  }

  Serial.printf("WiFi: trying %s\n", WIFI_SSIDS[ssid_idx]);
  WiFi.mode(WIFI_STA);
  WiFi.begin(WIFI_SSIDS[ssid_idx], WIFI_PASSPHRASES[ssid_idx]);

  last_attempt = millis();
}

// ── MQTT ───────────────────────────────────────────────────────
void connectMqtt() {
  if (mqtt.connected()) return;
  Serial.printf("MQTT: connecting as %s\n", clientId.c_str());
  if (mqtt.connect(clientId.c_str())) {
    Serial.println("MQTT: connected");
    mqttConnected = true;
    // Subscribe to all plug state topics
    for (int i = 0; i < PLUG_COUNT; i++) {
      mqtt.subscribe(plugs[i].topic);
    }
    mqtt.subscribe(PARTY_LOCK_TOPIC);
    mqtt.publish(PARTY_LOCK_TOPIC, partyLocked ? "ON" : "OFF");
    // Query current state from each plug via /get to bypass party-lock ACL.
    for (int i = 0; i < PLUG_COUNT; i++) {
      char get_topic[64];
      snprintf(get_topic, sizeof(get_topic), "%s/get", plugs[i].topic);
      mqtt.publish(get_topic, "{\"state\":\"\"}");
    }
  } else {
    Serial.printf("MQTT: failed rc=%d\n", mqtt.state());
    mqttConnected = false;
  }
  drawStatusBar();
}

void mqttCallback(char* topic, byte* payload, unsigned int length) {
  // Party lock state
  if (strcmp(topic, PARTY_LOCK_TOPIC) == 0) {
    char s[4];
    int len = min((unsigned int)3, length);
    memcpy(s, payload, len);
    s[len] = 0;
    bool new_state = (strcmp(s, "ON") == 0);
    if (partyLocked != new_state) {
      partyLocked = new_state;
      drawStatusBar();
    }
    return;
  }

  // Match topic to a plug
  int idx = -1;
  for (int i = 0; i < PLUG_COUNT; i++) {
    if (strcmp(topic, plugs[i].topic) == 0) { idx = i; break; }
  }
  if (idx < 0) return;

  // Parse JSON, look for "state"
  StaticJsonDocument<256> doc;
  DeserializationError err = deserializeJson(doc, payload, length);
  if (err) {
    Serial.printf("JSON parse err on %s: %s\n", topic, err.c_str());
    return;
  }
  const char* state = doc["state"] | "";
  bool new_state = (strcmp(state, "ON") == 0);
  if (plugs[idx].state != new_state) {
    Serial.printf("state update: %s -> %s\n", plugs[idx].name, state);
    plugs[idx].state = new_state;
  }
  drawButton(idx);
}

void publishToggle(int idx) {
  if (!mqtt.connected()) return;
  char topic[64];
  snprintf(topic, sizeof(topic), "%s/set", plugs[idx].topic);
  mqtt.publish(topic, "{\"state\":\"TOGGLE\"}");
  Serial.printf("published TOGGLE on %s\n", topic);
  // Schedule a state query 600ms later so display converges to real state
  queryIdx = idx;
  queryAt = millis() + 600;
}

// ── Drawing ────────────────────────────────────────────────────
void layoutButtons() {
  // 2x2 grid below status bar
  int grid_top = STATUS_BAR_H + BUTTON_PAD;
  int grid_h = tft.height() - grid_top - BUTTON_PAD;
  int grid_w = tft.width() - 2 * BUTTON_PAD;
  int btn_w = (grid_w - BUTTON_GAP) / 2;
  int btn_h = (grid_h - BUTTON_GAP) / 2;

  for (int i = 0; i < PLUG_COUNT; i++) {
    int col = i % 2;
    int row = i / 2;
    buttons[i] = {
      BUTTON_PAD + col * (btn_w + BUTTON_GAP),
      grid_top + row * (btn_h + BUTTON_GAP),
      btn_w, btn_h,
      i,
    };
  }
}

void drawStatusBar() {
  tft.fillRect(0, 0, tft.width(), STATUS_BAR_H, COL_STATUS_BG);
  tft.setTextSize(2);
  tft.setTextColor(COL_TEXT, COL_STATUS_BG);

  // WiFi dot
  tft.fillCircle(10, STATUS_BAR_H / 2, 4, wifiConnected ? COL_OK : COL_BAD);
  tft.setCursor(20, 6);
  tft.print(WiFi.localIP());

  // MQTT dot (left of screen-off button)
  int mqtt_x = tft.width() / 2 - 20;
  tft.fillCircle(mqtt_x, STATUS_BAR_H / 2, 4, mqttConnected ? COL_OK : COL_BAD);
  tft.setCursor(mqtt_x + 10, 6);
  tft.print("MQTT");

  // Party lock toggle (between MQTT and OFF)
  uint32_t lock_bg = partyLocked ? COL_BAD : COL_STATUS_BG;
  tft.fillRoundRect(LOCK_BTN_X, LOCK_BTN_Y, LOCK_BTN_W, LOCK_BTN_H, 4, lock_bg);
  tft.drawRoundRect(LOCK_BTN_X, LOCK_BTN_Y, LOCK_BTN_W, LOCK_BTN_H, 4, COL_BORDER);
  tft.setTextColor(COL_TEXT, lock_bg);
  tft.setTextSize(2);
  tft.setCursor(LOCK_BTN_X + 6, 4);
  tft.print(partyLocked ? "LOCK" : "UNLK");

  // Screen-off button (upper-right corner)
  tft.fillRoundRect(SCREEN_OFF_X, SCREEN_OFF_Y, SCREEN_OFF_W, SCREEN_OFF_H, 4, COL_BORDER);
  tft.setTextColor(COL_TEXT, COL_BORDER);
  tft.setCursor(SCREEN_OFF_X + 10, 4);
  tft.print("OFF");
}

void drawButton(int idx) {
  const Plug& p = plugs[idx];
  const Button& b = buttons[idx];

  uint32_t bg = buttonColor(p);
  tft.fillRoundRect(b.x, b.y, b.w, b.h, 8, bg);
  tft.drawRoundRect(b.x, b.y, b.w, b.h, 8, COL_BORDER);

  // Plug name (top) — textSize(2) keeps "Bamboo Lights" within ~237px button
  tft.setTextColor(COL_TEXT, bg);
  tft.setTextSize(2);
  tft.setCursor(b.x + 12, b.y + 10);
  tft.println(p.name);

  // State (bottom)
  tft.setTextSize(4);
  int state_y = b.y + b.h - 40;
  tft.setCursor(b.x + 12, state_y);
  tft.println(p.state ? "ON" : "OFF");
}

void drawAll() {
  drawStatusBar();
  for (int i = 0; i < PLUG_COUNT; i++) drawButton(i);
}

uint32_t buttonColor(const Plug& p) {
  if (p.state) return COL_BTN_ON;
  return COL_BTN_OFF;
}

int buttonAt(int tx, int ty) {
  if (ty < STATUS_BAR_H) return -1;
  for (int i = 0; i < PLUG_COUNT; i++) {
    const Button& b = buttons[i];
    if (tx >= b.x && tx < b.x + b.w && ty >= b.y && ty < b.y + b.h) return i;
  }
  return -1;
}