#include "esp_now_comm.h"
#include "esp_err.h"
#include "esp_event.h"
#include "esp_log.h"
#include "esp_mac.h"
#include "esp_netif.h"
#include "esp_now.h"
#include "esp_wifi.h"
#include "freertos/FreeRTOS.h"
#include "freertos/idf_additions.h"
#include "nvs_flash.h"
#include <stdio.h>
#include <string.h>

#ifndef POWERPOD_FW_VERSION
#define POWERPOD_FW_VERSION 1u
#endif

#define ESPNOW_MAGIC 0xA1
#define ESPNOW_MSG_DISCOVER 1
#define ESPNOW_MSG_SLAVE_INFO 2
#define ESPNOW_DISCOVER_INTERVAL_MS 500
#define ESPNOW_MAX_SLAVES 16

static const uint8_t ESPNOW_BCAST[ESP_NOW_ETH_ALEN] = {0xff, 0xff, 0xff,
                                                       0xff, 0xff, 0xff};

static const char *TAG = "[ESPNOW]";

typedef struct __attribute__((packed)) {
  uint8_t magic;
  uint8_t type;
  uint8_t network;
  uint8_t reserved;
} espnow_discover_packet_t;

typedef struct __attribute__((packed)) {
  uint8_t magic;
  uint8_t type;
  uint8_t network;
  uint8_t mac[ESP_NOW_ETH_ALEN];
  uint32_t version;
  uint32_t slave_id;
  uint8_t available;
  uint8_t used;
} espnow_slave_info_packet_t;

typedef struct {
  uint8_t mac[ESP_NOW_ETH_ALEN];
  uint32_t slave_id;
  uint32_t version;
  bool available;
  bool used;
  bool seen;
} slave_entry_t;

static app_config_t s_config;
static uint8_t s_wifi_channel;
static uint8_t s_own_mac[ESP_NOW_ETH_ALEN];
static slave_entry_t s_slaves[ESPNOW_MAX_SLAVES];

static uint8_t network_to_channel(uint8_t network) {
  if (network < 1 || network > 13) {
    return 1;
  }
  return network;
}

static bool mac_equal(const uint8_t *a, const uint8_t *b) {
  return memcmp(a, b, ESP_NOW_ETH_ALEN) == 0;
}

static void mac_to_str(const uint8_t *mac, char *out, size_t out_len) {
  snprintf(out, out_len, "%02x:%02x:%02x:%02x:%02x:%02x", mac[0], mac[1],
           mac[2], mac[3], mac[4], mac[5]);
}

static esp_err_t ensure_peer(const uint8_t *mac) {
  if (esp_now_is_peer_exist(mac)) {
    return ESP_OK;
  }

  esp_now_peer_info_t peer = {0};
  memcpy(peer.peer_addr, mac, ESP_NOW_ETH_ALEN);
  peer.channel = s_wifi_channel;
  peer.ifidx = WIFI_IF_STA;
  peer.encrypt = false;

  esp_err_t err = esp_now_add_peer(&peer);
  if (err != ESP_OK) {
    ESP_LOGW(TAG, "add peer failed: %s", esp_err_to_name(err));
  }
  return err;
}

static esp_err_t ensure_broadcast_peer(void) { return ensure_peer(ESPNOW_BCAST); }

static slave_entry_t *find_slave(const uint8_t *mac) {
  for (int i = 0; i < ESPNOW_MAX_SLAVES; i++) {
    if (s_slaves[i].seen && mac_equal(s_slaves[i].mac, mac)) {
      return &s_slaves[i];
    }
  }
  return NULL;
}

static slave_entry_t *alloc_slave(const uint8_t *mac) {
  slave_entry_t *existing = find_slave(mac);
  if (existing != NULL) {
    return existing;
  }
  for (int i = 0; i < ESPNOW_MAX_SLAVES; i++) {
    if (!s_slaves[i].seen) {
      memcpy(s_slaves[i].mac, mac, ESP_NOW_ETH_ALEN);
      s_slaves[i].seen = true;
      return &s_slaves[i];
    }
  }
  return NULL;
}

static void send_slave_info(const uint8_t *dest_mac) {
  espnow_slave_info_packet_t pkt = {
      .magic = ESPNOW_MAGIC,
      .type = ESPNOW_MSG_SLAVE_INFO,
      .network = s_config.network,
      .version = POWERPOD_FW_VERSION,
      .slave_id = s_own_mac[5],
      .available = 1,
      .used = 0,
  };
  memcpy(pkt.mac, s_own_mac, ESP_NOW_ETH_ALEN);

  if (ensure_peer(dest_mac) != ESP_OK) {
    return;
  }

  esp_err_t err =
      esp_now_send(dest_mac, (const uint8_t *)&pkt, sizeof(pkt));
  if (err != ESP_OK) {
    ESP_LOGW(TAG, "slave info send failed: %s", esp_err_to_name(err));
  }
}

static void handle_discover(const uint8_t *sender_mac,
                            const espnow_discover_packet_t *pkt) {
  if (pkt->network != s_config.network) {
    return;
  }
  ESP_LOGI(TAG, "discover from master on network %u", (unsigned)pkt->network);
  send_slave_info(sender_mac);
}

