// AWL Console — Networked v5
// XIAO ESP32C3 quiz console: WiFi + WebSocket + OLED + I2S audio + 5 buttons
//
// Per-team build: change TEAM_ID below to match the physical controller color.
//   Red    = P1
//   Yellow = P2
//   Blue   = P3
//
// Boot flow:
//   1. Boot screen "AWL READY?"  — player presses HAND to join
//   2. Then "TEAM: <color> / Ready" until a question opens
//   3. Press HAND during QUESTION_OPEN to buzz, then A/B/C/D to answer
//   4. Answer letter stays on OLED until result is shown (trophy / crying face)
//
// Library needed (install via Library Manager):
//   ArduinoWebsockets by Gil Maimon

#include <Wire.h>
#include <U8g2lib.h>
#include <driver/i2s.h>
#include <WiFi.h>
#include <ESPmDNS.h>
#include <ArduinoWebsockets.h>

using namespace websockets;

// ── PER-DEVICE CONFIG ────────────────────────────────
// Change ONLY this line when flashing each controller:
#define TEAM_ID   "blue"        // "red" | "yellow" | "blue"

// ── NETWORK CONFIG ───────────────────────────────────
const char* WIFI_SSID = "X.factory2.4G";
const char* WIFI_PASS = "make0314";
// The Mac's IP changes from day to day. Instead of hardcoding it, we use mDNS:
// at runtime we ask the network for the IP of this hostname.
const char* WS_HOST   = "Emilys-MacBook-Pro";   // Mac's mDNS name (without .local)
const int   WS_PORT   = 8080;
String      ws_url    = "";                      // built after mDNS resolves

// ── Pins ─────────────────────────────────────────────
#define LED_PIN          D3
#define I2C_SDA          D4
#define I2C_SCL          D5
#define I2S_DOUT         D0
#define I2S_BCLK         D1
#define I2S_LRC          D2
#define BTN_A            D6
#define BTN_B            D7
#define BTN_RAISE_HAND   D8
#define BTN_C            D9
#define BTN_D            D10

// ── OLED ─────────────────────────────────────────────
U8G2_SSD1309_128X64_NONAME0_F_HW_I2C u8g2(U8G2_R0, U8X8_PIN_NONE);

// ── WebSocket client ─────────────────────────────────
WebsocketsClient wsClient;
bool wsConnected = false;
unsigned long lastReconnectAttempt = 0;
const unsigned long RECONNECT_INTERVAL_MS = 3000;

// ── Game state (local to this controller) ───────────
bool ready = false;          // set to true after player presses HAND on boot
char lastChoice = 0;         // 'A'/'B'/'C'/'D' if we sent an answer, else 0

// ── Audio ─────────────────────────────────────────────
const int SAMPLE_RATE = 22050;
const int AMPLITUDE   = 4000;

void setupI2S() {
  i2s_config_t cfg = {
    .mode                 = (i2s_mode_t)(I2S_MODE_MASTER | I2S_MODE_TX),
    .sample_rate          = SAMPLE_RATE,
    .bits_per_sample      = I2S_BITS_PER_SAMPLE_16BIT,
    .channel_format       = I2S_CHANNEL_FMT_RIGHT_LEFT,
    .communication_format = I2S_COMM_FORMAT_STAND_I2S,
    .intr_alloc_flags     = 0,
    .dma_buf_count        = 8,
    .dma_buf_len          = 64,
    .use_apll             = false
  };
  i2s_pin_config_t pins = {
    .bck_io_num   = I2S_BCLK,
    .ws_io_num    = I2S_LRC,
    .data_out_num = I2S_DOUT,
    .data_in_num  = I2S_PIN_NO_CHANGE
  };
  i2s_driver_install(I2S_NUM_0, &cfg, 0, NULL);
  i2s_set_pin(I2S_NUM_0, &pins);
  i2s_zero_dma_buffer(I2S_NUM_0);
}

