/*
 * voice_assistant.ino — ESP32-S3 microphone-to-MQTT bridge
 *
 * Hub Fixes — Directive 3 hardening (2026-10-14).
 *
 * What changed and why:
 *
 *  1. MQTT auto-reconnect + re-subscribe on EVERY (re)connect.
 *     PubSubClient reconnects to a restarted broker but does NOT restore
 *     subscriptions. Without an explicit re-subscribe the client can hold a
 *     live socket with no subscription and go silently deaf. That is the same
 *     bug class that silenced voice_bridge on 2026-10-13.
 *
 *  2. Heartbeat every 5 minutes on voice/audio_stream.
 *     The Directive-1 watchdog watches that topic. The heartbeat is a valid
 *     silent PCM frame, so voice_bridge treats it as silence and ignores it,
 *     while the watchdog sees that the publisher is alive. This makes silence
 *     a reliable "the device is dead" signal instead of "nobody spoke".
 *
 *  3. LED shows the MQTT state, not just Wi-Fi.
 *     GREEN = MQTT connected and publishing (idle: steady green, streaming:
 *     bright green). Wi-Fi down = RED blink. Wi-Fi up but MQTT down = AMBER
 *     blink. Green is never shown when the device cannot publish.
 *
 *  4. The serial console logs the reason whenever the broker connection drops
 *     (the PubSubClient state code, or the Wi-Fi disconnect reason).
 *
 * Broker: the shared .13 voice broker on port 1883. That plane is plain and
 * shared with n8n and the voice path. Do not change it, and do not put TLS on
 * this client. config.h and led_control.h in this folder are stale; this
 * sketch is the single source of truth for the running firmware.
 */

#include <WiFi.h>
#include <PubSubClient.h>
#include <driver/i2s.h>
#include <Freenove_WS2812_Lib_for_ESP32.h>

// ---------------------------------------------------------------------------
// CONFIG — edit the credential placeholders before flashing.
// Never commit real credentials.
// ---------------------------------------------------------------------------

// --- WiFi ---
const char* WIFI_SSID = "Aussie Broadband 8729";
const char* WIFI_PASS = "WIFI_PASS_REDACTED";

// --- MQTT (.13 voice broker, plain 1883) ---
const char* MQTT_BROKER    = "192.168.20.13";
const int   MQTT_PORT      = 1883;
const char* MQTT_USER      = "MQTT_USER_REDACTED";
const char* MQTT_PASS      = "MQTT_PASS_REDACTED";
const char* MQTT_CLIENT_ID = "ESP32Mic";

// --- Topics ---
const char* TOPIC_AUDIO  = "voice/audio_stream"; // audio + heartbeat liveness
const char* TOPIC_STATUS = "voice/status";       // WAKE / OK / processing / presence

// --- Hardware (ESP32-S3) ---
#define I2S_SCK 4
#define I2S_WS  5
#define I2S_SD  6
#define I2S_PORT I2S_NUM_0
#define LED_PIN  1
#define NUM_LEDS 12

// --- Tuning ---
#define VAD_THRESHOLD 1000
#define SAMPLE_RATE   16000
#define BLOCK_SIZE    512

// --- Timing ---
#define HEARTBEAT_INTERVAL_MS      (5UL * 60UL * 1000UL) // 5 minutes
#define MQTT_RECONNECT_INTERVAL_MS 5000UL
#define WIFI_RECONNECT_INTERVAL_MS 5000UL
#define WIFI_CONNECT_TIMEOUT_MS    20000UL

int16_t sBuffer[BLOCK_SIZE];

WiFiClient espClient;
PubSubClient mqtt(espClient);
Freenove_ESP32_WS2812 strip(NUM_LEDS, LED_PIN, 0, TYPE_GRB);

// --- STATES ---
enum DeviceState {
  STATE_IDLE,        // connected, no voice
  STATE_STREAMING,   // voice detected, publishing audio (bright green)
  STATE_LISTENING,   // wake word confirmed, publishing audio (bright green)
  STATE_PROCESSING,  // waiting for the server (blue)
  STATE_SUCCESS,     // command executed (yellow)
  STATE_ERROR        // timeout (red blink)
};
DeviceState currentState = STATE_IDLE;
unsigned long stateTimer = 0;

// --- CONNECTION TRACKING ---
bool wasMqttConnected = false;
unsigned long lastMqttAttempt = 0;
unsigned long lastWifiAttempt = 0;
unsigned long lastHeartbeat = 0;
uint32_t ledValue = 0xFFFFFFFF; // sentinel so the first draw always happens

