#include <driver/i2s.h>
#include <WiFi.h>
#include <WiFiClientSecure.h>
#include <PubSubClient.h>
#include <ArduinoJson.h>
#include <TinyGPS++.h>
#include <HardwareSerial.h>
#include <math.h>

// ================================================================
//  PIN DEFINITIONS
// ================================================================
#define MIC_SCK    26
#define MIC_WS     20
#define MIC_SD     20
#define SPK_BCLK   19
#define SPK_LRC    33
#define SPK_DIN    11
#define GPS_RX     43
#define GPS_TX     44

#define BAT_ADC_PIN   34
#define CHRG_PIN      35
#define STDBY_PIN     36

// ================================================================
//  WIFI & MQTT (Updated for Secure HiveMQ Cloud)
// ================================================================
const char* WIFI_SSID   = "Wokwi-GUEST";
const char* WIFI_PASS   = "";

const char* MQTT_BROKER = "be07df79040549f0bdec3aed48f89f9c.s1.eu.hivemq.cloud"; 
const int   MQTT_PORT   = 8883; 

const char* DEVICE_ID   = "collar_001";

const char* MQTT_USER   = "Barker";
const char* MQTT_PASS   = "Skillarion@pet123";

// ================================================================
//  MQTT TOPICS
// ================================================================
const char* TOPIC_EMOTION        = "dogcollar/emotion";
const char* TOPIC_GPS            = "dogcollar/location";
const char* TOPIC_BATTERY        = "dogcollar/battery";
const char* TOPIC_AUDIO_META     = "dogcollar/audio/meta";
const char* TOPIC_AUDIO_STREAM   = "dogcollar/audio/stream";
const char* TOPIC_TTS_META       = "dogcollar/tts/meta";
const char* TOPIC_TTS_AUDIO      = "dogcollar/tts/audio";
const char* TOPIC_EMOTION_RESULT = "dogcollar/emotion/result";

// ================================================================
//  AUDIO SETTINGS
// ================================================================
#define SAMPLE_RATE         16000
#define AUDIO_CHUNK_SAMPLES 256
#define AUDIO_CHUNK_BYTES   (AUDIO_CHUNK_SAMPLES * 2)
#define AUDIO_TOTAL_SAMPLES 8000
#define MQTT_BUFFER_SIZE    4096  // Increased to support TLS overhead
#define SPK_BUFFER_SAMPLES  8000

// ================================================================
//  SIMULATION SETTINGS
// ================================================================
#define SIM_CYCLE_MS   5000
#define SIM_LAT        17.385044
#define SIM_LNG        78.486671

// ================================================================
//  GLOBAL VARIABLES
// ================================================================
char             g_emotion[16]   = "normal";
volatile long    g_audioRMS      = 0;
volatile float   g_lat           = SIM_LAT;
volatile float   g_lng           = SIM_LNG;
volatile bool    g_gpsValid      = false;

volatile float   g_batPercent    = 85.0;
volatile bool    g_isCharging    = false;
volatile bool    g_isStandby     = false;
volatile float   g_batVoltage    = 3.9;

int16_t* g_captureBuf    = nullptr;
volatile int     g_captureIdx    = 0;
volatile bool    g_capturing     = false;
volatile bool    g_captureReady  = false;
char             g_sessionId[16] = "sess_0001";
int              g_sessionCount  = 0;

int16_t* g_spkBuf        = nullptr;
volatile int     g_spkLen        = 0;
volatile int     g_spkExpected   = 0;
volatile int     g_spkReceived   = 0;
volatile bool    g_spkReady      = false;

SemaphoreHandle_t xMicMutex;
SemaphoreHandle_t xSpkMutex;

// ================================================================
//  OBJECTS (Updated for TLS/SSL Security)
// ================================================================
WiFiClientSecure wifiClient; 
PubSubClient     mqtt(wifiClient);
TinyGPSPlus      gps;
HardwareSerial   gpsSerial(1); // Use Serial1 or Serial2 cleanly without early hardware binding