void playTone(uint32_t freq_hz, uint32_t duration_ms) {
  uint32_t totalSamples = (SAMPLE_RATE * duration_ms) / 1000;
  if (freq_hz == 0) {
    int16_t silence[2] = {0, 0};
    for (uint32_t i = 0; i < totalSamples; i++) {
      size_t bw;
      i2s_write(I2S_NUM_0, silence, sizeof(silence), &bw, portMAX_DELAY);
    }
    return;
  }
  uint32_t halfPeriod = SAMPLE_RATE / (freq_hz * 2);
  if (halfPeriod < 1) halfPeriod = 1;
  bool     high       = true;
  uint32_t counter    = 0;
  for (uint32_t i = 0; i < totalSamples; i++) {
    int16_t s        = high ? AMPLITUDE : -AMPLITUDE;
    int16_t stereo[2] = {s, s};
    size_t bw;
    i2s_write(I2S_NUM_0, stereo, sizeof(stereo), &bw, portMAX_DELAY);
    if (++counter >= halfPeriod) { high = !high; counter = 0; }
  }
}

void flushAudio() {
  int16_t silence[2] = {0, 0};
  for (int i = 0; i < 512; i++) {
    size_t bw;
    i2s_write(I2S_NUM_0, silence, sizeof(silence), &bw, portMAX_DELAY);
  }
  i2s_zero_dma_buffer(I2S_NUM_0);
}

// ── Sound effects ─────────────────────────────────────
void soundRaiseHand() { playTone(1800, 80); playTone(0, 30); playTone(2400, 120); flushAudio(); }
void soundA()         { playTone(1500, 100); flushAudio(); }
void soundB()         { playTone(1800, 100); flushAudio(); }
void soundC()         { playTone(2100, 100); flushAudio(); }
void soundD()         { playTone(2400, 100); flushAudio(); }
void soundCorrect()   { playTone(1500, 80); playTone(2000, 80); playTone(2500, 150); flushAudio(); }
void soundWrong()     { playTone(800, 100); playTone(600, 200); flushAudio(); }
void soundStartup()   { playTone(1200, 60); playTone(0, 30); playTone(1600, 60); playTone(0, 30); playTone(2000, 100); flushAudio(); }
void soundConnected() { playTone(2000, 60); playTone(2400, 60); playTone(2800, 100); flushAudio(); }
void soundJoin()      { playTone(1000, 80); playTone(1500, 80); playTone(2000, 80); playTone(2500, 150); flushAudio(); }

// ── Display ───────────────────────────────────────────
void showOnOLED(const char* text) {
  u8g2.clearBuffer();
  u8g2.setFont(u8g2_font_logisoso32_tr);
  u8g2.drawStr((128 - u8g2.getStrWidth(text)) / 2, 48, text);
  u8g2.sendBuffer();
}

void showStatus(const char* line1, const char* line2) {
  u8g2.clearBuffer();
  u8g2.setFont(u8g2_font_helvB12_tr);   // Helvetica Bold (closest built-in to Karla)
  u8g2.drawStr((128 - u8g2.getStrWidth(line1)) / 2, 25, line1);
  u8g2.drawStr((128 - u8g2.getStrWidth(line2)) / 2, 50, line2);
  u8g2.sendBuffer();
}

// Uppercase helper — returns the TEAM_ID like "RED" instead of "red"
String teamUpper() {
  String s = String(TEAM_ID);
  s.toUpperCase();
  return s;
}

// "AWL READY?" boot screen — shown until player presses HAND to join
void showReadyScreen() {
  u8g2.clearBuffer();
  u8g2.setFont(u8g2_font_logisoso24_tr);
  const char* big = "AWL";
  u8g2.drawStr((128 - u8g2.getStrWidth(big)) / 2, 32, big);
  u8g2.setFont(u8g2_font_helvB12_tr);
  const char* sub = "READY?";
  u8g2.drawStr((128 - u8g2.getStrWidth(sub)) / 2, 56, sub);
  u8g2.sendBuffer();
}

void showIdle() {
  // If player hasn't joined yet, show the READY? screen instead
  if (!ready) { showReadyScreen(); return; }

  u8g2.clearBuffer();
  u8g2.setFont(u8g2_font_helvB12_tr);
  String header = String("TEAM: ") + teamUpper();   // → "TEAM: RED"
  u8g2.drawStr((128 - u8g2.getStrWidth(header.c_str())) / 2, 25, header.c_str());
  const char* sub = wsConnected ? "READY" : "CONNECTING...";
  u8g2.drawStr((128 - u8g2.getStrWidth(sub)) / 2, 50, sub);
  u8g2.sendBuffer();
}

