#include "client_registry.h"
#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_MSG_HEARTBEAT 3
#define ESPNOW_DISCOVER_INTERVAL_MS 500
#define ESPNOW_HEARTBEAT_INTERVAL_MS 1000
#define ESPNOW_HEARTBEAT_MISS_COUNT 3
#define ESPNOW_CLIENT_TIMEOUT_MS                                               \
  (ESPNOW_HEARTBEAT_INTERVAL_MS * ESPNOW_HEARTBEAT_MISS_COUNT)
#define SLAVE_MASTER_LOST_MS (ESPNOW_HEARTBEAT_INTERVAL_MS * 5)

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_packet_t;

static app_config_t s_config;
static uint8_t s_wifi_channel;
static uint8_t s_own_mac[ESP_NOW_ETH_ALEN];
static bool s_slave_joined;
static uint8_t s_master_mac[ESP_NOW_ETH_ALEN];
static uint32_t s_last_discover_ms;

static uint32_t now_ms(void) {
  return (uint32_t)(xTaskGetTickCount() * portTICK_PERIOD_MS);
}

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 void build_slave_packet(espnow_slave_packet_t *pkt, uint8_t type) {
  pkt->magic = ESPNOW_MAGIC;
  pkt->type = type;
  pkt->network = s_config.network;
  memcpy(pkt->mac, s_own_mac, ESP_NOW_ETH_ALEN);
  pkt->version = POWERPOD_FW_VERSION;
  pkt->slave_id = s_own_mac[5];
  pkt->available = 1;
  pkt->used = 0;
}

static void send_slave_packet(const uint8_t *dest_mac, uint8_t type) {
  espnow_slave_packet_t pkt;
  build_slave_packet(&pkt, type);

  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, "send type=%u failed: %s", (unsigned)type,
             esp_err_to_name(err));
  }
}

static void slave_reset_join(void) {
  s_slave_joined = false;
  memset(s_master_mac, 0, sizeof(s_master_mac));
  s_last_discover_ms = 0;
}

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

  bool is_new = false;
  bool reactivated = false;
  esp_err_t err = client_registry_heartbeat(
      pkt->mac, pkt->slave_id, pkt->version, pkt->used != 0, &is_new,
      &reactivated);
  if (err != ESP_OK) {
    ESP_LOGW(TAG, "client registry full");
    return;
  }

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

static void handle_discover(const uint8_t *sender_mac,
                            const espnow_discover_packet_t *pkt) {
  if (pkt->network != s_config.network) {
    return;
  }

  uint32_t now = now_ms();

  if (s_slave_joined) {
    if (!mac_equal(sender_mac, s_master_mac)) {
      return;
    }
    if ((now - s_last_discover_ms) <= SLAVE_MASTER_LOST_MS) {
      s_last_discover_ms = now;
      return;
    }
    ESP_LOGW(TAG, "master lost, rejoining");
    slave_reset_join();
  }

  memcpy(s_master_mac, sender_mac, ESP_NOW_ETH_ALEN);
  s_slave_joined = true;
  s_last_discover_ms = now;

  char mac_str[18];
  mac_to_str(sender_mac, mac_str, sizeof(mac_str));
  ESP_LOGI(TAG, "joined network %u, master %s", (unsigned)pkt->network, mac_str);

  send_slave_packet(sender_mac, ESPNOW_MSG_SLAVE_INFO);
}

static void slave_check_master_timeout(void) {
  if (!s_slave_joined) {
    return;
  }

  uint32_t now = now_ms();
  if (s_last_discover_ms == 0) {
    return;
  }

  if ((now - s_last_discover_ms) > SLAVE_MASTER_LOST_MS) {
    ESP_LOGW(TAG, "no master discover for %u ms, reconnecting",
             (unsigned)(now - s_last_discover_ms));
    slave_reset_join();
  }
}

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

  ESP_LOGI(TAG, "slave heartbeat task (interval %u ms)",
           (unsigned)ESPNOW_HEARTBEAT_INTERVAL_MS);

  while (1) {
    vTaskDelay(pdMS_TO_TICKS(ESPNOW_HEARTBEAT_INTERVAL_MS));

    slave_check_master_timeout();

    if (!s_slave_joined) {
      continue;
    }

    send_slave_packet(s_master_mac, ESPNOW_MSG_HEARTBEAT);
  }
}

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

  ESP_LOGI(TAG, "master monitor task (timeout %u ms)",
           (unsigned)ESPNOW_CLIENT_TIMEOUT_MS);

  while (1) {
    vTaskDelay(pdMS_TO_TICKS(ESPNOW_HEARTBEAT_INTERVAL_MS));
    client_registry_check_timeouts(ESPNOW_CLIENT_TIMEOUT_MS);
  }
}

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:
  case ESPNOW_MSG_HEARTBEAT:
    if (s_config.master && len >= (int)sizeof(espnow_slave_packet_t)) {
      handle_client_packet((const espnow_slave_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));
  client_registry_init();
  slave_reset_join();

  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;
    }
    if (xTaskCreate(master_monitor_task, "espnow_mon", 4096, NULL, 4, NULL) !=
        pdPASS) {
      ESP_LOGE(TAG, "failed to create monitor task");
      return ESP_FAIL;
    }
  } else {
    if (xTaskCreate(slave_heartbeat_task, "espnow_hb", 4096, NULL, 4, NULL) !=
        pdPASS) {
      ESP_LOGE(TAG, "failed to create heartbeat task");
      return ESP_FAIL;
    }
  }

  return ESP_OK;
}