const i2s_config_t i2s_config = {
  .mode = (i2s_mode_t)(I2S_MODE_MASTER | I2S_MODE_RX),
  .sample_rate = SAMPLE_RATE,
  .bits_per_sample = I2S_BITS_PER_SAMPLE_16BIT,
  .channel_format = I2S_CHANNEL_FMT_ONLY_LEFT,
  .communication_format = I2S_COMM_FORMAT_STAND_I2S,
  .intr_alloc_flags = ESP_INTR_FLAG_LEVEL1,
  .dma_buf_count = 8,
  .dma_buf_len = 64,
  .use_apll = true
};

const i2s_pin_config_t pin_config = {
  .bck_io_num = I2S_SCK,
  .ws_io_num = I2S_WS,
  .data_out_num = I2S_PIN_NO_CHANGE,
  .data_in_num = I2S_SD
};

// ---------------------------------------------------------------------------
// LED
// ---------------------------------------------------------------------------

// Only push a frame when the colour actually changed. strip.show() disables
// interrupts for the duration of the frame, so we must not call it every loop.
void applyLed(uint8_t r, uint8_t g, uint8_t b) {
  uint32_t value = ((uint32_t)r << 16) | ((uint32_t)g << 8) | (uint32_t)b;
  if (value == ledValue) return;
  ledValue = value;
  strip.setAllLedsColor(r, g, b);
  strip.show();
}

// The LED contract:
//   Wi-Fi down                       -> RED blink
//   Wi-Fi up, MQTT down              -> AMBER blink
//   MQTT connected + publishing      -> GREEN (steady when idle)
//   ... activity overlays            -> blue / yellow while still connected.
// Green therefore means "connected and publishing", never "Wi-Fi only".
void renderLed(bool mqttUp) {
  const bool blink = (millis() / 400) % 2 == 0;

  if (WiFi.status() != WL_CONNECTED) {
    applyLed(blink ? 90 : 0, 0, 0);
    return;
  }
  if (!mqttUp) {
    applyLed(blink ? 90 : 0, blink ? 45 : 0, 0);
    return;
  }

  switch (currentState) {
    case STATE_IDLE:       applyLed(0, 120, 0); break;             // healthy green
    case STATE_STREAMING:  applyLed(0, 255, 0); break;             // recording
    case STATE_LISTENING:  applyLed(0, 255, 0); break;             // talking to Jarvis
    case STATE_PROCESSING: applyLed(0, 0, 255); break;             // thinking
    case STATE_SUCCESS:    applyLed(255, 200, 0); break;           // done
    case STATE_ERROR:      applyLed(blink ? 255 : 0, 0, 0); break; // fault
  }
}

// ---------------------------------------------------------------------------
// CONNECTION LOGGING
// ---------------------------------------------------------------------------

const char* mqttStateStr(int state) {
  switch (state) {
    case -4: return "CONNECTION_TIMEOUT";
    case -3: return "CONNECTION_LOST";
    case -2: return "CONNECT_FAILED";
    case -1: return "DISCONNECTED";
    case  0: return "CONNECTED";
    case  1: return "BAD_PROTOCOL";
    case  2: return "BAD_CLIENT_ID";
    case  3: return "UNAVAILABLE";
    case  4: return "BAD_CREDENTIALS";
    case  5: return "UNAUTHORIZED";
    default: return "UNKNOWN";
  }
}

// Log the Wi-Fi reason from the driver, not a guess.
void onWifiEvent(WiFiEvent_t event, WiFiEventInfo_t info) {
  switch (event) {
    case ARDUINO_EVENT_WIFI_STA_CONNECTED:
      Serial.println("[wifi] joined access point");
      break;
    case ARDUINO_EVENT_WIFI_STA_GOT_IP:
      Serial.printf("[wifi] got IP %s\n", WiFi.localIP().toString().c_str());
      break;
    case ARDUINO_EVENT_WIFI_STA_DISCONNECTED:
      Serial.printf("[wifi] DISCONNECTED reason=%d\n", info.wifi_sta_disconnected.reason);
      break;
    default:
      break;
  }
}

// ---------------------------------------------------------------------------
// MQTT
// ---------------------------------------------------------------------------

