#include <WiFi.h>
#include <PubSubClient.h>
#include <kontrol_suara_inferencing.h>
#include "freertos/FreeRTOS.h"
#include "freertos/task.h"
#include "driver/i2s.h"

// === WiFi & MQTT ===
const char* ssid = "Koala";
const char* password = "joinprayerdulu";
const char* mqtt_server = "test.mosquitto.org"; 
WiFiClient espClient;
PubSubClient client(espClient);
bool mqtt_connected = false;


// === Variabel LED saat rekaman window audio ===
bool window_led_active = false;
unsigned long window_listen_start_time = 0;
const unsigned long window_duration_ms = (EI_CLASSIFIER_RAW_SAMPLE_COUNT * 1000UL) / EI_CLASSIFIER_FREQUENCY;
bool voiceAutoLockInProgress = false;
unsigned long voiceAutoLockStart = 0;
unsigned long ledper_last_blink = 0;
bool ledper_state = false;
const unsigned long ledper_blink_interval = 100;
bool skip_blink_active = false;


// === Mode Wake Word dan Perintah ===
bool wake_word_active = false;
unsigned long command_start_time = 0;
const unsigned long command_window_ms =20000;


// === Output Pin ===
#define PIN_LAMPU_KAMAR 15
#define PIN_LAMPU_TERAS 4
#define PIN_KIPAS 5
#define PIN_SOLENOID 18
#define PIN_LED 19
#define PIN_LEDPER 21

// === Inferensi Audio ===
typedef struct {
  signed short* buffers[2];
  unsigned char buf_select;
  unsigned char buf_ready;
  unsigned int buf_count;
  unsigned int n_samples;
} inference_t;

static inference_t inference;
static const uint32_t sample_buffer_size = 2048;
static signed short sampleBuffer[sample_buffer_size];
static bool debug_nn = false;
static int print_results = -(EI_CLASSIFIER_SLICES_PER_MODEL_WINDOW);
static bool record_status = true;
static TaskHandle_t captureTaskHandle = NULL;

// === I2S Audio ===
static int i2s_init(uint32_t sampling_rate) {
  i2s_config_t i2s_config = {
    .mode = (i2s_mode_t)(I2S_MODE_MASTER | I2S_MODE_RX),
    .sample_rate = sampling_rate,
    .bits_per_sample = I2S_BITS_PER_SAMPLE_16BIT,
    .channel_format = I2S_CHANNEL_FMT_ONLY_LEFT,
    .communication_format = I2S_COMM_FORMAT_I2S,
    .intr_alloc_flags = 0,
    .dma_buf_count = 8,
    .dma_buf_len = 512,
    .use_apll = false,
    .tx_desc_auto_clear = false,
    .fixed_mclk = -1,
  };
  i2s_pin_config_t pin_config = {
    .bck_io_num = 26,
    .ws_io_num = 32,
    .data_out_num = -1,
    .data_in_num = 33,
  };

  esp_err_t ret = i2s_driver_install(I2S_NUM_1, &i2s_config, 0, NULL);
  if (ret != ESP_OK) ei_printf("i2s_driver_install error\n");
  ret = i2s_set_pin(I2S_NUM_1, &pin_config);
  if (ret != ESP_OK) ei_printf("i2s_set_pin error\n");
  i2s_zero_dma_buffer(I2S_NUM_1);
  return ret == ESP_OK ? 0 : -1;
}

static void audio_inference_callback(uint32_t n_bytes) {
  for (int i = 0; i < n_bytes >> 1; i++) {
    inference.buffers[inference.buf_select][inference.buf_count++] = sampleBuffer[i];
    if (inference.buf_count >= inference.n_samples) {
      inference.buf_select ^= 1;
      inference.buf_count = 0;
      inference.buf_ready = 1;
    }
  }
}

static void capture_samples(void* arg) {
  size_t bytes_read = sample_buffer_size;
  while (record_status) {
    i2s_read(I2S_NUM_1, (void*)sampleBuffer, sample_buffer_size, &bytes_read, 100);
    if (bytes_read > 0) {
      for (int x = 0; x < bytes_read / 2; x++) {
        sampleBuffer[x] = (int16_t)(sampleBuffer[x]) * 16;
      }
      audio_inference_callback(bytes_read);
    }
  }
  vTaskDelete(NULL);
}

