#include "xiaomi_ble.h"
#include "esphome/core/helpers.h"
#include "esphome/core/log.h"

#ifdef USE_ESP32

#include <vector>
#include <esp_idf_version.h>
#if ESP_IDF_VERSION >= ESP_IDF_VERSION_VAL(6, 0, 0)
#include <psa/crypto.h>
#else
#include "mbedtls/ccm.h"
#endif

namespace esphome::xiaomi_ble {

static const char *const TAG = "xiaomi_ble";

// Maximum bytes to log in very verbose hex output (covers largest packet of ~24 bytes)
static constexpr size_t XIAOMI_MAX_LOG_BYTES = 32;

bool parse_xiaomi_value(uint16_t value_type, const uint8_t *data, uint8_t value_length, XiaomiParseResult &result) {
  // button pressed, 3 bytes, only byte 3 is used for supported devices so far
  if ((value_type == 0x1001) && (value_length == 3)) {
    result.button_press = data[2] == 0;
    return true;
  }
  // motion detection, 1 byte, 8-bit unsigned integer
  else if ((value_type == 0x0003) && (value_length == 1)) {
    result.has_motion = data[0];
  }
  // temperature, 2 bytes, 16-bit signed integer (LE), 0.1 °C
  else if ((value_type == 0x1004) && (value_length == 2)) {
    const int16_t temperature = encode_uint16(data[1], data[0]);
    result.temperature = temperature / 10.0f;
  }
  // humidity, 2 bytes, 16-bit signed integer (LE), 0.1 %
  else if ((value_type == 0x1006) && (value_length == 2)) {
    const int16_t humidity = encode_uint16(data[1], data[0]);
    result.humidity = humidity / 10.0f;
  }
  // illuminance (+ motion), 3 bytes, 24-bit unsigned integer (LE), 1 lx
  else if (((value_type == 0x1007) || (value_type == 0x000F)) && (value_length == 3)) {
    const uint32_t illuminance = encode_uint24(data[2], data[1], data[0]);
    result.illuminance = illuminance;
    result.is_light = illuminance >= 100;
    if (value_type == 0x0F)
      result.has_motion = true;
  }
  // soil moisture, 1 byte, 8-bit unsigned integer, 1 %
  else if ((value_type == 0x1008) && (value_length == 1)) {
    result.moisture = data[0];
  }
  // conductivity, 2 bytes, 16-bit unsigned integer (LE), 1 µS/cm
  else if ((value_type == 0x1009) && (value_length == 2)) {
    const uint16_t conductivity = encode_uint16(data[1], data[0]);
    result.conductivity = conductivity;
  }
  // battery / MiaoMiaoce battery, 1 byte, 8-bit unsigned integer, 1 %
  else if ((value_type == 0x100A || value_type == 0x4803) && (value_length == 1)) {
    result.battery_level = data[0];
  }
  // temperature + humidity, 4 bytes, 16-bit signed integer (LE) each, 0.1 °C, 0.1 %
  else if ((value_type == 0x100D) && (value_length == 4)) {
    const int16_t temperature = encode_uint16(data[1], data[0]);
    const int16_t humidity = encode_uint16(data[3], data[2]);
    result.temperature = temperature / 10.0f;
    result.humidity = humidity / 10.0f;
  }
  // formaldehyde, 2 bytes, 16-bit unsigned integer (LE), 0.01 mg / m3
  else if ((value_type == 0x1010) && (value_length == 2)) {
    const uint16_t formaldehyde = encode_uint16(data[1], data[0]);
    result.formaldehyde = formaldehyde / 100.0f;
  }
  // on/off state, 1 byte, 8-bit unsigned integer
  else if ((value_type == 0x1012) && (value_length == 1)) {
    result.is_active = data[0];
  }
  // mosquito tablet, 1 byte, 8-bit unsigned integer, 1 %
  else if ((value_type == 0x1013) && (value_length == 1)) {
    result.tablet = data[0];
  }
  // idle time since last motion, 4 byte, 32-bit unsigned integer, 1 min
  else if ((value_type == 0x1017) && (value_length == 4)) {
    const uint32_t idle_time = encode_uint32(data[3], data[2], data[1], data[0]);
    result.idle_time = idle_time / 60.0f;
    result.has_motion = !idle_time;
  } else if ((value_type == 0x1018) && (value_length == 1)) {
    result.is_light = data[0];
  }
  // MiaoMiaoce temperature, 4 bytes, float, 0.1 °C
  else if ((value_type == 0x4C01) && (value_length == 4)) {
    const uint32_t int_number = encode_uint32(data[3], data[2], data[1], data[0]);
    float temperature;
    std::memcpy(&temperature, &int_number, sizeof(temperature));
    result.temperature = temperature;
  }
  // MiaoMiaoce humidity, 1 byte, 8-bit unsigned integer, 1 %
  else if ((value_type == 0x4C02) && (value_length == 1)) {
    result.humidity = data[0];
  }
  // XMWSDJ04MMC humidity, 4 bytes, float, 0.1 °C
  else if ((value_type == 0x4C08) && (value_length == 4)) {
    const uint32_t int_number = encode_uint32(data[3], data[2], data[1], data[0]);
    float humidity;
    std::memcpy(&humidity, &int_number, sizeof(humidity));
    result.humidity = humidity;
  } else {
    return false;
  }

  return true;
}

bool parse_xiaomi_message(const std::vector<uint8_t> &message, XiaomiParseResult &result) {
  result.has_encryption = message[0] & 0x08;  // update encryption status
  if (result.has_encryption) {
    ESP_LOGVV(TAG, "parse_xiaomi_message(): payload is encrypted, stop reading message.");
    return false;
  }

  // Data point specs
  // Byte 0: type
  // Byte 1: fixed 0x10
  // Byte 2: length
  // Byte 3..3+len-1: data point value

  if (result.raw_offset < 0 || static_cast<size_t>(result.raw_offset) >= message.size()) {
    ESP_LOGVV(TAG, "parse_xiaomi_message(): raw_offset (%d) exceeds message size (%d)!", result.raw_offset,
              message.size());
    return false;
  }
  const uint8_t *payload = message.data() + result.raw_offset;
  uint8_t payload_length = message.size() - result.raw_offset;
  uint8_t payload_offset = 0;
  bool success = false;

  if (payload_length < 4) {
    ESP_LOGVV(TAG, "parse_xiaomi_message(): payload has wrong size (%d)!", payload_length);
    return false;
  }

  while (payload_length > 3) {
    if (payload[payload_offset + 1] != 0x10 && payload[payload_offset + 1] != 0x00 &&
        payload[payload_offset + 1] != 0x4C && payload[payload_offset + 1] != 0x48) {
      ESP_LOGVV(TAG, "parse_xiaomi_message(): fixed byte not found, stop parsing residual data.");
      break;
    }

    const uint8_t value_length = payload[payload_offset + 2];
    if ((value_length < 1) || (value_length > 4) || (payload_length < (3 + value_length))) {
      ESP_LOGVV(TAG, "parse_xiaomi_message(): value has wrong size (%d)!", value_length);
      break;
    }

    const uint16_t value_type = encode_uint16(payload[payload_offset + 1], payload[payload_offset + 0]);
    const uint8_t *data = &payload[payload_offset + 3];

    if (parse_xiaomi_value(value_type, data, value_length, result))
      success = true;

    payload_length -= 3 + value_length;
    payload_offset += 3 + value_length;
  }

  return success;
}

optional<XiaomiParseResult> parse_xiaomi_header(const esp32_ble_tracker::ServiceData &service_data) {
  XiaomiParseResult result;
  if (!service_data.uuid.contains(0x95, 0xFE)) {
    ESP_LOGVV(TAG, "parse_xiaomi_header(): no service data UUID magic bytes.");
    return {};
  }

  auto raw = service_data.data;
  if (raw.size() < 5) {
    ESP_LOGVV(TAG, "parse_xiaomi_header(): service data too short (%d).", raw.size());
    return {};
  }
  result.has_data = raw[0] & 0x40;
  result.has_capability = raw[0] & 0x20;
  result.has_encryption = raw[0] & 0x08;

  if (!result.has_data) {
    ESP_LOGVV(TAG, "parse_xiaomi_header(): service data has no DATA flag.");
    return {};
  }

  static uint8_t last_frame_count = 0;
  if (last_frame_count == raw[4]) {
    ESP_LOGVV(TAG, "parse_xiaomi_header(): duplicate data packet received (%d).", static_cast<int>(last_frame_count));
    result.is_duplicate = true;
    return {};
  }
  last_frame_count = raw[4];
  result.is_duplicate = false;
  result.raw_offset = result.has_capability ? 12 : 11;

  const uint16_t device_uuid = encode_uint16(raw[3], raw[2]);

  if (device_uuid == 0x0098) {  // MiFlora
    result.type = XiaomiParseResult::TYPE_HHCCJCY01;
    result.name = "HHCCJCY01";
  } else if (device_uuid == 0x01aa) {  // round body, segment LCD
    result.type = XiaomiParseResult::TYPE_LYWSDCGQ;
    result.name = "LYWSDCGQ";
  } else if (device_uuid == 0x015d) {  // FlowerPot, RoPot
    result.type = XiaomiParseResult::TYPE_HHCCPOT002;
    result.name = "HHCCPOT002";
  } else if (device_uuid == 0x02df) {  // Xiaomi (Honeywell) formaldehyde sensor, OLED display
    result.type = XiaomiParseResult::TYPE_JQJCY01YM;
    result.name = "JQJCY01YM";
  } else if (device_uuid == 0x03dd) {  // Philips/Xiaomi BLE nightlight
    result.type = XiaomiParseResult::TYPE_MUE4094RT;
    result.name = "MUE4094RT";
    result.raw_offset -= 6;
  } else if (device_uuid == 0x0347 ||  // ClearGrass-branded, round body, e-ink display
             device_uuid == 0x0B48) {  // Qingping-branded, round body, e-ink display — with bindkeys
    result.type = XiaomiParseResult::TYPE_CGG1;
    result.name = "CGG1";
  } else if (device_uuid == 0x03bc) {  // VegTrug Grow Care Garden
    result.type = XiaomiParseResult::TYPE_GCLS002;
    result.name = "GCLS002";
  } else if (device_uuid == 0x045b) {  // rectangular body, e-ink display
    result.type = XiaomiParseResult::TYPE_LYWSD02;
    result.name = "LYWSD02";
  } else if (device_uuid == 0x2542) {  // rectangular body, e-ink display — with bindkeys
    result.type = XiaomiParseResult::TYPE_LYWSD02MMC;
    result.name = "LYWSD02MMC";
    if (raw.size() == 19)
      result.raw_offset -= 6;
  } else if (device_uuid == 0x040a) {  // Mosquito Repellent Smart Version
    result.type = XiaomiParseResult::TYPE_WX08ZM;
    result.name = "WX08ZM";
  } else if (device_uuid == 0x0576) {  // Cleargrass (Qingping) alarm clock, segment LCD
    result.type = XiaomiParseResult::TYPE_CGD1;
    result.name = "CGD1";
  } else if (device_uuid == 0x066F) {  // Cleargrass (Qingping) Temp & RH Lite
    result.type = XiaomiParseResult::TYPE_CGDK2;
    result.name = "CGDK2";
  } else if (device_uuid == 0x055b) {  // small square body, segment LCD, encrypted
    result.type = XiaomiParseResult::TYPE_LYWSD03MMC;
    result.name = "LYWSD03MMC";
  } else if (device_uuid == 0x1203) {  // small square body, e-ink display, encrypted
    result.type = XiaomiParseResult::TYPE_XMWSDJ04MMC;
    result.name = "XMWSDJ04MMC";
    if (raw.size() == 19)
      result.raw_offset -= 6;
  } else if (device_uuid == 0x07f6) {  // Xiaomi-Yeelight BLE nightlight
    result.type = XiaomiParseResult::TYPE_MJYD02YLA;
    result.name = "MJYD02YLA";
    if (raw.size() == 19)
      result.raw_offset -= 6;
  } else if (device_uuid == 0x06d3) {  // rectangular body, e-ink display with alarm
    result.type = XiaomiParseResult::TYPE_MHOC303;
    result.name = "MHOC303";
  } else if (device_uuid == 0x0387) {  // square body, e-ink display
    result.type = XiaomiParseResult::TYPE_MHOC401;
    result.name = "MHOC401";
  } else if (device_uuid == 0x0A83) {  // Qingping-branded, motion & ambient light sensor
    result.type = XiaomiParseResult::TYPE_CGPR1;
    result.name = "CGPR1";
    if (raw.size() == 19)
      result.raw_offset -= 6;
  } else if (device_uuid == 0x0A8D) {  // Xiaomi Mi Motion Sensor 2
    result.type = XiaomiParseResult::TYPE_RTCGQ02LM;
    result.name = "RTCGQ02LM";
    if (raw.size() == 19)
      result.raw_offset -= 6;
  } else {
    ESP_LOGVV(TAG, "parse_xiaomi_header(): unknown device, no magic bytes.");
    return {};
  }

  return result;
}

bool decrypt_xiaomi_payload(std::vector<uint8_t> &raw, const uint8_t *bindkey, const uint64_t &address) {
  if ((raw.size() != 19) && ((raw.size() < 22) || (raw.size() > 24))) {
    ESP_LOGVV(TAG, "decrypt_xiaomi_payload(): data packet has wrong size (%d)!", raw.size());
#if ESPHOME_LOG_LEVEL >= ESPHOME_LOG_LEVEL_VERY_VERBOSE
    char hex_buf[format_hex_pretty_size(XIAOMI_MAX_LOG_BYTES)];
#endif
    ESP_LOGVV(TAG, "  Packet : %s", format_hex_pretty_to(hex_buf, raw.data(), raw.size()));
    return false;
  }

  uint8_t mac_reverse[6] = {0};
  mac_reverse[5] = (uint8_t) (address >> 40);
  mac_reverse[4] = (uint8_t) (address >> 32);
  mac_reverse[3] = (uint8_t) (address >> 24);
  mac_reverse[2] = (uint8_t) (address >> 16);
  mac_reverse[1] = (uint8_t) (address >> 8);
  mac_reverse[0] = (uint8_t) (address >> 0);

  XiaomiAESVector vector{.key = {0},
                         .plaintext = {0},
                         .ciphertext = {0},
                         .authdata = {0x11},
                         .iv = {0},
                         .tag = {0},
                         .keysize = 16,
                         .authsize = 1,
                         .datasize = 0,
                         .tagsize = 4,
                         .ivsize = 12};

  vector.datasize = (raw.size() == 19) ? raw.size() - 12 : raw.size() - 18;
  int cipher_pos = (raw.size() == 19) ? 5 : 11;

  const uint8_t *v = raw.data();

  memcpy(vector.key, bindkey, vector.keysize);
  memcpy(vector.ciphertext, v + cipher_pos, vector.datasize);
  memcpy(vector.tag, v + raw.size() - vector.tagsize, vector.tagsize);
  memcpy(vector.iv, mac_reverse, 6);             // MAC address reverse
  memcpy(vector.iv + 6, v + 2, 3);               // sensor type (2) + packet id (1)
  memcpy(vector.iv + 9, v + raw.size() - 7, 3);  // payload counter

#if ESP_IDF_VERSION >= ESP_IDF_VERSION_VAL(6, 0, 0)
  // PSA AEAD expects ciphertext + tag concatenated
  uint8_t ct_with_tag[sizeof(vector.ciphertext) + sizeof(vector.tag)];
  memcpy(ct_with_tag, vector.ciphertext, vector.datasize);
  memcpy(ct_with_tag + vector.datasize, vector.tag, vector.tagsize);
  size_t ct_with_tag_size = vector.datasize + vector.tagsize;

  psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
  psa_set_key_type(&attributes, PSA_KEY_TYPE_AES);
  psa_set_key_bits(&attributes, vector.keysize * 8);
  psa_set_key_usage_flags(&attributes, PSA_KEY_USAGE_DECRYPT);
  psa_set_key_algorithm(&attributes, PSA_ALG_AEAD_WITH_SHORTENED_TAG(PSA_ALG_CCM, vector.tagsize));

  mbedtls_svc_key_id_t key_id;
  if (psa_import_key(&attributes, vector.key, vector.keysize, &key_id) != PSA_SUCCESS) {
    ESP_LOGVV(TAG, "decrypt_xiaomi_payload(): psa_import_key() failed.");
    return false;
  }

  size_t plaintext_length;
  psa_status_t status = psa_aead_decrypt(key_id, PSA_ALG_AEAD_WITH_SHORTENED_TAG(PSA_ALG_CCM, vector.tagsize),
                                         vector.iv, vector.ivsize, vector.authdata, vector.authsize, ct_with_tag,
                                         ct_with_tag_size, vector.plaintext, vector.datasize, &plaintext_length);
  psa_destroy_key(key_id);
  bool decrypt_ok = (status == PSA_SUCCESS && plaintext_length == vector.datasize);
#else
  mbedtls_ccm_context ctx;
  mbedtls_ccm_init(&ctx);

  int ret = mbedtls_ccm_setkey(&ctx, MBEDTLS_CIPHER_ID_AES, vector.key, vector.keysize * 8);
  if (ret) {
    ESP_LOGVV(TAG, "decrypt_xiaomi_payload(): mbedtls_ccm_setkey() failed.");
    mbedtls_ccm_free(&ctx);
    return false;
  }

  ret = mbedtls_ccm_auth_decrypt(&ctx, vector.datasize, vector.iv, vector.ivsize, vector.authdata, vector.authsize,
                                 vector.ciphertext, vector.plaintext, vector.tag, vector.tagsize);
  mbedtls_ccm_free(&ctx);
  bool decrypt_ok = (ret == 0);
#endif

  if (!decrypt_ok) {
    uint8_t mac_address[6] = {0};
    memcpy(mac_address, mac_reverse + 5, 1);
    memcpy(mac_address + 1, mac_reverse + 4, 1);
    memcpy(mac_address + 2, mac_reverse + 3, 1);
    memcpy(mac_address + 3, mac_reverse + 2, 1);
    memcpy(mac_address + 4, mac_reverse + 1, 1);
    memcpy(mac_address + 5, mac_reverse, 1);
    ESP_LOGVV(TAG, "decrypt_xiaomi_payload(): authenticated decryption failed.");
#if ESPHOME_LOG_LEVEL >= ESPHOME_LOG_LEVEL_VERY_VERBOSE
    char mac_buf[MAC_ADDRESS_PRETTY_BUFFER_SIZE];
    format_mac_addr_upper(mac_address, mac_buf);
    char hex_buf[format_hex_pretty_size(XIAOMI_MAX_LOG_BYTES)];
#endif
    ESP_LOGVV(TAG, "  MAC address : %s", mac_buf);
    ESP_LOGVV(TAG, "       Packet : %s", format_hex_pretty_to(hex_buf, raw.data(), raw.size()));
    ESP_LOGVV(TAG, "          Key : %s", format_hex_pretty_to(hex_buf, vector.key, vector.keysize));
    ESP_LOGVV(TAG, "           Iv : %s", format_hex_pretty_to(hex_buf, vector.iv, vector.ivsize));
    ESP_LOGVV(TAG, "       Cipher : %s", format_hex_pretty_to(hex_buf, vector.ciphertext, vector.datasize));
    ESP_LOGVV(TAG, "          Tag : %s", format_hex_pretty_to(hex_buf, vector.tag, vector.tagsize));
    return false;
  }

  // replace encrypted payload with plaintext
  uint8_t *p = vector.plaintext;
  for (std::vector<uint8_t>::iterator it = raw.begin() + cipher_pos; it != raw.begin() + cipher_pos + vector.datasize;
       ++it) {
    *it = *(p++);
  }

  // clear encrypted flag
  raw[0] &= ~0x08;

  ESP_LOGVV(TAG, "decrypt_xiaomi_payload(): authenticated decryption passed.");
#if ESPHOME_LOG_LEVEL >= ESPHOME_LOG_LEVEL_VERY_VERBOSE
  char hex_buf[format_hex_pretty_size(XIAOMI_MAX_LOG_BYTES)];
#endif
  ESP_LOGVV(TAG, "  Plaintext : %s, Packet : %d",
            format_hex_pretty_to(hex_buf, raw.data() + cipher_pos, vector.datasize), static_cast<int>(raw[4]));

  return true;
}

bool report_xiaomi_results(const optional<XiaomiParseResult> &result, const char *address) {
  if (!result.has_value()) {
    ESP_LOGVV(TAG, "report_xiaomi_results(): no results available.");
    return false;
  }

  ESP_LOGD(TAG, "Got Xiaomi %s (%s):", result->name.c_str(), address);

  if (result->temperature.has_value()) {
    ESP_LOGD(TAG, "  Temperature: %.1f°C", *result->temperature);
  }
  if (result->humidity.has_value()) {
    ESP_LOGD(TAG, "  Humidity: %.1f%%", *result->humidity);
  }
  if (result->battery_level.has_value()) {
    ESP_LOGD(TAG, "  Battery Level: %.0f%%", *result->battery_level);
  }
  if (result->conductivity.has_value()) {
    ESP_LOGD(TAG, "  Conductivity: %.0fµS/cm", *result->conductivity);
  }
  if (result->illuminance.has_value()) {
    ESP_LOGD(TAG, "  Illuminance: %.0flx", *result->illuminance);
  }
  if (result->moisture.has_value()) {
    ESP_LOGD(TAG, "  Moisture: %.0f%%", *result->moisture);
  }
  if (result->tablet.has_value()) {
    ESP_LOGD(TAG, "  Mosquito tablet: %.0f%%", *result->tablet);
  }
  if (result->is_active.has_value()) {
    ESP_LOGD(TAG, "  Repellent: %s", (*result->is_active) ? "on" : "off");
  }
  if (result->has_motion.has_value()) {
    ESP_LOGD(TAG, "  Motion: %s", (*result->has_motion) ? "yes" : "no");
  }
  if (result->is_light.has_value()) {
    ESP_LOGD(TAG, "  Light: %s", (*result->is_light) ? "on" : "off");
  }
  if (result->button_press.has_value()) {
    ESP_LOGD(TAG, "  Button: %s", (*result->button_press) ? "pressed" : "");
  }

  return true;
}

bool XiaomiListener::parse_device(const esp32_ble_tracker::ESPBTDevice &device) {
  // Previously the message was parsed twice per packet, once by XiaomiListener::parse_device()
  // and then again by the respective device class's parse_device() function. Parsing the header
  // here and then for each device seems to be unnecessary and complicates the duplicate packet filtering.
  // Hence I disabled the call to parse_xiaomi_header() here and the message parsing is done entirely
  // in the respective device instance. The XiaomiListener class is defined in __init__.py and I was not
  // able to remove it entirely.

  return false;  // with true it's not showing device scans
}

}  // namespace esphome::xiaomi_ble

#endif
