#include "bbnow_driver.hpp"

#include "esp_log.h"
#include "esp_now.h"
#include "esp_wifi.h"
#include "freertos/FreeRTOS.h"
#include "freertos/queue.h"

#include <algorithm>
#include <cstring>
#include <new>
#include <utility>
#include <variant>

namespace {

[[maybe_unused]] const static char TAG[] = "bbnow_driver";
constexpr size_t RX_QUEUE_LENGTH = 16;

struct ReceivedPacket {
  BBNowDriver::MacAddress source{};
  std::variant<BBNowPubPacket, BBNowPrivPacket> packet;
};

QueueHandle_t pub_queue = nullptr;
QueueHandle_t priv_queue = nullptr;

constexpr bool is_broadcast(const BBNowDriver::MacAddress &mac) {
  return mac == BBNowDriver::MacAddress{0xff, 0xff, 0xff, 0xff, 0xff, 0xff};
}

constexpr bool is_zero(const BBNowDriver::MacAddress &mac) {
  return mac == BBNowDriver::MacAddress{};
}

constexpr BBNowDriver::MacAddress BROADCAST_MAC = {0xff, 0xff, 0xff,
                                                   0xff, 0xff, 0xff};

bool is_broadcast(const uint8_t *mac) {
  return mac != nullptr &&
         std::memcmp(mac, BROADCAST_MAC.data(), BROADCAST_MAC.size()) == 0;
}

bool is_valid_expander(BlkBoxExpantion type) {
  return type != BlkBoxExpantion::NONE && type <= BlkBoxExpantion::MAX;
}

void destroy_queues() {
  if (pub_queue != nullptr) {
    ReceivedPacket *packet = nullptr;
    while (xQueueReceive(pub_queue, &packet, 0) == pdTRUE) {
      delete packet;
    }
    vQueueDelete(pub_queue);
    pub_queue = nullptr;
  }
  if (priv_queue != nullptr) {
    ReceivedPacket *packet = nullptr;
    while (xQueueReceive(priv_queue, &packet, 0) == pdTRUE) {
      delete packet;
    }
    vQueueDelete(priv_queue);
    priv_queue = nullptr;
  }
}

void espnow_receive_callback(
  const esp_now_recv_info_t *receive_info,
  const uint8_t *payload, int length
) {
  if (receive_info == nullptr || receive_info->src_addr == nullptr ||
      receive_info->des_addr == nullptr || payload == nullptr || length <= 0 ||
      static_cast<size_t>(length) > ESP_NOW_MAX_DATA_LEN) {
    ESP_LOGW(TAG, "dropping ESP-NOW packet with invalid callback arguments");
    return;
  }

  const size_t payload_length = static_cast<size_t>(length);
  const bool is_pub = is_broadcast(receive_info->des_addr);
  QueueHandle_t queue = is_pub ? pub_queue : priv_queue;
  if (queue == nullptr) {
    ESP_LOGW(TAG, "dropping ESP-NOW packet before driver initialization");
    return;
  }

  ReceivedPacket *packet = nullptr;
  if (is_pub) {
    auto decoded = BBNowPubPacket::deserialize(
        std::span<const uint8_t>(payload, payload_length));
    if (!decoded.has_value()) {
      ESP_LOGW(TAG, "dropping invalid public ESP-NOW packet");
      return;
    }

    packet = new (std::nothrow) ReceivedPacket{
        .source = {},
        .packet = std::move(*decoded),
    };
  } else {
    auto decoded = BBNowPrivPacket::deserialize(
        std::span<const uint8_t>(payload, payload_length));
    if (!decoded.has_value()) {
      ESP_LOGW(TAG, "dropping invalid private ESP-NOW packet");
      return;
    }

    packet = new (std::nothrow) ReceivedPacket{
        .source = {},
        .packet = std::move(*decoded),
    };
  }

  if (packet == nullptr) {
    ESP_LOGW(TAG, "dropping ESP-NOW packet: no memory for receive queue");
    return;
  }
  std::memcpy(packet->source.data(), receive_info->src_addr,
              packet->source.size());

  // The callback cannot wait for application code to drain the queue.  A
  // warning makes queue overflow visible while preserving the non-blocking
  // callback contract.
  ReceivedPacket *queued_packet = packet;
  if (xQueueSend(queue, &queued_packet, 0) != pdTRUE) {
    ESP_LOGW(TAG, "dropping ESP-NOW packet: receive queue is full");
    delete packet;
    return;
  }
}

bool receive_packet(
  QueueHandle_t queue, ReceivedPacket *&packet,
  TickType_t max_delay
) {
  return queue != nullptr && xQueueReceive(queue, &packet, max_delay) == pdTRUE;
}

} // namespace

bool BBNowDriver::initialized_ = false;
std::vector<BBNowDriver::PairedExpander> BBNowDriver::paired_expanders;

esp_err_t BBNowDriver::init() {
  if (initialized_) {
    return ESP_OK;
  }

  pub_queue = xQueueCreate(RX_QUEUE_LENGTH, sizeof(ReceivedPacket *));
  priv_queue = xQueueCreate(RX_QUEUE_LENGTH, sizeof(ReceivedPacket *));
  if (pub_queue == nullptr || priv_queue == nullptr) {
    destroy_queues();
    return ESP_ERR_NO_MEM;
  }

  esp_err_t err = esp_now_init();
  if (err != ESP_OK) {
    destroy_queues();
    return err;
  }

  esp_now_peer_info_t broadcast_peer{};
  std::memcpy(broadcast_peer.peer_addr, BROADCAST_MAC.data(),
              BROADCAST_MAC.size());
  broadcast_peer.channel = 0;
  broadcast_peer.ifidx = WIFI_IF_STA;
  broadcast_peer.encrypt = false;

  err = esp_now_add_peer(&broadcast_peer);
  if (err != ESP_OK) {
    (void)esp_now_deinit();
    destroy_queues();
    return err;
  }

  err = esp_now_register_recv_cb(espnow_receive_callback);
  if (err != ESP_OK) {
    (void)esp_now_deinit();
    destroy_queues();
    return err;
  }

  initialized_ = true;
  return ESP_OK;
}