static bool microphone_inference_start(uint32_t n_samples) {
  inference.buffers[0] = (short*)malloc(n_samples * sizeof(short));
  inference.buffers[1] = (short*)malloc(n_samples * sizeof(short));
  if (!inference.buffers[0] || !inference.buffers[1]) return false;

  inference.buf_select = 0;
  inference.buf_count = 0;
  inference.n_samples = n_samples;
  inference.buf_ready = 0;

  if (i2s_init(EI_CLASSIFIER_FREQUENCY) != 0) return false;
  record_status = true;

  if (captureTaskHandle == NULL) {
    xTaskCreate(capture_samples, "CaptureSamples", 1024 * 32, NULL, 10, &captureTaskHandle);
    ei_printf("🎙 Task audio capture dimulai\n");
  }

  return true;
}

static bool microphone_inference_record(void) {
  if (inference.buf_ready == 1) return false;
  while (inference.buf_ready == 0) delay(1);
  inference.buf_ready = 0;
  return true;
}

static int microphone_audio_signal_get_data(size_t offset, size_t length, float* out_ptr) {
  numpy::int16_to_float(&inference.buffers[inference.buf_select ^ 1][offset], out_ptr, length);
  return 0;
}

// === WiFi Setup ===
void setup_wifi() {
  WiFi.begin(ssid, password);
  while (WiFi.status() != WL_CONNECTED) {
    delay(500);
    Serial.print(".");
  }
  Serial.println("\n✅ WiFi terhubung!");
}

// === MQTT Callback ===
bool autoLockInProgress = false;
unsigned long autoLockStart = 0;

void publishMQTT(const char* topic, const char* message) {
  if (client.connected()) client.publish(topic, message, true);
}

void publishStatus(const char* device, bool state) {
  String topic = String("cristine/status/") + device;
  const char* message = state ? "on" : "off";
  publishMQTT(topic.c_str(), message);
}

void callback(char* topic, byte* payload, unsigned int length) {
  String message;
  for (unsigned int i = 0; i < length; i++) message += (char)payload[i];
  message.trim();
  message.toLowerCase();

  if (String(topic) == "cristine/control/lampu_kamar") {
    bool state = (message == "on");
    digitalWrite(PIN_LAMPU_KAMAR, state ? HIGH : LOW);
    publishStatus("lampu_kamar", state);
    ei_printf("📡 Lampu kamar %s via MQTT\n", state ? "dinyalakan" : "dimatikan");
  } else if (String(topic) == "cristine/control/lampu_teras") {
    bool state = (message == "on");
    digitalWrite(PIN_LAMPU_TERAS, state ? HIGH : LOW);
    publishStatus("lampu_teras", state);
    ei_printf("📡 Lampu teras %s via MQTT\n", state ? "dinyalakan" : "dimatikan");
  } else if (String(topic) == "cristine/control/kipas") {
    bool state = (message == "on");
    digitalWrite(PIN_KIPAS, state ? HIGH : LOW);
    publishStatus("kipas", state);
    ei_printf("📡 Kipas %s via MQTT\n", state ? "dinyalakan" : "dimatikan");
  } else if (String(topic) == "cristine/control/pintu") {
    if (message == "on" && !autoLockInProgress) {
      ei_printf("📡 Pintu dibuka via MQTT\n");
      digitalWrite(PIN_SOLENOID, HIGH);
      publishStatus("pintu", true);
      autoLockStart = millis();
      autoLockInProgress = true;
    }
  } else if (String(topic) == "cristine/request/status") {
    String statusPayload = "{";
    statusPayload += "\"lampu kamar\": " + String(digitalRead(PIN_LAMPU_KAMAR) == HIGH ? "true" : "false") + ",";
    statusPayload += "\"lampu teras\": " + String(digitalRead(PIN_LAMPU_TERAS) == HIGH ? "true" : "false") + ",";
    statusPayload += "\"kipas\": " + String(digitalRead(PIN_KIPAS) == HIGH ? "true" : "false") + ",";
    statusPayload += "\"pintu\": \"" + String(digitalRead(PIN_SOLENOID) == HIGH ? "dibuka" : "dikunci") + "\"";
    statusPayload += "}";
    client.publish("cristine/status", statusPayload.c_str());
    ei_printf("📤 Status terkirim ke bot WhatsApp\n");
  }
}

// === MQTT Reconnect ===
void reconnect() {
  while (!client.connected()) {
    Serial.println("🔁 Reconnecting MQTT...");
    String clientId = "ESP32Client-" + String(random(0xffff), HEX);
    if (client.connect(clientId.c_str())) {
      Serial.println("✅ MQTT Connected!");
      mqtt_connected = true;

      client.subscribe("cristine/control/lampu_kamar");
      client.subscribe("cristine/control/lampu_teras");
      client.subscribe("cristine/control/kipas");
      client.subscribe("cristine/control/pintu");
      client.subscribe("cristine/request/status");
    } else {
      Serial.print("❌ MQTT gagal, rc=");
      Serial.println(client.state());
      mqtt_connected = false;
      delay(5000);
    }
  }
}

