#include "bosch456.h"
#include "client_registry.h"
#include "esp_now_comm.h"
#include "esp_now_proto.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_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]";

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 fill_presence(alox_EspNowSlavePresence *presence) {
  presence->network = s_config.network;
  presence->version = POWERPOD_FW_VERSION;
  presence->slave_id = s_own_mac[5];
  presence->available = true;
  presence->used = false;
  esp_now_proto_setup_presence_encode(presence, s_own_mac);
}

static esp_err_t send_message(const uint8_t *dest_mac,
                              const alox_EspNowMessage *msg) {
  uint8_t buf[ESPNOW_PB_MAX_SIZE];
  size_t len = 0;

  esp_err_t err = esp_now_proto_encode(msg, buf, sizeof(buf), &len);
  if (err != ESP_OK) {
    ESP_LOGW(TAG, "encode failed");
    return err;
  }

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

  err = esp_now_send(dest_mac, buf, len);
  if (err != ESP_OK) {
    ESP_LOGW(TAG, "send type=%u failed: %s", (unsigned)msg->type,
             esp_err_to_name(err));
  }
  return err;
}

static esp_err_t send_accel_deadzone(const uint8_t *dest_mac, uint32_t client_id,
                                     uint32_t deadzone) {
  alox_EspNowMessage msg = alox_EspNowMessage_init_zero;

  msg.type = alox_EspNowMessageType_ESPNOW_SET_ACCEL_DEADZONE;
  msg.which_payload = alox_EspNowMessage_accel_deadzone_tag;
  msg.payload.accel_deadzone.deadzone = deadzone;
  msg.payload.accel_deadzone.client_id = client_id;

  return send_message(dest_mac, &msg);
}

static esp_err_t send_unicast_test(const uint8_t *dest_mac, uint32_t seq) {
  alox_EspNowMessage msg = alox_EspNowMessage_init_zero;

  msg.type = alox_EspNowMessageType_ESPNOW_UNICAST_TEST;
  msg.which_payload = alox_EspNowMessage_unicast_test_tag;
  msg.payload.unicast_test.seq = seq;

  return send_message(dest_mac, &msg);
}

esp_err_t esp_now_comm_send_unicast_test(const uint8_t mac[CLIENT_MAC_LEN],
                                         uint32_t seq) {
  if (mac == NULL || !s_config.master) {
    return ESP_ERR_INVALID_STATE;
  }

  char mac_str[18];
  mac_to_str(mac, mac_str, sizeof(mac_str));
  esp_err_t err = send_unicast_test(mac, seq);
  if (err == ESP_OK) {
    ESP_LOGI(TAG, "unicast TEST to %s seq=%lu", mac_str, (unsigned long)seq);
  } else {
    ESP_LOGW(TAG, "unicast TEST to %s failed: %s", mac_str, esp_err_to_name(err));
  }
  return err;
}

esp_err_t esp_now_comm_send_accel_deadzone(const uint8_t mac[CLIENT_MAC_LEN],
                                           uint32_t client_id, uint32_t deadzone) {
  if (mac == NULL || !s_config.master) {
    return ESP_ERR_INVALID_STATE;
  }

  char mac_str[18];
  mac_to_str(mac, mac_str, sizeof(mac_str));
  esp_err_t err = send_accel_deadzone(mac, client_id, deadzone);
  if (err == ESP_OK) {
    ESP_LOGI(TAG, "unicast SET_ACCEL_DEADZONE to %s: deadzone=%lu client_id=%lu",
             mac_str, (unsigned long)deadzone, (unsigned long)client_id);
  } else {
    ESP_LOGW(TAG, "unicast SET_ACCEL_DEADZONE to %s failed: %s", mac_str,
             esp_err_to_name(err));
  }
  return err;
}