esp_err_t BBNowDriver::deinit() {
  if (!initialized_) {
    return ESP_OK;
  }

  esp_err_t err = esp_now_deinit();
  if (err != ESP_OK) {
    return err;
  }

  initialized_ = false;
  paired_expanders.clear();
  destroy_queues();
  return ESP_OK;
}

esp_err_t BBNowDriver::send_pub(const BBNowPubPacket &packet) {
  if (!initialized_) {
    return ESP_ERR_INVALID_STATE;
  }

  const std::vector<uint8_t> payload = packet.serialize();
  if (payload.empty()) {
    return ESP_ERR_INVALID_ARG;
  }
  if (payload.size() > ESP_NOW_MAX_DATA_LEN) {
    return ESP_ERR_INVALID_SIZE;
  }

  return esp_now_send(BROADCAST_MAC.data(), payload.data(), payload.size());
}

bool BBNowDriver::has_pub() {
  return pub_queue != nullptr && uxQueueMessagesWaiting(pub_queue) != 0;
}

bool BBNowDriver::recv_pub(
  MacAddress &source,
  BBNowPubPacket &packet,
  TickType_t max_delay
) {
  ReceivedPacket *received = nullptr;
  if (receive_packet(pub_queue, received, max_delay)) {
    source = received->source;
    packet = std::move(std::get<BBNowPubPacket>(received->packet));
    delete received;
    return true;
  }
  return false;
}

esp_err_t BBNowDriver::send_priv(
  const MacAddress &destination,                               
  const BBNowPrivPacket &packet
) {
  if (!initialized_) {
    return ESP_ERR_INVALID_STATE;
  }
  if (is_zero(destination)) {
    return ESP_ERR_INVALID_ARG;
  }

  const std::vector<uint8_t> payload = packet.serialize();
  if (payload.empty()) {
    return ESP_ERR_INVALID_ARG;
  }
  if (payload.size() > ESP_NOW_MAX_DATA_LEN) {
    return ESP_ERR_INVALID_SIZE;
  }

  return esp_now_send(destination.data(), payload.data(), payload.size());
}

bool BBNowDriver::has_priv() {
  return priv_queue != nullptr && uxQueueMessagesWaiting(priv_queue) != 0;
}

bool BBNowDriver::recv_priv(
  MacAddress &source,
  BBNowPrivPacket &packet,
  TickType_t max_delay
) {
  ReceivedPacket *received = nullptr;
  if (receive_packet(priv_queue, received, max_delay)) {
    source = received->source;
    packet = std::move(std::get<BBNowPrivPacket>(received->packet));
    delete received;
    return true;
  }
  return false;
}

size_t BBNowDriver::send_discovery(BlkBoxExpantion type) {
  if (!is_valid_expander(type)) {
    return 0;
  }

  BlkBoxNowPubDiscoveryData data{.expantion_type = type};
  BBNowPubPacket packet = BBNowPubPacket::new_discovery(data);
  const size_t size = packet.serialize_size();
  return send_pub(packet) == ESP_OK ? size : 0;
}

esp_err_t BBNowDriver::pair_expander(BlkBoxExpantion type,
                                     const MacAddress &mac) {
  if (!initialized_) {
    return ESP_ERR_INVALID_STATE;
  }
  if (!is_valid_expander(type) || is_zero(mac) || is_broadcast(mac)) {
    return ESP_ERR_INVALID_ARG;
  }

  for (auto &expander : paired_expanders) {
    if (expander.mac == mac) {
      expander.type = type;
      return ESP_OK;
    }
  }

  esp_now_peer_info_t peer{};
  std::memcpy(peer.peer_addr, mac.data(), mac.size());
  peer.channel = 0;
  peer.ifidx = WIFI_IF_STA;
  peer.encrypt = false;

  esp_err_t err = esp_now_add_peer(&peer);
  if (err != ESP_OK) {
    return err;
  }

  paired_expanders.push_back(PairedExpander{
      .mac = mac,
      .name = {},
      .type = type,
  });

  return ESP_OK;
}

esp_err_t BBNowDriver::unpair_expander(const MacAddress &mac) {
  if (!initialized_) {
    return ESP_ERR_INVALID_STATE;
  }
  if (is_zero(mac) || is_broadcast(mac)) {
    return ESP_ERR_INVALID_ARG;
  }

  auto it = std::find_if(
      paired_expanders.begin(), paired_expanders.end(),
      [&mac](const PairedExpander &expander) { return expander.mac == mac; });
  if (it == paired_expanders.end()) {
    return ESP_ERR_NOT_FOUND;
  }

  esp_err_t err = esp_now_del_peer(mac.data());
  if (err != ESP_OK) {
    return err;
  }

  paired_expanders.erase(it);
  return ESP_OK;
}

std::span<const BBNowDriver::PairedExpander>
BBNowDriver::get_paired_expanders() {
  return paired_expanders;
}