#include "esphome/core/log.h"
#include "esphome/core/application.h"
#include "esphome/core/helpers.h"
#include "packet_transport.h"

#include <ranges>

#include "esphome/components/xxtea/xxtea.h"

namespace esphome {
namespace packet_transport {

// Maximum bytes to log in hex output (168 * 3 = 504, under TX buffer size of 512)
static constexpr size_t PACKET_MAX_LOG_BYTES = 168;
/**
 * Structure of a data packet; everything is little-endian
 *
 * --- In clear text ---
 * MAGIC_NUMBER: 16 bits
 * host name length: 1 byte
 * host name: (length) bytes
 * padding: 0 or more null bytes to a 4 byte boundary
 *
 * --- Encrypted (if key set) ----
 * DATA_KEY: 1 byte: OR ROLLING_CODE_KEY:
 *  Rolling code (if enabled): 8 bytes
 * Ping keys: if any
 * repeat:
 *      PING_KEY: 1 byte
 *      ping code: 4 bytes
 * Sensors:
 * repeat:
 *      SENSOR_KEY: 1 byte
 *      float value: 4 bytes
 *      name length: 1 byte
 *      name
 * Binary Sensors:
 * repeat:
 *      BINARY_SENSOR_KEY: 1 byte
 *      bool value: 1 bytes
 *      name length: 1 byte
 *      name
 *
 * Padded to a 4 byte boundary with nulls
 *
 * Structure of a ping request packet:
 * --- In clear text ---
 * MAGIC_PING: 16 bits
 * host name length: 1 byte
 * host name: (length) bytes
 * Ping key (4 bytes)
 *
 */
static const char *const TAG = "packet_transport";

static size_t round4(size_t value) { return (value + 3) & ~3; }

union FuData {
  uint32_t u32;
  float f32;
};

static constexpr uint16_t MAGIC_NUMBER = 0x4553;
static constexpr uint16_t MAGIC_PING = 0x5048;
static constexpr uint32_t PREF_HASH = 0x45535043;
enum DataKey {
  ZERO_FILL_KEY,
  DATA_KEY,
  SENSOR_KEY,
  BINARY_SENSOR_KEY,
  PING_KEY,
  ROLLING_CODE_KEY,
};

enum DecodeResult {
  DECODE_OK,
  DECODE_UNMATCHED,
  DECODE_ERROR,
  DECODE_EMPTY,
};

static constexpr size_t MAX_PING_KEYS = 4;

static inline void add(std::vector<uint8_t> &vec, uint32_t data) {
  vec.push_back(data & 0xFF);
  vec.push_back((data >> 8) & 0xFF);
  vec.push_back((data >> 16) & 0xFF);
  vec.push_back((data >> 24) & 0xFF);
}

class PacketDecoder {
 public:
  PacketDecoder(const uint8_t *buffer, size_t len) : buffer_(buffer), len_(len) {}

  DecodeResult decode_string(char *data, size_t maxlen) {
    if (this->position_ == this->len_)
      return DECODE_EMPTY;
    auto len = this->buffer_[this->position_];
    if (len == 0 || this->position_ + 1 + len > this->len_ || len >= maxlen)
      return DECODE_ERROR;
    this->position_++;
    memcpy(data, this->buffer_ + this->position_, len);
    data[len] = 0;
    this->position_ += len;
    return DECODE_OK;
  }

  template<typename T> DecodeResult get(T &data) {
    if (this->position_ + sizeof(T) > this->len_)
      return DECODE_ERROR;
    T value = 0;
    for (size_t i = 0; i != sizeof(T); ++i) {
      value += this->buffer_[this->position_++] << (i * 8);
    }
    data = value;
    return DECODE_OK;
  }

  template<typename T> DecodeResult decode(uint8_t key, T &data) {
    if (this->position_ == this->len_)
      return DECODE_EMPTY;
    if (this->buffer_[this->position_] != key)
      return DECODE_UNMATCHED;
    if (this->position_ + 1 + sizeof(T) > this->len_)
      return DECODE_ERROR;
    this->position_++;
    T value = 0;
    for (size_t i = 0; i != sizeof(T); ++i) {
      value += this->buffer_[this->position_++] << (i * 8);
    }
    data = value;
    return DECODE_OK;
  }

  template<typename T> DecodeResult decode(uint8_t key, char *buf, size_t buflen, T &data) {
    if (this->position_ == this->len_)
      return DECODE_EMPTY;
    if (this->buffer_[this->position_] != key)
      return DECODE_UNMATCHED;
    if (this->position_ + 1 + sizeof(T) > this->len_)
      return DECODE_ERROR;
    this->position_++;
    T value = 0;
    for (size_t i = 0; i != sizeof(T); ++i) {
      value += this->buffer_[this->position_++] << (i * 8);
    }
    data = value;
    return this->decode_string(buf, buflen);
  }