// ── Happy face (per-round CORRECT answer) ────────────
void showHappyFace() {
  u8g2.clearBuffer();

  // "YEAY!" label at top
  u8g2.setFont(u8g2_font_helvB18_tr);
  const char* label = "YEAY!";
  u8g2.drawStr((128 - u8g2.getStrWidth(label)) / 2, 18, label);

  int cx = 64;
  int cy = 44;
  int r  = 16;

  // Face circle (double line for emphasis)
  u8g2.drawCircle(cx, cy, r);
  u8g2.drawCircle(cx, cy, r - 1);

  // Eyes — small filled dots
  u8g2.drawDisc(cx - 6, cy - 4, 2);
  u8g2.drawDisc(cx + 6, cy - 4, 2);

  // Big smile — lower half of a small circle, opens upward
  u8g2.drawCircle(cx, cy + 2, 8, U8G2_DRAW_LOWER_LEFT | U8G2_DRAW_LOWER_RIGHT);
  u8g2.drawCircle(cx, cy + 2, 7, U8G2_DRAW_LOWER_LEFT | U8G2_DRAW_LOWER_RIGHT);

  u8g2.sendBuffer();
}

// ── #1 + Trophy (END OF GAME winner) ─────────────────
void showWinnerFinal() {
  u8g2.clearBuffer();

  // "#1 WINNER!" label at top
  u8g2.setFont(u8g2_font_helvB12_tr);
  const char* label = "#1 WINNER!";
  u8g2.drawStr((128 - u8g2.getStrWidth(label)) / 2, 13, label);

  int cx = 64;

  // Cup body (rectangle frame so we can put "1" inside)
  u8g2.drawFrame(cx - 14, 18, 28, 20);
  // "1" inside the cup
  u8g2.setFont(u8g2_font_helvB14_tr);
  const char* one = "1";
  u8g2.drawStr(cx - u8g2.getStrWidth(one) / 2, 34, one);

  // Cup bottom edge
  u8g2.drawBox(cx - 16, 38, 32, 3);
  // Side handles
  u8g2.drawCircle(cx - 18, 25, 5);
  u8g2.drawCircle(cx + 18, 25, 5);
  // Stem
  u8g2.drawBox(cx - 3, 41, 6, 7);
  // Two-tier base
  u8g2.drawBox(cx - 10, 48, 20, 4);
  u8g2.drawBox(cx - 14, 52, 28, 4);

  u8g2.sendBuffer();
}

// ── Crying face graphic (for loser) ──────────────────
void showCryingFace() {
  u8g2.clearBuffer();

  // "TOO BAD" label at top
  u8g2.setFont(u8g2_font_helvB12_tr);
  const char* label = "TOO BAD";
  u8g2.drawStr((128 - u8g2.getStrWidth(label)) / 2, 12, label);

  int cx = 64;
  int cy = 42;
  int r  = 18;

  // Face circle outline
  u8g2.drawCircle(cx, cy, r);
  u8g2.drawCircle(cx, cy, r - 1);  // double line for emphasis

  // Eyes - X marks (closed crying eyes)
  u8g2.drawLine(cx - 9, cy - 6, cx - 5, cy - 2);
  u8g2.drawLine(cx - 9, cy - 2, cx - 5, cy - 6);
  u8g2.drawLine(cx + 5, cy - 6, cx + 9, cy - 2);
  u8g2.drawLine(cx + 5, cy - 2, cx + 9, cy - 6);

  // Frown (upper half of small circle, so the curve opens downward)
  u8g2.drawCircle(cx, cy + 12, 6, U8G2_DRAW_UPPER_LEFT | U8G2_DRAW_UPPER_RIGHT);

  // Tear drops
  u8g2.drawDisc(cx - 9, cy + 4, 2);
  u8g2.drawDisc(cx + 9, cy + 4, 2);
  // Tear streams
  u8g2.drawLine(cx - 9, cy + 2, cx - 9, cy + 7);
  u8g2.drawLine(cx + 9, cy + 2, cx + 9, cy + 7);

  u8g2.sendBuffer();
}

// ── WebSocket send ────────────────────────────────────
void sendEvent(const char* event, const char* choice = nullptr) {
  if (!wsConnected) {
    Serial.println("[ws] not connected, skipping send");
    return;
  }
  String msg = String("{\"team\":\"") + TEAM_ID +
               "\",\"event\":\"" + event + "\"";
  if (choice) {
    msg += String(",\"choice\":\"") + choice + "\"";
  }
  msg += String(",\"t\":") + millis() + "}";
  wsClient.send(msg);
  Serial.print("[ws->] "); Serial.println(msg);
}