// ================================================================
//  HELPERS
// ================================================================
void newSession() {
  g_sessionCount++;
  snprintf(g_sessionId, sizeof(g_sessionId), "sess_%04d", g_sessionCount);
}

float adcToVoltage(int raw) {
  return (raw / 4095.0f) * 3.3f * 2.0f;
}

float voltageToPercent(float v) {
  return constrain(((v - 3.0f) / 1.2f) * 100.0f, 0.0f, 100.0f);
}

// ================================================================
//  TASK 0 — BATTERY MONITOR
// ================================================================
void taskBattery(void* param) {
  pinMode(CHRG_PIN,  INPUT_PULLUP);
  pinMode(STDBY_PIN, INPUT_PULLUP);
  Serial.println("[BAT] Task started");

  float simPct   = 85.0f;
  bool  simCharg = false;

  while (true) {
    int   raw  = analogRead(BAT_ADC_PIN);
    float vReal = adcToVoltage(raw);
    float pReal = voltageToPercent(vReal);

    if (simCharg) {
      simPct += 0.5f;
      if (simPct >= 100.0f) { simPct = 100.0f; simCharg = false; }
    } else {
      simPct -= 0.3f;
      if (simPct <= 20.0f)  { simPct = 20.0f;  simCharg = true;  }
    }

    g_batPercent = (raw < 100) ? simPct : pReal;
    g_batVoltage = 3.0f + (g_batPercent / 100.0f) * 1.2f;
    g_isCharging = simCharg;
    g_isStandby  = (simPct >= 99.5f);

    Serial.printf("[BAT] %.1f%%  %.2fV  %s\n",
      g_batPercent, g_batVoltage,
      g_isStandby  ? "STANDBY" :
      g_isCharging ? "CHARGING" : "ON_BATTERY");

    vTaskDelay(10000 / portTICK_PERIOD_MS);
  }
}

// ================================================================
//  SIMULATED MIC RMS
// ================================================================
long getSimulatedRMS() {
  int step = (millis() / SIM_CYCLE_MS) % 5;
  switch (step) {
    case 0: return 200;
    case 1: return 600;
    case 2: return 2000;
    case 3: return 5000;
    case 4: return 9000;
    default: return 200;
  }
}

// ================================================================
//  MQTT CALLBACK
// ================================================================
void onMQTTMessage(char* topic, byte* payload, unsigned int len) {
  if (strcmp(topic, TOPIC_TTS_META) == 0) {
    StaticJsonDocument<128> doc;
    if (deserializeJson(doc, payload, len) != DeserializationError::Ok) return;
    const char* sid = doc["session"];
    if (strcmp(sid, g_sessionId) != 0) return;
    if (xSemaphoreTake(xSpkMutex, pdMS_TO_TICKS(100)) == pdTRUE) {
      g_spkExpected = doc["total_chunks"];
      g_spkReceived = 0; g_spkLen = 0; g_spkReady = false;
      xSemaphoreGive(xSpkMutex);
    }
    return;
  }

  if (strcmp(topic, TOPIC_TTS_AUDIO) == 0) {
    if (len < 17) return;
    char sid[16] = {0};
    memcpy(sid, payload, 16);
    if (strcmp(sid, g_sessionId) != 0) return;
    int audioSamples = (len - 16) / 2;
    if (xSemaphoreTake(xSpkMutex, pdMS_TO_TICKS(100)) == pdTRUE) {
      int copy = min(audioSamples, SPK_BUFFER_SAMPLES - (int)g_spkLen);
      memcpy((void*)(g_spkBuf + g_spkLen), payload + 16, copy * 2);
      g_spkLen += copy; g_spkReceived++;
      if (g_spkReceived >= g_spkExpected) g_spkReady = true;
      xSemaphoreGive(xSpkMutex);
    }
    return;
  }

  if (strcmp(topic, TOPIC_EMOTION_RESULT) == 0) {
    StaticJsonDocument<128> doc;
    if (deserializeJson(doc, payload, len) != DeserializationError::Ok) return;
    const char* emo = doc["emotion"];
    if (emo) {
      strncpy(g_emotion, emo, sizeof(g_emotion));
      g_emotion[sizeof(g_emotion)-1] = '\0';
      Serial.printf("[CLOUD] AI emotion: %s (%.0f%%)\n",
        g_emotion, (float)(doc["confidence"] | 0.0f) * 100);
    }
  }
}