void mqttCallback(char* topic, byte* payload, unsigned int length) {
  char msg[length + 1];
  memcpy(msg, payload, length);
  msg[length] = '\0';

  String message = String(msg);

  if (message == "WAKE") {
    currentState = STATE_LISTENING;
    stateTimer = millis();
  } else if (message == "OK") {
    currentState = STATE_SUCCESS;
    stateTimer = millis();
  }
}

// Connect (or re-connect) and re-subscribe. Called on every (re)connect so the
// client is never left holding a live socket with no subscription.
bool connectMqtt() {
  lastMqttAttempt = millis();
  Serial.printf("[mqtt] connecting to %s:%d (wifi=%d)\n",
                MQTT_BROKER, MQTT_PORT, (int)WiFi.status());

  // Retained Last Will on voice/status: a hard drop marks the device offline.
  // voice_bridge only acts on "processing"/"WAKE"/"OK", so this is inert to it.
  bool ok = mqtt.connect(MQTT_CLIENT_ID, MQTT_USER, MQTT_PASS,
                         TOPIC_STATUS, 0, true, "offline");
  if (!ok) {
    Serial.printf("[mqtt] connect FAILED state=%d (%s)\n",
                  mqtt.state(), mqttStateStr(mqtt.state()));
    return false;
  }

  Serial.println("[mqtt] connected");
  bool sub = mqtt.subscribe(TOPIC_STATUS);
  Serial.printf("[mqtt] re-subscribe %s -> %s\n", TOPIC_STATUS, sub ? "ok" : "FAILED");

  bool online = mqtt.publish(TOPIC_STATUS, "online", true);
  Serial.printf("[mqtt] presence online -> %s\n", online ? "ok" : "FAILED");

  lastHeartbeat = millis(); // do not fire a heartbeat immediately after connect
  return true;
}

// ---------------------------------------------------------------------------
// WIFI
// ---------------------------------------------------------------------------

void wifiBegin() {
  Serial.printf("[wifi] connecting to SSID \"%s\"\n", WIFI_SSID);
  WiFi.mode(WIFI_STA);
  WiFi.begin(WIFI_SSID, WIFI_PASS);
  lastWifiAttempt = millis();
}

bool wifiEnsure() {
  if (WiFi.status() == WL_CONNECTED) return true;
  if (millis() - lastWifiAttempt >= WIFI_RECONNECT_INTERVAL_MS) {
    Serial.printf("[wifi] down (status=%d); reconnecting\n", (int)WiFi.status());
    WiFi.disconnect();
    WiFi.begin(WIFI_SSID, WIFI_PASS);
    lastWifiAttempt = millis();
  }
  return false;
}

// ---------------------------------------------------------------------------
// SETUP
// ---------------------------------------------------------------------------

void setup() {
  delay(2000); // let USB CDC serial come up
  Serial.begin(115200);
  Serial.println();
  Serial.println("[boot] voice_assistant starting (Directive 3 firmware)");

  strip.begin();
  strip.setBrightness(30);
  applyLed(60, 60, 60); // white boot

  i2s_driver_install(I2S_PORT, &i2s_config, 0, NULL);
  i2s_set_pin(I2S_PORT, &pin_config);

  mqtt.setBufferSize(2048);
  mqtt.setServer(MQTT_BROKER, MQTT_PORT);
  mqtt.setCallback(mqttCallback);
  mqtt.setKeepAlive(15);

  WiFi.onEvent(onWifiEvent);
  wifiBegin();

  // Bounded first-join wait. If it fails, loop() keeps retrying non-stop.
  const unsigned long start = millis();
  while (WiFi.status() != WL_CONNECTED &&
         millis() - start < WIFI_CONNECT_TIMEOUT_MS) {
    delay(250);
    applyLed(0, 0, ((millis() / 400) % 2) ? 80 : 0); // blue blink = joining
  }
  if (WiFi.status() == WL_CONNECTED) {
    Serial.printf("[wifi] connected, ip=%s\n", WiFi.localIP().toString().c_str());
  } else {
    Serial.println("[wifi] no connection yet; will keep retrying");
  }
}

// ---------------------------------------------------------------------------
// LOOP
// ---------------------------------------------------------------------------