// ── WebSocket message handler ────────────────────────
void onMessage(WebsocketsMessage message) {
  String data = message.data();
  Serial.print("[ws<-] "); Serial.println(data);

  if (data.indexOf("\"type\":\"lockout\"") >= 0) {
    if (data.indexOf(String("\"team\":\"") + TEAM_ID + "\"") >= 0) {
      // We won the buzz
      digitalWrite(LED_PIN, HIGH);
      showOnOLED("GO!");
    } else {
      showStatus("Buzzed:", "other team");
    }
  }
  else if (data.indexOf("\"type\":\"result\"") >= 0) {
    bool wasUs = data.indexOf(String("\"team\":\"") + TEAM_ID + "\"") >= 0;
    bool correct = data.indexOf("\"correct\":true") >= 0;
    if (wasUs) {
      if (correct) { showHappyFace();  soundCorrect(); }   // "YEAY!" + smiley
      else         { showCryingFace(); soundWrong();   }   // tears
      delay(2500);
    } else {
      showStatus("OTHER TEAM", correct ? "SCORED" : "MISSED");
      delay(800);
    }
    digitalWrite(LED_PIN, LOW);
    lastChoice = 0;
    showIdle();
  }
  else if (data.indexOf("\"type\":\"gameover\"") >= 0) {
    // Server tells us who won the whole game
    bool weWon = data.indexOf(String("\"winner\":\"") + TEAM_ID + "\"") >= 0;
    if (weWon) {
      showWinnerFinal();
      soundCorrect();
      delay(150);
      soundCorrect();   // double celebration jingle
    } else {
      showCryingFace();
      soundWrong();
    }
    // Stay on this screen until host presses 'r' (reset) or device is power-cycled.
  }
  else if (data.indexOf("\"type\":\"reset\"") >= 0) {
    digitalWrite(LED_PIN, LOW);
    lastChoice = 0;
    showIdle();
  }
}

void onEvent(WebsocketsEvent event, String data) {
  if (event == WebsocketsEvent::ConnectionOpened) {
    wsConnected = true;
    Serial.println("[ws] connected");
    soundConnected();
    showIdle();
    sendEvent("hello");
  } else if (event == WebsocketsEvent::ConnectionClosed) {
    wsConnected = false;
    Serial.println("[ws] disconnected");
    showIdle();
  }
}

// Resolve the Mac's hostname (Emilys-MacBook-Pro.local) to an IP via mDNS.
// Returns true if we got an IP and built ws_url.
bool resolveServer() {
  Serial.print("[mdns] resolving "); Serial.print(WS_HOST); Serial.println(".local ...");
  showStatus("Finding", "server...");
  IPAddress ip = MDNS.queryHost(WS_HOST, 3000);   // 3-second timeout
  if (ip == IPAddress(0, 0, 0, 0)) {
    Serial.println("[mdns] resolve FAILED");
    return false;
  }
  ws_url = "ws://" + ip.toString() + ":" + String(WS_PORT);
  Serial.print("[mdns] resolved -> "); Serial.println(ws_url);
  return true;
}

void connectWebSocket() {
  // Make sure we have an IP for the server. If not, resolve via mDNS.
  if (ws_url.length() == 0) {
    if (!resolveServer()) return;   // try again next reconnect tick
  }
  Serial.print("[ws] connecting to "); Serial.println(ws_url);
  showStatus("Connecting", "to server...");
  wsClient.onMessage(onMessage);
  wsClient.onEvent(onEvent);
  bool ok = wsClient.connect(ws_url);
  if (!ok) {
    Serial.println("[ws] connect failed - clearing URL to re-resolve next time");
    ws_url = "";   // force re-resolution (Mac's IP may have changed)
  }
}

// ── Edge detection ────────────────────────────────────
bool lastA = HIGH, lastB = HIGH, lastH = HIGH, lastC = HIGH, lastD = HIGH;

bool justPressed(bool curr, bool& last) {
  bool fired = (last == HIGH && curr == LOW);
  last = curr;
  return fired;
}

// ── Button reaction (local feedback + network send) ──
// persist=true means the OLED label STAYS until a result/reset message clears it.
void buttonReact(const char* label, void (*soundFn)(),
                 const char* eventName, const char* choice = nullptr,
                 bool persist = false) {
  digitalWrite(LED_PIN, HIGH);
  showOnOLED(label);
  Serial.print("Button: "); Serial.println(label);
  soundFn();
  sendEvent(eventName, choice);
  if (!persist) {
    digitalWrite(LED_PIN, LOW);
    showIdle();
  }
  // If persist=true: leave LED on + label showing until the server replies.
}