static void handle_slave_info(const espnow_slave_info_packet_t *pkt) {
  if (pkt->network != s_config.network) {
    return;
  }

  bool is_new = (find_slave(pkt->mac) == NULL);
  slave_entry_t *entry = alloc_slave(pkt->mac);
  if (entry == NULL) {
    ESP_LOGW(TAG, "slave table full");
    return;
  }

  entry->slave_id = pkt->slave_id;
  entry->version = pkt->version;
  entry->available = pkt->available != 0;
  entry->used = pkt->used != 0;

  char mac_str[18];
  mac_to_str(pkt->mac, mac_str, sizeof(mac_str));
  if (is_new) {
    ESP_LOGI(TAG, "slave joined id=%lu mac=%s ver=%lu",
             (unsigned long)pkt->slave_id, mac_str,
             (unsigned long)pkt->version);
  }
}

static void espnow_recv_cb(const esp_now_recv_info_t *info, const uint8_t *data,
                           int len) {
  if (info == NULL || data == NULL || len < 3) {
    return;
  }

  if (data[0] != ESPNOW_MAGIC) {
    return;
  }

  switch (data[1]) {
  case ESPNOW_MSG_DISCOVER:
    if (!s_config.master && len >= (int)sizeof(espnow_discover_packet_t)) {
      handle_discover(info->src_addr,
                      (const espnow_discover_packet_t *)data);
    }
    break;
  case ESPNOW_MSG_SLAVE_INFO:
    if (s_config.master && len >= (int)sizeof(espnow_slave_info_packet_t)) {
      handle_slave_info((const espnow_slave_info_packet_t *)data);
    }
    break;
  default:
    break;
  }
}

static void master_discover_task(void *param) {
  (void)param;

  espnow_discover_packet_t pkt = {
      .magic = ESPNOW_MAGIC,
      .type = ESPNOW_MSG_DISCOVER,
      .network = s_config.network,
      .reserved = 0,
  };

  ESP_LOGI(TAG, "master discover task on network %u ch %u",
           (unsigned)s_config.network, (unsigned)s_wifi_channel);

  while (1) {
    esp_err_t err =
        esp_now_send(ESPNOW_BCAST, (const uint8_t *)&pkt, sizeof(pkt));
    if (err != ESP_OK) {
      ESP_LOGW(TAG, "discover broadcast failed: %s", esp_err_to_name(err));
    }
    vTaskDelay(pdMS_TO_TICKS(ESPNOW_DISCOVER_INTERVAL_MS));
  }
}

static esp_err_t init_wifi_stack(uint8_t channel) {
  esp_err_t err;

  err = nvs_flash_init();
  if (err == ESP_ERR_NVS_NO_FREE_PAGES || err == ESP_ERR_NVS_NEW_VERSION_FOUND) {
    ESP_ERROR_CHECK(nvs_flash_erase());
    err = nvs_flash_init();
  }
  if (err != ESP_OK) {
    return err;
  }

  ESP_ERROR_CHECK(esp_netif_init());
  ESP_ERROR_CHECK(esp_event_loop_create_default());

  wifi_init_config_t cfg = WIFI_INIT_CONFIG_DEFAULT();
  ESP_ERROR_CHECK(esp_wifi_init(&cfg));

  wifi_config_t wifi_config = {0};
  wifi_config.sta.channel = channel;
  wifi_config.sta.scan_method = WIFI_ALL_CHANNEL_SCAN;
  wifi_config.sta.sort_method = WIFI_CONNECT_AP_BY_SIGNAL;

  ESP_ERROR_CHECK(esp_wifi_set_mode(WIFI_MODE_STA));
  ESP_ERROR_CHECK(esp_wifi_set_config(WIFI_IF_STA, &wifi_config));
  ESP_ERROR_CHECK(esp_wifi_start());
  ESP_ERROR_CHECK(esp_wifi_set_channel(channel, WIFI_SECOND_CHAN_NONE));

  return ESP_OK;
}

esp_err_t esp_now_comm_init(const app_config_t *config) {
  if (config == NULL) {
    return ESP_ERR_INVALID_ARG;
  }

  memset(&s_config, 0, sizeof(s_config));
  memcpy(&s_config, config, sizeof(s_config));
  memset(s_slaves, 0, sizeof(s_slaves));

  s_wifi_channel = network_to_channel(config->network);
  ESP_ERROR_CHECK(esp_read_mac(s_own_mac, ESP_MAC_WIFI_STA));

  char mac_str[18];
  mac_to_str(s_own_mac, mac_str, sizeof(mac_str));
  ESP_LOGI(TAG, "role=%s network=%u channel=%u mac=%s",
           config->master ? "master" : "slave", (unsigned)config->network,
           (unsigned)s_wifi_channel, mac_str);

  esp_err_t err = init_wifi_stack(s_wifi_channel);
  if (err != ESP_OK) {
    ESP_LOGE(TAG, "wifi init failed: %s", esp_err_to_name(err));
    return err;
  }

  ESP_ERROR_CHECK(esp_now_init());
  ESP_ERROR_CHECK(esp_now_register_recv_cb(espnow_recv_cb));

  if (config->master) {
    ESP_ERROR_CHECK(ensure_broadcast_peer());
    if (xTaskCreate(master_discover_task, "espnow_disc", 4096, NULL, 4,
                    NULL) != pdPASS) {
      ESP_LOGE(TAG, "failed to create discover task");
      return ESP_FAIL;
    }
  } else {
    ESP_LOGI(TAG, "slave listening for master discover");
  }

  return ESP_OK;
}