void loop() {
  // 1) Wi-Fi first: MQTT cannot work without it.
  const bool wifiUp = wifiEnsure();
  bool nowConnected = false;

  // 2) MQTT connect / reconnect with a non-blocking backoff.
  if (wifiUp) {
    if (!mqtt.connected() &&
        (millis() - lastMqttAttempt >= MQTT_RECONNECT_INTERVAL_MS)) {
      connectMqtt();
    }
    if (mqtt.connected()) {
      mqtt.loop();
      nowConnected = true;
    }
  }

  // 3) Log every drop exactly once, and name the reason.
  if (wasMqttConnected && !nowConnected) {
    if (WiFi.status() != WL_CONNECTED) {
      Serial.printf("[mqtt] connection DROPPED: Wi-Fi lost (wifi=%d). "
                    "Reconnect + re-subscribe is automatic.\n",
                    (int)WiFi.status());
    } else {
      Serial.printf("[mqtt] connection DROPPED: state=%d (%s). "
                    "Reconnect + re-subscribe is automatic.\n",
                    mqtt.state(), mqttStateStr(mqtt.state()));
    }
    currentState = STATE_IDLE; // never leave a stale green/streaming state
  } else if (!wasMqttConnected && nowConnected) {
    Serial.println("[mqtt] link up: connected and subscribed");
  }
  wasMqttConnected = nowConnected;

  // 4) Heartbeat: prove the publisher is alive so topic silence means "dead".
  if (nowConnected &&
      currentState != STATE_STREAMING && currentState != STATE_LISTENING &&
      (millis() - lastHeartbeat >= HEARTBEAT_INTERVAL_MS)) {
    lastHeartbeat = millis();
    static const uint8_t silence[256] = {0}; // 128 samples of silent PCM
    bool ok = mqtt.publish(TOPIC_AUDIO, silence, sizeof(silence));
    Serial.printf("[heartbeat] %u bytes on %s -> %s\n",
                  (unsigned)sizeof(silence), TOPIC_AUDIO, ok ? "ok" : "FAILED");
    if (!ok) {
      Serial.println("[heartbeat] publish failed; forcing MQTT reconnect");
      mqtt.disconnect();
      lastMqttAttempt = millis();
    }
  }

  // 5) Device-state timers.
  if (currentState == STATE_PROCESSING && millis() - stateTimer > 6000) {
    for (int i = 0; i < 3; i++) {
      applyLed(255, 0, 0); delay(200);
      applyLed(0, 0, 0);   delay(200);
    }
    currentState = STATE_IDLE;
  }
  if (currentState == STATE_SUCCESS && millis() - stateTimer > 3000) {
    currentState = STATE_IDLE;
  }

  renderLed(nowConnected);

  // 6) Read audio with a bounded timeout. portMAX_DELAY would stall Wi-Fi and
  //    MQTT housekeeping forever if the I2S peripheral stopped producing data.
  size_t bytesRead = 0;
  i2s_read(I2S_PORT, &sBuffer, sizeof(sBuffer), &bytesRead, pdMS_TO_TICKS(50));
  if (bytesRead < 2) {
    return;
  }

  const int samples = bytesRead / 2;

  // Gain boost (x6).
  for (int i = 0; i < samples; i++) {
    int32_t boosted = sBuffer[i] * 6;
    if (boosted > 32767) boosted = 32767;
    if (boosted < -32768) boosted = -32768;
    sBuffer[i] = (int16_t)boosted;
  }

  long sum = 0;
  for (int i = 0; i < samples; i++) {
    sum += abs(sBuffer[i]);
  }
  const int avg = samples > 0 ? (int)(sum / samples) : 0;

  // Never publish audio into a dead socket.
  if (!nowConnected) {
    currentState = STATE_IDLE;
    return;
  }

  // Speech detected while idle -> start streaming (bright green).
  if (currentState == STATE_IDLE && avg > VAD_THRESHOLD) {
    currentState = STATE_STREAMING;
    stateTimer = millis();
    mqtt.publish(TOPIC_AUDIO, (const uint8_t*)sBuffer, bytesRead);
  }

  // Stream while active.
  if (currentState == STATE_STREAMING || currentState == STATE_LISTENING) {
    mqtt.publish(TOPIC_AUDIO, (const uint8_t*)sBuffer, bytesRead);

    if (avg > VAD_THRESHOLD) stateTimer = millis();

    // Silence timeout (1.5 s).
    if (millis() - stateTimer > 1500) {
      if (currentState == STATE_LISTENING) {
        // We were talking to Jarvis; now wait for the reply.
        currentState = STATE_PROCESSING;
        stateTimer = millis();
        mqtt.publish(TOPIC_STATUS, "processing");
      } else {
        // Just noise; never woke up.
        currentState = STATE_IDLE;
      }
    }
  }
}