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

// --- CONNECTIVITY ---
const char* WIFI_SSID = "Aussie Broadband 8729";
const char* WIFI_PASS = "Ffdfmunfca";

const char* MQTT_BROKER = "192.168.20.13"; 
const int   MQTT_PORT   = 1883;
const char* MQTT_USER   = "mqtt-user";
const char* MQTT_PASS   = "sam4jo";

const char* TOPIC_AUDIO = "voice/audio_stream";
const char* TOPIC_STATUS = "voice/status"; 

// --- HARDWARE ---
#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 
const int BLOCK_SIZE = 512;
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, 
  STATE_STREAMING,   // Noise detected (Green)
  STATE_LISTENING,   // Wake Word Confirmed (Blue)
  STATE_PROCESSING,  // Sending to Server (Red)
  STATE_SUCCESS,     // Done (Yellow)
  STATE_ERROR        // Timeout (Blink Red)
};
DeviceState currentState = STATE_IDLE;
unsigned long stateTimer = 0;

const i2s_config_t i2s_config = {
  .mode = (i2s_mode_t)(I2S_MODE_MASTER | I2S_MODE_RX),
  .sample_rate = 16000,
  .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 HELPER ---
void setRingColor(uint8_t r, uint8_t g, uint8_t b) {
  strip.setAllLedsColor(r, g, b);
  strip.show();
}

// --- MQTT HANDLING ---
void mqttCallback(char* topic, byte* payload, unsigned int length) {
  char msg[length + 1];
  memcpy(msg, payload, length);
  msg[length] = '\0';
  
  String message = String(msg);
  // Serial.print("Recv: "); Serial.println(message); // Debug

  if (message == "WAKE") {
    // Server heard "Hey Jarvis" -> Turn BLUE
    currentState = STATE_LISTENING;
    stateTimer = millis(); // Reset timer to prevent premature timeout
    setRingColor(0, 0, 255); 
  }
  else if (message == "OK") {
    // Command executed -> Turn YELLOW
    currentState = STATE_SUCCESS;
    stateTimer = millis();
    setRingColor(255, 200, 0); // Yellow/Orange
  }
}

void setup() {
  delay(2000);
  Serial.begin(115200);
  
  mqtt.setBufferSize(2048);
  
  strip.begin();
  strip.setBrightness(30);
  setRingColor(50, 50, 50); // White Boot

  i2s_driver_install(I2S_PORT, &i2s_config, 0, NULL);
  i2s_set_pin(I2S_PORT, &pin_config);
  
  WiFi.begin(WIFI_SSID, WIFI_PASS);
  while (WiFi.status() != WL_CONNECTED) {
    delay(500);
    setRingColor(0, 0, 100); delay(100); setRingColor(0,0,0);
  }
  
  mqtt.setServer(MQTT_BROKER, MQTT_PORT);
  mqtt.setCallback(mqttCallback);
  
  setRingColor(0,0,0); // Idle
}

void loop() {
  if (!mqtt.connected()) {
    if (mqtt.connect("ESP32Mic", MQTT_USER, MQTT_PASS)) {
      mqtt.subscribe(TOPIC_STATUS);
    }
  }
  mqtt.loop();

  // --- TIMEOUT LOGIC (Fixes the "Lock Up") ---
  if (currentState == STATE_PROCESSING && millis() - stateTimer > 6000) {
    // We waited 6 seconds for an answer, but got nothing.
    // Blink Red 3 times then Reset.
    for(int i=0; i<3; i++) {
        setRingColor(255, 0, 0); delay(200);
        setRingColor(0, 0, 0); delay(200);
    }
    currentState = STATE_IDLE;
  }
  
  // Reset "Success" yellow light after 3 seconds
  if (currentState == STATE_SUCCESS && millis() - stateTimer > 3000) {
    currentState = STATE_IDLE;
    setRingColor(0,0,0);
  }

  size_t bytesRead;
  i2s_read(I2S_PORT, &sBuffer, sizeof(sBuffer), &bytesRead, portMAX_DELAY);

  // Gain Boost (x6)
  for (int i = 0; i < bytesRead / 2; 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 < bytesRead / 2; i++) {
    sum += abs(sBuffer[i]);
  }
  int avg = sum / (bytesRead / 2);

  // --- AUDIO LOGIC ---
  
  // 1. Detect Noise -> Green
  if (currentState == STATE_IDLE && avg > VAD_THRESHOLD) {
    currentState = STATE_STREAMING;
    stateTimer = millis();
    setRingColor(0, 255, 0); // Green (I hear you)
    
    // Send trigger buffer
    mqtt.publish(TOPIC_AUDIO, (const uint8_t*)sBuffer, bytesRead);
  }

  // 2. Stream Audio
  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.5s)
    if (millis() - stateTimer > 1500) {
      if (currentState == STATE_LISTENING) {
        // We were talking to Jarvis, now we wait for reply
        currentState = STATE_PROCESSING;
        stateTimer = millis(); // Start timeout timer
        mqtt.publish(TOPIC_STATUS, "processing"); 
        setRingColor(255, 0, 0); // Red (Thinking/Wait)
      } else {
        // Just noise, never woke up
        currentState = STATE_IDLE;
        setRingColor(0, 0, 0);
      }
    }
  }
}