#include <WiFi.h>
#include <WiFiClientSecure.h>
#include <PubSubClient.h>
#include <time.h>
#include "DHT.h"
#include "mbedtls/md.h"

#include "secrets.h"
#include "ca_cert.h"

// ================== PIN ==================

#define PIR_PIN      12
#define DHT_PIN      17
#define LDR_PIN      34
#define RELAY_PIN    15
#define LED_XANH     2
#define LED_DO       4

#define DHTTYPE DHT22

// ================== CẤU HÌNH ==================

const unsigned long PUBLISH_INTERVAL_MS = 5000;
const unsigned long RECONNECT_INTERVAL_MS = 5000;
const size_t MQTT_BUFFER_SIZE = 512;

DHT dht(DHT_PIN, DHTTYPE);

WiFiClientSecure secureClient;
PubSubClient mqttClient(secureClient);

// ================== DEVICE ==================

char deviceId[32];
char topicTelemetry[100];
char topicStatus[100];

unsigned long lastPublishTime = 0;
unsigned long lastReconnectAttempt = 0;

uint32_t messageSequence = 0;
bool timeReady = false;

// =====================================================
// TẠO DEVICE ID TỪ CHIP ID CỦA ESP32
// =====================================================

void createDeviceIdentity() {
  uint64_t chipId = ESP.getEfuseMac();

  snprintf(
    deviceId,
    sizeof(deviceId),
    "esp32-%04X%08X",
    (uint16_t)(chipId >> 32),
    (uint32_t)chipId
  );

  snprintf(
    topicTelemetry,
    sizeof(topicTelemetry),
    "vbank/room1/%s/telemetry",
    deviceId
  );

  snprintf(
    topicStatus,
    sizeof(topicStatus),
    "vbank/room1/%s/status",
    deviceId
  );

  Serial.println();
  Serial.println("========== THÔNG TIN THIẾT BỊ ==========");

  Serial.print("[DEVICE ID]: ");
  Serial.println(deviceId);

  Serial.print("[TELEMETRY TOPIC]: ");
  Serial.println(topicTelemetry);

  Serial.print("[STATUS TOPIC]: ");
  Serial.println(topicStatus);

  Serial.println("========================================");
}

// =====================================================
// KẾT NỐI WIFI
// =====================================================

bool setupWiFi() {
  if (WiFi.status() == WL_CONNECTED) {
    return true;
  }

  Serial.print("[WIFI]: Đang kết nối");

  WiFi.mode(WIFI_STA);
  WiFi.begin(WIFI_SSID, WIFI_PASSWORD);

  unsigned long startTime = millis();

  while (WiFi.status() != WL_CONNECTED) {
    delay(500);
    Serial.print(".");

    if (millis() - startTime >= 20000) {
      Serial.println();
      Serial.println("[WIFI]: Kết nối thất bại.");
      return false;
    }
  }

  Serial.println();
  Serial.println("[WIFI]: Kết nối thành công.");

  Serial.print("[WIFI]: IP = ");
  Serial.println(WiFi.localIP());

  return true;
}

// =====================================================
// ĐỒNG BỘ THỜI GIAN NTP
//
// Thời gian chính xác cần thiết để kiểm tra thời hạn
// chứng chỉ TLS và hỗ trợ chống Replay Attack.
// =====================================================

bool syncClock() {
  Serial.print("[TIME]: Đang đồng bộ thời gian NTP");

  configTime(
    0,
    0,
    "pool.ntp.org",
    "time.google.com",
    "time.cloudflare.com"
  );

  unsigned long startTime = millis();

  while (time(nullptr) < 1700000000) {
    delay(500);
    Serial.print(".");

    if (millis() - startTime >= 20000) {
      Serial.println();
      Serial.println("[TIME]: Đồng bộ thời gian thất bại.");
      return false;
    }
  }

  Serial.println();
  Serial.println("[TIME]: Đồng bộ thành công.");

  Serial.print("[TIME]: Unix timestamp = ");
  Serial.println((unsigned long)time(nullptr));

  return true;
}

// =====================================================
// TẠO CHỮ KÝ HMAC-SHA256
// =====================================================