static void send_presence(const uint8_t *dest_mac,
                          alox_EspNowMessageType type) {
  alox_EspNowMessage msg = alox_EspNowMessage_init_zero;
  alox_EspNowSlavePresence *presence = NULL;

  msg.type = type;
  if (type == alox_EspNowMessageType_ESPNOW_SLAVE_INFO) {
    msg.which_payload = alox_EspNowMessage_slave_info_tag;
    presence = &msg.payload.slave_info;
  } else {
    msg.which_payload = alox_EspNowMessage_heartbeat_tag;
    presence = &msg.payload.heartbeat;
  }

  fill_presence(presence);
  send_message(dest_mac, &msg);
}

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_slave_unicast_test(const uint8_t *master_mac,
                                      const alox_EspNowUnicastTest *test) {
  char mac_str[18];
  mac_to_str(master_mac, mac_str, sizeof(mac_str));

  ESP_LOGI(TAG, "UNICAST TEST OK from master %s seq=%lu (joined=%d)",
           mac_str, (unsigned long)test->seq, (int)s_slave_joined);
}

static void handle_slave_accel_deadzone(const uint8_t *master_mac,
                                        const alox_EspNowAccelDeadzone *cfg) {
  uint32_t my_id = s_own_mac[5];

  if (cfg->client_id != 0 && cfg->client_id != my_id) {
    return;
  }

  if (s_slave_joined && !mac_equal(master_mac, s_master_mac)) {
    return;
  }

  char mac_str[18];
  mac_to_str(master_mac, mac_str, sizeof(mac_str));

  ESP_LOGI(TAG,
           "accel deadzone from master %s: %lu LSB id=%lu (sensor %s)",
           mac_str, (unsigned long)cfg->deadzone, (unsigned long)my_id,
           bma456_is_ready() ? "ok" : "not installed");

  bma456_set_accel_deadzone(cfg->deadzone);
}

static void handle_client_presence(const alox_EspNowSlavePresence *presence,
                                   const uint8_t mac[CLIENT_MAC_LEN]) {
  if (presence->network != s_config.network) {
    return;
  }

  ensure_peer(mac);

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

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

static void handle_discover(const uint8_t *sender_mac,
                            const alox_EspNowDiscover *discover) {
  if (discover->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;
  ensure_peer(sender_mac);

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

  send_presence(sender_mac, alox_EspNowMessageType_ESPNOW_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_presence(s_master_mac, alox_EspNowMessageType_ESPNOW_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 <= 0) {
    return;
  }

  if (!s_config.master) {
    alox_EspNowMessage msg = alox_EspNowMessage_init_zero;

    if (esp_now_proto_decode(data, (size_t)len, &msg) != ESP_OK) {
      ESP_LOGW(TAG, "slave: ESP-NOW decode failed (%d bytes)", len);
      return;
    }

    if (s_slave_joined && mac_equal(info->src_addr, s_master_mac)) {
      ensure_peer(info->src_addr);
    }

    switch (msg.which_payload) {
    case alox_EspNowMessage_discover_tag:
      handle_discover(info->src_addr, &msg.payload.discover);
      break;
    case alox_EspNowMessage_unicast_test_tag:
      handle_slave_unicast_test(info->src_addr, &msg.payload.unicast_test);
      break;
    case alox_EspNowMessage_accel_deadzone_tag:
      handle_slave_accel_deadzone(info->src_addr, &msg.payload.accel_deadzone);
      break;
    default:
      ESP_LOGW(TAG, "slave: unhandled ESP-NOW which=%u type=%u", msg.which_payload,
               (unsigned)msg.type);
      break;
    }
    return;
  }

  alox_EspNowMessage msg = alox_EspNowMessage_init_zero;

  if (esp_now_proto_decode(data, (size_t)len, &msg) != ESP_OK) {
    ESP_LOGW(TAG, "master: ESP-NOW decode failed (%d bytes)", len);
    return;
  }

  const alox_EspNowSlavePresence *presence = esp_now_proto_get_presence(&msg);
  if (presence != NULL) {
    /* Registry key is the ESP-NOW sender MAC, not the optional protobuf mac field. */
    ensure_peer(info->src_addr);
    handle_client_presence(presence, info->src_addr);
  }
}

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

  alox_EspNowMessage msg = alox_EspNowMessage_init_zero;
  msg.type = alox_EspNowMessageType_ESPNOW_DISCOVER;
  msg.which_payload = alox_EspNowMessage_discover_tag;
  msg.payload.discover.network = s_config.network;

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

  while (1) {
    send_message(ESPNOW_BCAST, &msg);
    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_ps(WIFI_PS_NONE));
  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;
}