// ================================================================
//  TASK 1 — MICROPHONE (SIMULATED)
// ================================================================
void taskMicrophone(void* param) {
  i2s_config_t micCfg = {
    .mode                 = (i2s_mode_t)(I2S_MODE_MASTER | I2S_MODE_RX),
    .sample_rate          = SAMPLE_RATE,
    .bits_per_sample      = I2S_BITS_PER_SAMPLE_32BIT,
    .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             = false
  };
  i2s_pin_config_t micPins = {
    .bck_io_num   = MIC_SCK, .ws_io_num = MIC_WS,
    .data_out_num = I2S_PIN_NO_CHANGE, .data_in_num = MIC_SD
  };
  i2s_driver_install(I2S_NUM_0, &micCfg, 0, NULL);
  i2s_set_pin(I2S_NUM_0, &micPins);
  Serial.println("[MIC] Ready (SIMULATED)");

  while (true) {
    g_audioRMS = getSimulatedRMS();

    if      (g_audioRMS > 8000) strncpy(g_emotion, "aggressive", sizeof(g_emotion));
    else if (g_audioRMS > 4000) strncpy(g_emotion, "pain",       sizeof(g_emotion));
    else if (g_audioRMS > 1500) strncpy(g_emotion, "fearful",    sizeof(g_emotion));
    else if (g_audioRMS > 400)  strncpy(g_emotion, "happy",      sizeof(g_emotion));
    else                        strncpy(g_emotion, "normal",     sizeof(g_emotion));

    if (g_audioRMS > 400 && !g_capturing && !g_captureReady) {
      if (xSemaphoreTake(xMicMutex, pdMS_TO_TICKS(50)) == pdTRUE) {
        g_captureIdx = 0; g_capturing = true;
        newSession();
        Serial.printf("[MIC] Capture started — %s\n", g_sessionId);
        xSemaphoreGive(xMicMutex);
      }
    }

    if (g_capturing) {
      if (xSemaphoreTake(xMicMutex, pdMS_TO_TICKS(50)) == pdTRUE) {
        float freq = 200.0f + (g_audioRMS / 50.0f);
        for (int i = 0; i < 64 && g_captureIdx < AUDIO_TOTAL_SAMPLES; i++) {
          g_captureBuf[g_captureIdx] = (int16_t)(g_audioRMS *
            sinf(2.0f * M_PI * freq * g_captureIdx / SAMPLE_RATE));
          g_captureIdx++;
        }
        if (g_captureIdx >= AUDIO_TOTAL_SAMPLES) {
          g_capturing = false; g_captureReady = true;
          Serial.printf("[MIC] Capture done — %d samples\n", g_captureIdx);
        }
        xSemaphoreGive(xMicMutex);
      }
    }

    Serial.printf("[MIC] RMS:%ld  Emotion:%s\n", g_audioRMS, g_emotion);
    vTaskDelay(200 / portTICK_PERIOD_MS);
  }
}

// ================================================================
//  TASK 2 — SPEAKER
// ================================================================
void playTone(float freq, int ms) {
  int total = SAMPLE_RATE * ms / 1000;
  int16_t buf[128]; int idx = 0;
  for (int i = 0; i < total; i++) {
    int16_t v = (int16_t)(8000 * sinf(2 * M_PI * freq * i / SAMPLE_RATE));
    buf[idx++] = v; buf[idx++] = v;
    if (idx == 128) { size_t w; i2s_write(I2S_NUM_0, buf, sizeof(buf), &w, portMAX_DELAY); idx = 0; }
  }
  if (idx > 0) { size_t w; i2s_write(I2S_NUM_0, buf, idx*2, &w, portMAX_DELAY); }
}