// === Stop Inference ===
void stop_inference() {
  record_status = false;
  if (captureTaskHandle != NULL) {
    vTaskDelete(captureTaskHandle);
    captureTaskHandle = NULL;
  }
  i2s_driver_uninstall(I2S_NUM_1);
  free(inference.buffers[0]);
  free(inference.buffers[1]);
  ei_printf("🛑 Inferensi dihentikan\n");
}

void setup() {
  Serial.begin(115200);
  setup_wifi();
  client.setServer(mqtt_server, 1883);
  client.setCallback(callback);

  pinMode(PIN_LAMPU_KAMAR, OUTPUT);
  pinMode(PIN_LAMPU_TERAS, OUTPUT);
  pinMode(PIN_KIPAS, OUTPUT);
  pinMode(PIN_SOLENOID, OUTPUT);
  pinMode(PIN_LED, OUTPUT);
  pinMode(PIN_LEDPER, OUTPUT);
  digitalWrite(PIN_LED, LOW);
  digitalWrite(PIN_LEDPER, LOW);

  run_classifier_init();
  ei_printf("\n🎬 Menunggu koneksi MQTT...\n");
}

void loop() {
  client.loop();
  if (!client.connected()) {
    mqtt_connected = false;
    if (captureTaskHandle != NULL) stop_inference();
    reconnect();
  }

  if (autoLockInProgress && millis() - autoLockStart >= 10000) {
    digitalWrite(PIN_SOLENOID, LOW);
    publishStatus("pintu", false);
    ei_printf("🔒 Pintu dikunci otomatis setelah 10 detik\n");
    autoLockInProgress = false;
  }

  if (voiceAutoLockInProgress && millis() - voiceAutoLockStart >= 10000) {
  digitalWrite(PIN_SOLENOID, LOW);
  publishStatus("pintu", false);
  ei_printf("🔒 Pintu dikunci otomatis setelah 10 detik (via suara)\n");
  voiceAutoLockInProgress = false;
  }

  if (mqtt_connected && captureTaskHandle == NULL) {
    ei_printf("🎙 Mulai inferensi\n");
    if (!microphone_inference_start(EI_CLASSIFIER_SLICE_SIZE)) return;
  }

  if (!microphone_inference_record()) return;

  // === LED Perintah: Berkedip saat skip aktif
  if (skip_blink_active) {
  if (millis() - ledper_last_blink >= ledper_blink_interval) {
    ledper_last_blink = millis();
    ledper_state = !ledper_state;
    digitalWrite(PIN_LEDPER, ledper_state);
      }
    }

  // === LED Indikator Window Audio
  if (!window_led_active) {
    window_listen_start_time = millis();
    window_led_active = true;
    digitalWrite(PIN_LED, HIGH);
  }
  if (window_led_active && millis() - window_listen_start_time >= window_duration_ms) {
    window_led_active = false;
    digitalWrite(PIN_LED, LOW);
  }

  // === Inference Voice Command
  signal_t signal;
  signal.total_length = EI_CLASSIFIER_SLICE_SIZE;
  signal.get_data = &microphone_audio_signal_get_data;
  ei_impulse_result_t result = { 0 };
  EI_IMPULSE_ERROR r = run_classifier_continuous(&signal, &result, debug_nn);
  if (r != EI_IMPULSE_OK) return;

  static int print_results = -(EI_CLASSIFIER_SLICES_PER_MODEL_WINDOW);
  static bool skip_next_command_window = false;
  static unsigned long perintah_start_time = 0;

  if (++print_results >= EI_CLASSIFIER_SLICES_PER_MODEL_WINDOW) {
    print_results = 0;

    const char* top_label = "";
    float top_score = 0.0f;
    float silent_score = 0.0f, unknown_score = 0.0f;

    for (size_t i = 0; i < EI_CLASSIFIER_LABEL_COUNT; i++) {
      float score = result.classification[i].value;
      if (score > top_score) {
        top_score = score;
        top_label = result.classification[i].label;
      }

      if (strcmp(result.classification[i].label, "silent") == 0) silent_score = score;
      if (strcmp(result.classification[i].label, "unknown") == 0) unknown_score = score;
    }

    bool noise_only = (silent_score > 0.7f || unknown_score > 0.7f);

    // === Logging Confidence
    if (!wake_word_active && !skip_next_command_window) {
      ei_printf("[Menunggu Wake Word] %s: %.2f\n", top_label, top_score);
    } else if (wake_word_active && skip_next_command_window && print_results == 0) {
      ei_printf("[Lewati 1 Window] %s: %.2f\n", top_label, top_score);
    } else if (wake_word_active && !skip_next_command_window) {
      ei_printf("[Mode Perintah] %s: %.2f\n", top_label, top_score);
    }

    // === Wake Word Detection
    if (!wake_word_active && strcmp(top_label, "wake_word") == 0 && top_score > 0.90f && !noise_only) {
      ei_printf("🟢 Wake word terdeteksi, masuk ke mode perintah...\n");
      wake_word_active = true;
      skip_next_command_window = true;
      skip_blink_active = true;
      perintah_start_time = millis();
    }

    // === Mode Perintah
    else if (wake_word_active) {
      unsigned long elapsed = millis() - perintah_start_time;

      // LED perintah: berkedip saat skip, nyala tetap saat mode perintah aktif
      if (skip_next_command_window) {
        if (millis() - ledper_last_blink >= ledper_blink_interval) {
          ledper_last_blink = millis();
          ledper_state = !ledper_state;
          digitalWrite(PIN_LEDPER, ledper_state);
        }

        // Pastikan hanya 1 window dilewati (durasi window)
if (elapsed >= window_duration_ms) {
  skip_next_command_window = false;
  skip_blink_active = false; 
  ei_printf("✅ Mode perintah aktif, siap menerima perintah\n");
  digitalWrite(PIN_LEDPER, HIGH); // nyala tetap setelah skip
}
      } else {
        digitalWrite(PIN_LEDPER, HIGH);
      }

      if (elapsed >= command_window_ms) {
        ei_printf("⏱ Timeout! Kembali ke mode wake word\n");
        wake_word_active = false;
        skip_next_command_window = false;
        digitalWrite(PIN_LEDPER, LOW);
      }

      else if (!skip_next_command_window) {
        if (strcmp(top_label, "silent") != 0 &&
            strcmp(top_label, "unknown") != 0 &&
            strcmp(top_label, "wake_word") != 0 &&
            top_score > 0.88f) {

          ei_printf("🗣 Perintah terdeteksi: %s (%.2f)\n", top_label, top_score);

          if (strcmp(top_label, "nyalakan_lampu_kamar") == 0) {
            digitalWrite(PIN_LAMPU_KAMAR, HIGH);
            publishMQTT("cristine/control/lampu_kamar", "on");
            publishStatus("lampu_kamar", true);
          } else if (strcmp(top_label, "matikan_lampu_kamar") == 0) {
            digitalWrite(PIN_LAMPU_KAMAR, LOW);
            publishMQTT("cristine/control/lampu_kamar", "off");
            publishStatus("lampu_kamar", false);
          } else if (strcmp(top_label, "nyalakan_lampu_teras") == 0) {
            digitalWrite(PIN_LAMPU_TERAS, HIGH);
            publishMQTT("cristine/control/lampu_teras", "on");
            publishStatus("lampu_teras", true);
          } else if (strcmp(top_label, "matikan_lampu_teras") == 0) {
            digitalWrite(PIN_LAMPU_TERAS, LOW);
            publishMQTT("cristine/control/lampu_teras", "off");
            publishStatus("lampu_teras", false);
          } else if (strcmp(top_label, "nyalakan_kipas") == 0) {
            digitalWrite(PIN_KIPAS, HIGH);
            publishMQTT("cristine/control/kipas", "on");
            publishStatus("kipas", true);
          } else if (strcmp(top_label, "matikan_kipas") == 0) {
            digitalWrite(PIN_KIPAS, LOW);
            publishMQTT("cristine/control/kipas", "off");
            publishStatus("kipas", false);
          } else if (strcmp(top_label, "buka_pintu") == 0) {
              digitalWrite(PIN_SOLENOID, HIGH);
              publishMQTT("cristine/control/pintu", "on");
              publishStatus("pintu", true);
              voiceAutoLockStart = millis();
              voiceAutoLockInProgress = true;
          } else {
            ei_printf("❓ Perintah '%s' tidak dikenali\n", top_label);
          }

          wake_word_active = false;
          digitalWrite(PIN_LEDPER, LOW);
        }
      }
    }
  }
}