#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";

// POINT TO VOICE SERVER (Not Home Assistant)
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_ACK   = "voice/status"; 

// --- HARDWARE ---
#define I2S_SCK 4
#define I2S_WS  5
#define I2S_SD  6
#define I2S_PORT I2S_NUM_0

#define LED_PIN  5
#define NUM_LEDS 12

// --- AUDIO TUNING ---
// Increased to 2000 to ignore footsteps/typing (Range is 0-32768)
#define VAD_THRESHOLD 1000 
const int BLOCK_SIZE = 512;
int16_t sBuffer[BLOCK_SIZE];

// ... (Rest of Objects/Config remains standard) ...
WiFiClient espClient;
PubSubClient mqtt(espClient);
Freenove_ESP32_WS2812 strip(NUM_LEDS, LED_PIN, 0, TYPE_GRB);

enum DeviceState { STATE_IDLE, STATE_RECORDING, STATE_TRANSMITTING, STATE_ACK_RECEIVED };
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 Helpers remain the same) ...
void setRingColor(uint8_t r, uint8_t g, uint8_t b) {
  strip.setAllLedsColor(r, g, b);
  strip.show();
}
void updateLEDs() {
  if (currentState == STATE_ACK_RECEIVED) {
      for (int j = 0; j < 255; j += 10) {
        for (int i = 0; i < NUM_LEDS; i++) {
           byte pos = (i * 256 / NUM_LEDS + j) & 255;
           if(pos < 85) { strip.setLedColorData(i, pos * 3, 255 - pos * 3, 0); }
           else if(pos < 170) { pos -= 85; strip.setLedColorData(i, 255 - pos * 3, 0, pos * 3); }
           else { pos -= 170; strip.setLedColorData(i, 0, pos * 3, 255 - pos * 3); }
        }
        strip.show();
        delay(5);
      }
      currentState = STATE_IDLE;
      setRingColor(0,0,0);
  }
}
void mqttCallback(char* topic, byte* payload, unsigned int length) {
  // If Bridge sends a "WAKE" or "OK" message
  currentState = STATE_ACK_RECEIVED;
}

void setup() {
  delay(2000);
  Serial.begin(115200);
  
  // High Buffer for Audio
  mqtt.setBufferSize(2048);
  
  strip.begin();
  strip.setBrightness(30);
  setRingColor(50, 50, 50);

  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, 255); delay(100); setRingColor(0,0,0);
  }
  
  mqtt.setServer(MQTT_BROKER, MQTT_PORT);
  mqtt.setCallback(mqttCallback);
  
  setRingColor(0,0,0);
}

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

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

  // --- GAIN BOOST (x8) ---
  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);
  
  // Debug (View this to tune threshold)
  // Serial.printf("Vol: %d\n", avg); 

  if (currentState == STATE_IDLE && avg > VAD_THRESHOLD) {
    currentState = STATE_RECORDING;
    stateTimer = millis();
    setRingColor(0, 255, 0); // Green
    mqtt.publish(TOPIC_AUDIO, (const uint8_t*)sBuffer, bytesRead);
  }

  if (currentState == STATE_RECORDING) {
    mqtt.publish(TOPIC_AUDIO, (const uint8_t*)sBuffer, bytesRead);
    
    if (avg > VAD_THRESHOLD) stateTimer = millis();
    
    // Silence Timeout (1.5s)
    if (millis() - stateTimer > 1500) {
      currentState = STATE_TRANSMITTING;
      // Send marker to Voice Server
      mqtt.publish("voice/status", "processing"); 
      setRingColor(0, 0, 255); // Blue
    }
  }
  
  if (currentState == STATE_TRANSMITTING && millis() - stateTimer > 4000) {
    currentState = STATE_IDLE;
    setRingColor(0,0,0);
  }
}