void taskSpeaker(void* param) {
  i2s_config_t spkCfg = {
    .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     = ESP_INTR_FLAG_LEVEL1,
    .dma_buf_count        = 8, .dma_buf_len = 64, .use_apll = false
  };
  i2s_pin_config_t spkPins = {
    .bck_io_num = SPK_BCLK, .ws_io_num = SPK_LRC,
    .data_out_num = SPK_DIN, .data_in_num = I2S_PIN_NO_CHANGE
  };
  i2s_driver_install(I2S_NUM_0, &spkCfg, 0, NULL);
  i2s_set_pin(I2S_NUM_0, &spkPins);
  Serial.println("[SPK] Ready");
  playTone(880, 300);
  vTaskDelay(80 / portTICK_PERIOD_MS);
  playTone(1100, 200);

  while (true) {
    if (xSemaphoreTake(xSpkMutex, pdMS_TO_TICKS(100)) == pdTRUE) {
      if (g_spkReady && g_spkLen > 0) {
        int16_t stereo[128]; int pos = 0;
        while (pos < g_spkLen) {
          int chunk = min(64, g_spkLen - pos);
          for (int i = 0; i < chunk; i++) {
            stereo[i*2] = g_spkBuf[pos+i]; stereo[i*2+1] = g_spkBuf[pos+i];
          }
          size_t w; i2s_write(I2S_NUM_0, stereo, chunk*4, &w, portMAX_DELAY);
          pos += chunk;
        }
        g_spkLen = 0; g_spkReady = false;
        Serial.println("[SPK] Playback done");
      }
      xSemaphoreGive(xSpkMutex);
    }
    vTaskDelay(50 / portTICK_PERIOD_MS);
  }
}

// ================================================================
//  TASK 3 — GPS (SIMULATED)
// ================================================================
void taskGPS(void* param) {
  gpsSerial.begin(9600, SERIAL_8N1, GPS_RX, GPS_TX);
  Serial.println("[GPS] Ready (SIMULATED)");
  int simCount = 0;
  while (true) {
    while (gpsSerial.available()) gps.encode(gpsSerial.read());
    if (gps.location.isValid()) {
      g_lat = gps.location.lat(); g_lng = gps.location.lng(); g_gpsValid = true;
      Serial.printf("[GPS] REAL — %.6f, %.6f\n", g_lat, g_lng);
    } else {
      simCount++;
      g_lat = SIM_LAT + (simCount * 0.0001f);
      g_lng = SIM_LNG + (simCount * 0.0001f);
      g_gpsValid = true;
      Serial.printf("[GPS] SIM  — %.6f, %.6f\n", g_lat, g_lng);
    }
    vTaskDelay(5000 / portTICK_PERIOD_MS);
  }
}

// ================================================================
//  TASK 4 — AUDIO STREAM
// ================================================================
void taskAudioStream(void* param) {
  vTaskDelay(8000 / portTICK_PERIOD_MS);
  while (true) {
    bool doSend = false;
    if (xSemaphoreTake(xMicMutex, pdMS_TO_TICKS(100)) == pdTRUE) {
      doSend = g_captureReady; xSemaphoreGive(xMicMutex);
    }
    if (doSend && mqtt.connected()) {
      int totalChunks = (AUDIO_TOTAL_SAMPLES + AUDIO_CHUNK_SAMPLES - 1) / AUDIO_CHUNK_SAMPLES;
      StaticJsonDocument<128> metaDoc;
      metaDoc["session"] = g_sessionId; metaDoc["total_chunks"] = totalChunks;
      metaDoc["sample_rate"] = SAMPLE_RATE; metaDoc["bits"] = 16; metaDoc["device"] = DEVICE_ID;
      char metaBuf[128]; serializeJson(metaDoc, metaBuf);
      mqtt.publish(TOPIC_AUDIO_META, metaBuf);
      vTaskDelay(100 / portTICK_PERIOD_MS);

      uint8_t chunkPayload[16 + AUDIO_CHUNK_BYTES];
      memset(chunkPayload, 0, 16);
      memcpy(chunkPayload, g_sessionId, strlen(g_sessionId));

      if (xSemaphoreTake(xMicMutex, portMAX_DELAY) == pdTRUE) {
        int sent = 0, chunkNum = 0;
        while (sent < AUDIO_TOTAL_SAMPLES) {
          int copy = min(AUDIO_CHUNK_SAMPLES, AUDIO_TOTAL_SAMPLES - sent);
          memcpy(chunkPayload + 16, g_captureBuf + sent, copy * 2);
          mqtt.publish(TOPIC_AUDIO_STREAM, chunkPayload, 16 + copy * 2, false);
          sent += copy; chunkNum++;
          vTaskDelay(30 / portTICK_PERIOD_MS);
        }
        g_captureReady = false; xSemaphoreGive(xMicMutex);
        Serial.printf("[STREAM] Done — %d chunks sent\n", chunkNum);
      }
    }
    vTaskDelay(200 / portTICK_PERIOD_MS);
  }
}