String createHmacSHA256(const String& message) {
  unsigned char hmacResult[32];

  mbedtls_md_context_t context;
  mbedtls_md_init(&context);

  const mbedtls_md_info_t* mdInfo =
      mbedtls_md_info_from_type(MBEDTLS_MD_SHA256);

  if (mdInfo == nullptr) {
    Serial.println("[HMAC]: Không tìm thấy SHA-256.");
    mbedtls_md_free(&context);
    return "";
  }

  int result = mbedtls_md_setup(&context, mdInfo, 1);

  if (result != 0) {
    Serial.print("[HMAC]: Lỗi khởi tạo, mã = ");
    Serial.println(result);

    mbedtls_md_free(&context);
    return "";
  }

  result = mbedtls_md_hmac_starts(
    &context,
    reinterpret_cast<const unsigned char*>(DEVICE_HMAC_KEY),
    strlen(DEVICE_HMAC_KEY)
  );

  if (result == 0) {
    result = mbedtls_md_hmac_update(
      &context,
      reinterpret_cast<const unsigned char*>(message.c_str()),
      message.length()
    );
  }

  if (result == 0) {
    result = mbedtls_md_hmac_finish(
      &context,
      hmacResult
    );
  }

  mbedtls_md_free(&context);

  if (result != 0) {
    Serial.print("[HMAC]: Không thể tạo chữ ký, mã = ");
    Serial.println(result);
    return "";
  }

  char hexResult[65];

  for (int i = 0; i < 32; i++) {
    snprintf(
      hexResult + (i * 2),
      3,
      "%02x",
      hmacResult[i]
    );
  }

  hexResult[64] = '\0';

  return String(hexResult);
}

// =====================================================
// KẾT NỐI MQTT
// =====================================================

bool connectMQTT() {
  if (WiFi.status() != WL_CONNECTED) {
    Serial.println("[MQTT]: Wi-Fi chưa kết nối.");
    return false;
  }

  if (!timeReady) {
    Serial.println("[MQTT]: Chưa có thời gian chuẩn, chưa kết nối TLS.");
    return false;
  }

  Serial.print("[MQTT]: Đang kết nối TLS tới ");
  Serial.print(MQTT_SERVER);
  Serial.print(":");
  Serial.println(MQTT_PORT);

  /*
   * Last Will and Testament:
   * Nếu thiết bị mất mạng hoặc mất nguồn đột ngột,
   * broker sẽ tự publish trạng thái "offline".
   */
  bool connected = mqttClient.connect(
    deviceId,
    MQTT_USERNAME,
    MQTT_PASSWORD,
    topicStatus,
    1,
    true,
    "offline"
  );

  if (!connected) {
    Serial.print("[MQTT]: Kết nối thất bại. Mã lỗi = ");
    Serial.println(mqttClient.state());
    return false;
  }

  bool statusPublished = mqttClient.publish(
    topicStatus,
    "online",
    true
  );

  if (!statusPublished) {
    Serial.println("[MQTT]: Không publish được trạng thái online.");
  }

  Serial.println("[SECURITY]: Xác minh CA thành công.");
  Serial.println("[SECURITY]: TLS và MQTT authentication thành công.");
  Serial.println("[MQTT]: Thiết bị đã kết nối broker.");

  return true;
}

// =====================================================
// ĐIỀU KHIỂN THIẾT BỊ
// =====================================================

void controlDevices(
  int motion,
  int lightValue,
  float temperature
) {
  // Có người và phòng tối thì bật LED xanh
  bool turnOnLight =
      motion == HIGH &&
      lightValue > 2000;

  digitalWrite(
    LED_XANH,
    turnOnLight ? HIGH : LOW
  );

  // Nhiệt độ trên 35 độ thì bật relay và LED đỏ
  bool overTemperature =
      temperature > 35.0;

  digitalWrite(
    RELAY_PIN,
    overTemperature ? HIGH : LOW
  );

  digitalWrite(
    LED_DO,
    overTemperature ? HIGH : LOW
  );
}

// =====================================================
// GỬI DỮ LIỆU CÓ CHỮ KÝ HMAC
// =====================================================