  DecodeResult decode(uint8_t key) {
    if (this->position_ == this->len_)
      return DECODE_EMPTY;
    if (this->buffer_[this->position_] != key)
      return DECODE_UNMATCHED;
    this->position_++;
    return DECODE_OK;
  }

  size_t get_remaining_size() const { return this->len_ - this->position_; }

  // align the pointer to the given byte boundary
  bool bump_to(size_t boundary) {
    auto newpos = this->position_;
    auto offset = this->position_ % boundary;
    if (offset != 0) {
      newpos += boundary - offset;
    }
    if (newpos >= this->len_)
      return false;
    this->position_ = newpos;
    return true;
  }

  bool decrypt(const uint32_t *key) const {
    if (this->get_remaining_size() % 4 != 0) {
      return false;
    }
    xxtea::decrypt((uint32_t *) (this->buffer_ + this->position_), this->get_remaining_size() / 4, key);
    return true;
  }

 protected:
  const uint8_t *buffer_;
  size_t len_;
  size_t position_{};
};

static inline void add(std::vector<uint8_t> &vec, uint8_t data) { vec.push_back(data); }
static inline void add(std::vector<uint8_t> &vec, uint16_t data) {
  vec.push_back((uint8_t) data);
  vec.push_back((uint8_t) (data >> 8));
}
static inline void add(std::vector<uint8_t> &vec, DataKey data) { vec.push_back(data); }
static void add(std::vector<uint8_t> &vec, const char *str) {
  auto len = strlen(str);
  vec.push_back(len);
  for (size_t i = 0; i != len; i++) {
    vec.push_back(*str++);
  }
}

void PacketTransport::setup() {
  this->name_ = App.get_name().c_str();
  if (strlen(this->name_) > 255) {
    this->status_set_error(LOG_STR("Device name exceeds 255 chars"));
    this->mark_failed();
    return;
  }
  this->resend_ping_key_ = this->ping_pong_enable_;
  this->pref_ = global_preferences->make_preference<uint32_t>(PREF_HASH, true);
  if (this->rolling_code_enable_) {
    // restore the upper 32 bits of the rolling code, increment and save.
    this->pref_.load(&this->rolling_code_[1]);
    this->rolling_code_[1]++;
    this->pref_.save(&this->rolling_code_[1]);
    // must make sure it's saved immediately
    global_preferences->sync();
    this->ping_key_ = random_uint32();
    ESP_LOGV(TAG, "Rolling code incremented, upper part now %u", (unsigned) this->rolling_code_[1]);
  }
#ifdef USE_SENSOR
  for (auto &sensor : this->sensors_) {
    // [&sensor] is safe: sensor refers to a FixedVector element that never reallocates,
    // so the reference remains valid for the component's lifetime.
    sensor.sensor->add_on_state_callback([&sensor](float x) {
      sensor.parent->updated_ = true;
      sensor.updated = true;
    });
  }
#endif
#ifdef USE_BINARY_SENSOR
  for (auto &sensor : this->binary_sensors_) {
    // [&sensor] is safe: sensor refers to a FixedVector element that never reallocates,
    // so the reference remains valid for the component's lifetime.
    sensor.sensor->add_on_state_callback([&sensor](bool value) {
      sensor.parent->updated_ = true;
      sensor.updated = true;
    });
  }
#endif
  // initialise the header. This is invariant.
  add(this->header_, MAGIC_NUMBER);
  add(this->header_, this->name_);
  // pad to a multiple of 4 bytes
  while (this->header_.size() & 0x3)
    this->header_.push_back(0);
}

void PacketTransport::init_data_() {
  this->data_.clear();
  if (this->rolling_code_enable_) {
    add(this->data_, ROLLING_CODE_KEY);
    add(this->data_, this->rolling_code_[0]);
    add(this->data_, this->rolling_code_[1]);
    this->increment_code_();
  } else {
    add(this->data_, DATA_KEY);
  }
  for (auto &value : this->ping_keys_ | std::views::values) {
    add(this->data_, PING_KEY);
    add(this->data_, value);
  }
}

void PacketTransport::flush_() {
  if (!this->should_send() || this->data_.empty())
    return;
  auto header_len = round4(this->header_.size());
  auto len = round4(data_.size());
  auto encode_buffer = std::vector<uint8_t>(round4(header_len + len));
  memcpy(encode_buffer.data(), this->header_.data(), this->header_.size());
  memcpy(encode_buffer.data() + header_len, this->data_.data(), this->data_.size());
  if (this->is_encrypted_()) {
    xxtea::encrypt((uint32_t *) (encode_buffer.data() + header_len), len / 4,
                   (uint32_t *) this->encryption_key_.data());
  }
  char hex_buf[format_hex_pretty_size(PACKET_MAX_LOG_BYTES)];
  ESP_LOGVV(TAG, "Sending packet %s", format_hex_pretty_to(hex_buf, encode_buffer.data(), encode_buffer.size()));
  this->send_packet(encode_buffer);
}

void PacketTransport::add_binary_data_(uint8_t key, const char *id, bool data) {
  auto len = 1 + 1 + 1 + strlen(id);
  if (round4(this->header_.size()) + round4(this->data_.size() + len) > this->get_max_packet_size()) {
    this->flush_();
    this->init_data_();
  }
  add(this->data_, key);
  add(this->data_, (uint8_t) data);
  add(this->data_, id);
}
void PacketTransport::add_data_(uint8_t key, const char *id, float data) {
  FuData udata{.f32 = data};
  this->add_data_(key, id, udata.u32);
}

void PacketTransport::add_data_(uint8_t key, const char *id, uint32_t data) {
  auto len = 4 + 1 + 1 + strlen(id);
  if (round4(this->header_.size()) + round4(this->data_.size() + len) > this->get_max_packet_size()) {
    this->flush_();
    this->init_data_();
  }
  add(this->data_, key);
  add(this->data_, data);
  add(this->data_, id);
}
void PacketTransport::send_data_(bool all) {
  if (!this->should_send())
    return;
  this->init_data_();
#ifdef USE_SENSOR
  for (auto &sensor : this->sensors_) {
    if (all || sensor.updated) {
      sensor.updated = false;
      this->add_data_(SENSOR_KEY, sensor.id, sensor.sensor->get_state());
    }
  }
#endif
#ifdef USE_BINARY_SENSOR
  for (auto &sensor : this->binary_sensors_) {
    if (all || sensor.updated) {
      sensor.updated = false;
      this->add_binary_data_(BINARY_SENSOR_KEY, sensor.id, sensor.sensor->state);
    }
  }
#endif
  this->flush_();
  this->updated_ = false;
}

void PacketTransport::update() {
  // resend all sensors if required
  if (this->is_provider_)
    this->send_data_(true);
  if (!this->ping_pong_enable_) {
    return;
  }
  uint32_t now = millis();
  uint32_t ping_request_age = now - this->last_key_time_;
  if (ping_request_age > this->ping_pong_recyle_time_ * 1000u) {
    this->resend_ping_key_ = this->ping_pong_enable_;
    ESP_LOGV(TAG, "Ping request, age %" PRIu32, ping_request_age);
    this->last_key_time_ = now;
  }
  for (const auto &provider : this->providers_) {
    uint32_t key_response_age = now - provider.second.last_key_response_time;
    if (key_response_age > (this->ping_pong_recyle_time_ * 2000u)) {
#ifdef USE_STATUS_SENSOR
      if (provider.second.status_sensor != nullptr && provider.second.status_sensor->state) {
        ESP_LOGI(TAG, "Ping status for %s timeout at %" PRIu32 " with age %" PRIu32, provider.first.c_str(), now,
                 key_response_age);
        provider.second.status_sensor->publish_state(false);
      }
#endif
#ifdef USE_SENSOR
      auto it = this->remote_sensors_.find(provider.first);
      if (it != this->remote_sensors_.end()) {
        for (auto &val : it->second | std::views::values) {
          val->publish_state(NAN);
        }
      }
#endif
#ifdef USE_BINARY_SENSOR
      auto bs_it = this->remote_binary_sensors_.find(provider.first);
      if (bs_it != this->remote_binary_sensors_.end()) {
        for (auto &val : bs_it->second | std::views::values) {
          val->invalidate_state();
        }
      }
#endif
    } else {
#ifdef USE_STATUS_SENSOR
      if (provider.second.status_sensor != nullptr && !provider.second.status_sensor->state) {
        ESP_LOGI(TAG, "Ping status for %s restored at %" PRIu32 " with age %" PRIu32, provider.first.c_str(), now,
                 key_response_age);
        provider.second.status_sensor->publish_state(true);
      }
#endif
    }
  }
}

void PacketTransport::add_key_(const char *name, uint32_t key) {
  if (!this->is_encrypted_())
    return;
  auto it = this->ping_keys_.find(name);
  if (it == this->ping_keys_.end()) {
    if (this->ping_keys_.size() == MAX_PING_KEYS) {
      ESP_LOGW(TAG, "Ping key from %s discarded", name);
      return;
    }
    this->ping_keys_.emplace(name, key);  // allocates string key once only
  } else {
    it->second = key;  // key string already exists in map, no allocation
  }
  this->updated_ = true;
  ESP_LOGV(TAG, "Ping key from %s now %X", name, (unsigned) key);
}

static bool process_rolling_code(Provider &provider, PacketDecoder &decoder) {
  uint32_t code0, code1;
  if (decoder.get(code0) != DECODE_OK || decoder.get(code1) != DECODE_OK) {
    ESP_LOGW(TAG, "Rolling code requires 8 bytes");
    return false;
  }
  if (code1 < provider.last_code[1] || (code1 == provider.last_code[1] && code0 <= provider.last_code[0])) {
    ESP_LOGW(TAG, "Rolling code for %s %08lX:%08lX is old", provider.name, (unsigned long) code1,
             (unsigned long) code0);
    return false;
  }
  provider.last_code[0] = code0;
  provider.last_code[1] = code1;
  ESP_LOGV(TAG, "Saw new rolling code for %s %08lX:%08lX", provider.name, (unsigned long) code1, (unsigned long) code0);
  return true;
}

/**
 * Process a received packet
 */
void PacketTransport::process_(std::span<const uint8_t> data) {
  auto ping_key_seen = !this->ping_pong_enable_;
  PacketDecoder decoder(data.data(), data.size());
  char namebuf[256]{};
  uint8_t byte;
  FuData rdata{};
  uint16_t magic;
  if (decoder.get(magic) != DECODE_OK) {
    ESP_LOGD(TAG, "Short buffer");
    return;
  }
  if (magic != MAGIC_NUMBER && magic != MAGIC_PING) {
    ESP_LOGV(TAG, "Bad magic %X", magic);
    return;
  }

  if (decoder.decode_string(namebuf, sizeof namebuf) != DECODE_OK) {
    ESP_LOGV(TAG, "Bad hostname length");
    return;
  }
  if (strcmp(this->name_, namebuf) == 0) {
    ESP_LOGVV(TAG, "Ignoring our own data");
    return;
  }
  if (magic == MAGIC_PING) {
    uint32_t key;
    if (decoder.get(key) != DECODE_OK) {
      ESP_LOGW(TAG, "Bad ping request");
      return;
    }
    this->add_key_(namebuf, key);
    ESP_LOGV(TAG, "Updated ping key for %s to %08X", namebuf, (unsigned) key);
    return;
  }

  auto it = this->providers_.find(namebuf);
  if (it == this->providers_.end()) {
    ESP_LOGVV(TAG, "Unknown hostname %s", namebuf);
    return;
  }
  auto &provider = it->second;
  ESP_LOGV(TAG, "Found hostname %s", namebuf);

#ifdef USE_SENSOR
  auto &sensors = this->remote_sensors_.try_emplace(namebuf).first->second;
#endif
#ifdef USE_BINARY_SENSOR
  auto &binary_sensors = this->remote_binary_sensors_.try_emplace(namebuf).first->second;
#endif

  if (!decoder.bump_to(4)) {
    ESP_LOGW(TAG, "Bad packet length %zu", data.size());
  }
  auto len = decoder.get_remaining_size();
  if (round4(len) != len) {
    ESP_LOGW(TAG, "Bad payload length %zu", len);
    return;
  }

  // if encryption not used with this host, ping check is pointless since it would be easily spoofed.
  if (provider.encryption_key.empty())
    ping_key_seen = true;

  if (!provider.encryption_key.empty()) {
    decoder.decrypt((const uint32_t *) provider.encryption_key.data());
  }
  if (decoder.get(byte) != DECODE_OK) {
    ESP_LOGV(TAG, "No key byte");
    return;
  }

  if (byte == ROLLING_CODE_KEY) {
    if (!process_rolling_code(provider, decoder))
      return;
  } else if (byte != DATA_KEY) {
    ESP_LOGV(TAG, "Expected rolling_key or data_key, got %X", byte);
    return;
  }
  uint32_t key;
  while (decoder.get_remaining_size() != 0) {
    if (decoder.decode(ZERO_FILL_KEY) == DECODE_OK)
      continue;
    if (decoder.decode(PING_KEY, key) == DECODE_OK) {
      if (key == this->ping_key_) {
        ping_key_seen = true;
        provider.last_key_response_time = millis();
        ESP_LOGV(TAG, "Found good ping key %X at timestamp %" PRIu32, (unsigned) key, provider.last_key_response_time);
      } else {
        ESP_LOGV(TAG, "Unknown ping key %X", (unsigned) key);
      }
      continue;
    }
    if (!ping_key_seen) {
      ESP_LOGW(TAG, "Ping key not seen");
      this->resend_ping_key_ = true;
      break;
    }
    if (decoder.decode(BINARY_SENSOR_KEY, namebuf, sizeof(namebuf), byte) == DECODE_OK) {
      ESP_LOGV(TAG, "Got binary sensor %s %d", namebuf, byte);
#ifdef USE_BINARY_SENSOR
      auto bs = binary_sensors.find(namebuf);
      if (bs != binary_sensors.end()) {
        bs->second->publish_state(byte != 0);
      }
#endif
      continue;
    }
    if (decoder.decode(SENSOR_KEY, namebuf, sizeof(namebuf), rdata.u32) == DECODE_OK) {
      ESP_LOGV(TAG, "Got sensor %s %f", namebuf, rdata.f32);
#ifdef USE_SENSOR
      auto sensor_it = sensors.find(namebuf);
      if (sensor_it != sensors.end())
        sensor_it->second->publish_state(rdata.f32);
#endif
      continue;
    }
    if (decoder.get(byte) == DECODE_OK) {
      ESP_LOGW(TAG, "Unknown key %X", byte);
      char hex_buf[format_hex_pretty_size(PACKET_MAX_LOG_BYTES)];
      ESP_LOGD(TAG, "Buffer pos: %zu contents: %s", data.size() - decoder.get_remaining_size(),
               format_hex_pretty_to(hex_buf, data.data(), data.size()));
    }
    break;
  }
}

void PacketTransport::dump_config() {
  ESP_LOGCONFIG(TAG,
                "Packet Transport:\n"
                "  Platform: %s\n"
                "  Encrypted: %s\n"
                "  Ping-pong: %s",
                this->platform_name_, YESNO(this->is_encrypted_()), YESNO(this->ping_pong_enable_));
#ifdef USE_SENSOR
  for (const auto &sensor : this->sensors_)
    ESP_LOGCONFIG(TAG, "  Sensor: %s", sensor.id);
#endif
#ifdef USE_BINARY_SENSOR
  for (const auto &sensor : this->binary_sensors_)
    ESP_LOGCONFIG(TAG, "  Binary Sensor: %s", sensor.id);
#endif
  for (const auto &host : this->providers_) {
    ESP_LOGCONFIG(TAG, "  Remote host: %s", host.first.c_str());
    ESP_LOGCONFIG(TAG, "    Encrypted: %s", YESNO(!host.second.encryption_key.empty()));
#ifdef USE_SENSOR
    auto rs = this->remote_sensors_.find(host.first.c_str());
    if (rs != this->remote_sensors_.end()) {
      for (const auto &key : rs->second | std::views::keys)
        ESP_LOGCONFIG(TAG, "    Sensor: %s", key.c_str());
    }
#endif
#ifdef USE_BINARY_SENSOR
    auto rbs = this->remote_binary_sensors_.find(host.first.c_str());
    if (rbs != this->remote_binary_sensors_.end()) {
      for (const auto &key : rbs->second | std::views::keys)
        ESP_LOGCONFIG(TAG, "    Binary Sensor: %s", key.c_str());
    }
#endif
  }
}
void PacketTransport::increment_code_() {
  if (this->rolling_code_enable_) {
    if (++this->rolling_code_[0] == 0) {
      this->rolling_code_[1]++;
      this->pref_.save(&this->rolling_code_[1]);
      // must make sure it's saved immediately
      global_preferences->sync();
    }
  }
}

void PacketTransport::loop() {
  if (this->resend_ping_key_)
    this->send_ping_pong_request_();
  if (this->updated_) {
    this->send_data_(false);
  }
}

void PacketTransport::send_ping_pong_request_() {
  if (!this->ping_pong_enable_ || !this->should_send())
    return;
  this->ping_key_ = random_uint32();
  this->ping_header_.clear();
  add(this->ping_header_, MAGIC_PING);
  add(this->ping_header_, this->name_);
  add(this->ping_header_, this->ping_key_);
  this->send_packet(this->ping_header_);
  this->resend_ping_key_ = false;
  ESP_LOGV(TAG, "Sent new ping request %08X", (unsigned) this->ping_key_);
}
}  // namespace packet_transport
}  // namespace esphome