// ================================================================
//  MQTT RECONNECT (Updated for Access Management)
// ================================================================
void reconnectMQTT() {
  while (!mqtt.connected()) {
    Serial.println("[MQTT] Connecting to HiveMQ Cloud...");
    
    // Authenticating via custom username and password
    if (mqtt.connect(DEVICE_ID, MQTT_USER, MQTT_PASS)) { 
      Serial.println("[MQTT] Connected securely!");
      mqtt.subscribe(TOPIC_TTS_META);
      mqtt.subscribe(TOPIC_TTS_AUDIO);
      mqtt.subscribe(TOPIC_EMOTION_RESULT);
      Serial.println("[MQTT] Subscribed to cloud topics");
    } else {
      Serial.printf("[MQTT] Failed rc=%d — retry in 3s\n", mqtt.state());
      vTaskDelay(3000 / portTICK_PERIOD_MS);
    }
  }
}

// ================================================================
//  TASK 5 — MQTT
// ================================================================
void taskMQTT(void* param) {
  mqtt.setServer(MQTT_BROKER, MQTT_PORT);
  mqtt.setCallback(onMQTTMessage);
  mqtt.setBufferSize(MQTT_BUFFER_SIZE);
  reconnectMQTT();

  while (true) {
    if (!mqtt.connected()) reconnectMQTT();
    mqtt.loop();

    // Publish emotion
    StaticJsonDocument<256> eDoc;
    eDoc["device"] = DEVICE_ID; eDoc["emotion"] = (const char*)g_emotion;
    eDoc["audio_rms"] = g_audioRMS; eDoc["session"] = g_sessionId;
    char eBuf[256]; serializeJson(eDoc, eBuf);
    mqtt.publish(TOPIC_EMOTION, eBuf);
    Serial.printf("[MQTT] Emotion: %s\n", eBuf);

    // Publish GPS
    if (g_gpsValid) {
      StaticJsonDocument<128> gDoc;
      gDoc["device"] = DEVICE_ID;
      gDoc["lat"] = serialized(String(g_lat, 6));
      gDoc["lng"] = serialized(String(g_lng, 6));
      char gBuf[128]; serializeJson(gDoc, gBuf);
      mqtt.publish(TOPIC_GPS, gBuf);
    }

    // Publish battery
    {
      StaticJsonDocument<128> bDoc;
      bDoc["device"] = DEVICE_ID;
      bDoc["percent"]  = serialized(String(g_batPercent, 1));
      bDoc["voltage"]  = serialized(String(g_batVoltage, 2));
      bDoc["charging"] = g_isCharging;
      bDoc["standby"]  = g_isStandby;
      char bBuf[128]; serializeJson(bDoc, bBuf);
      mqtt.publish(TOPIC_BATTERY, bBuf);
      Serial.printf("[MQTT] Battery: %.1f%% %s\n", g_batPercent,
        g_isCharging ? "CHARGING" : "ON_BATTERY");
    }

    vTaskDelay(2000 / portTICK_PERIOD_MS);
  }
}