// ── Handle a HAND press (context-dependent) ──────────
void handleHandPress() {
  if (!ready) {
    // First press = "join". Don't send a buzz to the server, just go to Ready state.
    ready = true;
    Serial.println("[ready] player joined");
    showStatus("WELCOME", String(String("TEAM: ") + teamUpper()).c_str());
    soundJoin();
    delay(800);
    showIdle();
  } else {
    // Real buzz
    buttonReact("HAND", soundRaiseHand, "hand");
  }
}

// ── Setup ─────────────────────────────────────────────
void setup() {
  Serial.begin(115200);
  delay(500);
  Serial.println("\nAWL Console v5 — TEAM: " TEAM_ID);

  pinMode(LED_PIN,        OUTPUT);
  digitalWrite(LED_PIN,   LOW);
  pinMode(BTN_A,          INPUT_PULLUP);
  pinMode(BTN_B,          INPUT_PULLUP);
  pinMode(BTN_RAISE_HAND, INPUT_PULLUP);
  pinMode(BTN_C,          INPUT_PULLUP);
  pinMode(BTN_D,          INPUT_PULLUP);

  Wire.begin(I2C_SDA, I2C_SCL);
  u8g2.begin();
  setupI2S();

  showStatus("AWL", "Booting...");
  soundStartup();

  // WiFi
  Serial.print("[wifi] connecting to "); Serial.println(WIFI_SSID);
  showStatus("WiFi", "connecting...");
  WiFi.mode(WIFI_STA);
  WiFi.begin(WIFI_SSID, WIFI_PASS);
  unsigned long t0 = millis();
  while (WiFi.status() != WL_CONNECTED && millis() - t0 < 15000) {
    delay(250);
    Serial.print(".");
  }
  Serial.println();
  if (WiFi.status() == WL_CONNECTED) {
    Serial.print("[wifi] OK, ip="); Serial.println(WiFi.localIP());
    showStatus("WiFi OK", WiFi.localIP().toString().c_str());
    delay(800);
  } else {
    Serial.println("[wifi] FAILED (will keep retrying)");
    showStatus("WiFi", "FAILED");
  }

  // Start mDNS so we can resolve "Emilys-MacBook-Pro.local" to an IP.
  // The name we pass here is what THIS device would advertise — not what we're querying.
  if (!MDNS.begin("awl-client")) {
    Serial.println("[mdns] start FAILED");
  } else {
    Serial.println("[mdns] started");
  }

  connectWebSocket();
  // After WS connect, onEvent() will fire and call showIdle() — which now
  // shows "AWL READY?" because ready==false. Player must press HAND to join.
  showIdle();
}

// ── Loop ──────────────────────────────────────────────
void loop() {
  // Maintain WiFi
  if (WiFi.status() != WL_CONNECTED) {
    wsConnected = false;
    showStatus("WiFi", "reconnecting");
    WiFi.reconnect();
    delay(500);
    return;
  }

  // Maintain WebSocket
  if (wsConnected) {
    wsClient.poll();
  } else if (millis() - lastReconnectAttempt > RECONNECT_INTERVAL_MS) {
    lastReconnectAttempt = millis();
    connectWebSocket();
  }

  // Buttons
  bool a = digitalRead(BTN_A);
  bool b = digitalRead(BTN_B);
  bool h = digitalRead(BTN_RAISE_HAND);
  bool c = digitalRead(BTN_C);
  bool d = digitalRead(BTN_D);

  if (justPressed(h, lastH)) handleHandPress();

  // Answer buttons only respond after player has joined (ready==true)
  if (ready) {
    if (justPressed(a, lastA)) { lastChoice = 'A'; buttonReact("A", soundA, "answer", "A", true); }
    if (justPressed(b, lastB)) { lastChoice = 'B'; buttonReact("B", soundB, "answer", "B", true); }
    if (justPressed(c, lastC)) { lastChoice = 'C'; buttonReact("C", soundC, "answer", "C", true); }
    if (justPressed(d, lastD)) { lastChoice = 'D'; buttonReact("D", soundD, "answer", "D", true); }
  } else {
    // Still update the edge-detect history so we don't fire on join transition
    lastA = a; lastB = b; lastC = c; lastD = d;
  }

  delay(20);
}