void publishSensorData() {
  int motion = digitalRead(PIR_PIN);
  int lightValue = analogRead(LDR_PIN);

  float humidity = dht.readHumidity();
  float temperature = dht.readTemperature();

  if (isnan(humidity) || isnan(temperature)) {
    Serial.println("[SENSOR]: Không đọc được DHT22.");
    return;
  }

  controlDevices(
    motion,
    lightValue,
    temperature
  );

  messageSequence++;

  unsigned long timestamp =
      (unsigned long)time(nullptr);

  /*
   * Chuỗi chuẩn dùng để ký.
   *
   * Backend phải ghép các trường theo đúng thứ tự này:
   *
   * deviceId|seq|timestamp|temp|humi|light|motion
   */
  String signedData;
  signedData.reserve(160);

  signedData =
      String(deviceId) + "|" +
      String(messageSequence) + "|" +
      String(timestamp) + "|" +
      String(temperature, 1) + "|" +
      String(humidity, 1) + "|" +
      String(lightValue) + "|" +
      String(motion);

  String signature =
      createHmacSHA256(signedData);

  if (signature.length() == 0) {
    Serial.println("[SECURITY]: Không tạo được chữ ký HMAC.");
    return;
  }

  String payload;
  payload.reserve(320);

  payload =
      "{"
        "\"version\":1,"
        "\"deviceId\":\"" + String(deviceId) + "\","
        "\"seq\":" + String(messageSequence) + ","
        "\"timestamp\":" + String(timestamp) + ","
        "\"temp\":" + String(temperature, 1) + ","
        "\"humi\":" + String(humidity, 1) + ","
        "\"light\":" + String(lightValue) + ","
        "\"motion\":" + String(motion) + ","
        "\"sigAlg\":\"HMAC-SHA256\","
        "\"signature\":\"" + signature + "\""
      "}";

  bool published = mqttClient.publish(
    topicTelemetry,
    payload.c_str(),
    false
  );

  if (!published) {
    Serial.println("[MQTT]: Publish thất bại.");
    Serial.println("[MQTT]: Kiểm tra kết nối, quyền topic hoặc buffer.");
    return;
  }

  Serial.println();
  Serial.println("[SECURITY]: Dữ liệu đã được ký HMAC-SHA256.");

  Serial.print("[MQTTS TOPIC]: ");
  Serial.println(topicTelemetry);

  Serial.print("[MQTTS PAYLOAD]: ");
  Serial.println(payload);
}

// =====================================================
// SETUP
// =====================================================

void setup() {
  Serial.begin(115200);
  delay(500);

  pinMode(RELAY_PIN, OUTPUT);
  pinMode(LED_XANH, OUTPUT);
  pinMode(LED_DO, OUTPUT);

  pinMode(PIR_PIN, INPUT);
  pinMode(LDR_PIN, INPUT);

  // Trạng thái an toàn ban đầu
  digitalWrite(RELAY_PIN, LOW);
  digitalWrite(LED_XANH, LOW);
  digitalWrite(LED_DO, LOW);

  dht.begin();

  createDeviceIdentity();

  /*
   * Xác minh chứng chỉ broker.
   * Không dùng setInsecure().
   */
  secureClient.setCACert(HIVEMQ_ROOT_CA);

  mqttClient.setServer(
    MQTT_SERVER,
    MQTT_PORT
  );

  mqttClient.setKeepAlive(30);
  mqttClient.setSocketTimeout(15);

  /*
   * Payload có HMAC dài hơn 256 byte khi cộng thêm
   * topic và MQTT header, nên tăng buffer.
   */
  bool bufferReady =
      mqttClient.setBufferSize(MQTT_BUFFER_SIZE);

  if (!bufferReady) {
    Serial.println("[MQTT]: Không cấp phát được buffer 512 byte.");
  }

  if (setupWiFi()) {
    timeReady = syncClock();
  }

  if (timeReady) {
    connectMQTT();
  }
}

// =====================================================
// LOOP
// =====================================================

void loop() {
  unsigned long currentTime = millis();

  // Khôi phục Wi-Fi khi mất mạng
  if (WiFi.status() != WL_CONNECTED) {
    if (
      currentTime - lastReconnectAttempt >=
      RECONNECT_INTERVAL_MS
    ) {
      lastReconnectAttempt = currentTime;

      Serial.println("[WIFI]: Mất kết nối, đang kết nối lại...");

      if (setupWiFi()) {
        timeReady = syncClock();
      }
    }

    return;
  }

  // Đảm bảo đã đồng bộ thời gian trước TLS
  if (!timeReady) {
    timeReady = syncClock();

    if (!timeReady) {
      delay(1000);
      return;
    }
  }

  // Khôi phục MQTT
  if (!mqttClient.connected()) {
    if (
      currentTime - lastReconnectAttempt >=
      RECONNECT_INTERVAL_MS
    ) {
      lastReconnectAttempt = currentTime;
      connectMQTT();
    }

    return;
  }

  mqttClient.loop();

  // Gửi dữ liệu mỗi 5 giây
  if (
    currentTime - lastPublishTime >=
    PUBLISH_INTERVAL_MS
  ) {
    lastPublishTime = currentTime;
    publishSensorData();
  }
}