// ================================================================
//  SETUP — WiFi & Routing Parameters
// ================================================================
void setup() {
  Serial.begin(115200);
  
  // Set up secure SSL network attributes for port 8883
  wifiClient.setInsecure();
  wifiClient.setHandshakeTimeout(30); 
  /*delay(500);

  Serial.println("\n============================================");
  Serial.println("  SMART DOG COLLAR — WOKWI SIMULATION");
  Serial.println("============================================\n");

  WiFi.mode(WIFI_STA);
  WiFi.begin("Wokwi-GUEST", "");   

  Serial.print("[WIFI] Connecting");
  int tries = 0;
  while (WiFi.status() != WL_CONNECTED && tries < 20) {
    delay(500);
    Serial.print(".");
    tries++;
  }

  if (WiFi.status() == WL_CONNECTED) {
    Serial.printf("\n[WIFI] Connected — IP: %s\n", WiFi.localIP().toString().c_str());
  } else {
    Serial.println("\n[WIFI] Still connecting — MQTT task will retry");
  }

  // Bind core MQTT instance configurations
  mqtt.setServer(MQTT_BROKER, MQTT_PORT);
  mqtt.setCallback(onMQTTMessage);
  mqtt.setBufferSize(MQTT_BUFFER_SIZE); 

  xMicMutex = xSemaphoreCreateMutex();
  xSpkMutex = xSemaphoreCreateMutex();
  if (!xMicMutex || !xSpkMutex) { Serial.println("[ERROR] Mutex failed!"); while(true); }

  g_captureBuf = (int16_t*) malloc(AUDIO_TOTAL_SAMPLES * sizeof(int16_t));
  g_spkBuf     = (int16_t*) malloc(SPK_BUFFER_SAMPLES  * sizeof(int16_t));
  if (!g_captureBuf || !g_spkBuf) { Serial.println("[ERROR] Buffer alloc failed!"); while(true); }
  Serial.println("[SETUP] Buffers OK");

  // Spin up parallel FreeRTOS tasks
  xTaskCreate(taskBattery,     "BatTask",    2048, NULL, 1, NULL);
  xTaskCreate(taskMicrophone,  "MicTask",    4096, NULL, 3, NULL);
  xTaskCreate(taskSpeaker,     "SpkTask",    4096, NULL, 3, NULL);
  xTaskCreate(taskGPS,         "GPSTask",    4096, NULL, 1, NULL);
  xTaskCreate(taskAudioStream, "StreamTask", 8192, NULL, 2, NULL);
  xTaskCreate(taskMQTT,        "MQTTTask",   8192, NULL, 2, NULL);

  Serial.println("[SETUP] All 6 tasks started\n");*/
  Serial.println("[SETUP] Buffers OK");

  // Spin up parallel FreeRTOS tasks with brief delays for simulation stability
  xTaskCreate(taskBattery,     "BatTask",    2048, NULL, 1, NULL);
  vTaskDelay(100 / portTICK_PERIOD_MS);
  
  xTaskCreate(taskMicrophone,  "MicTask",    4096, NULL, 3, NULL);
  vTaskDelay(100 / portTICK_PERIOD_MS);
  
  xTaskCreate(taskSpeaker,     "SpkTask",    4096, NULL, 3, NULL);
  vTaskDelay(100 / portTICK_PERIOD_MS);
  
  xTaskCreate(taskGPS,         "GPSTask",    4096, NULL, 1, NULL);
  vTaskDelay(100 / portTICK_PERIOD_MS);
  
  xTaskCreate(taskAudioStream, "StreamTask", 8192, NULL, 2, NULL);
  vTaskDelay(100 / portTICK_PERIOD_MS);
  
  xTaskCreate(taskMQTT,        "MQTTTask",   8192, NULL, 2, NULL);

  Serial.println("[SETUP] All 6 tasks started safely\n");
}


void loop() {
  vTaskDelay(1000 / portTICK_PERIOD_MS);
}