#include "ota_espnow.h"
#include "app_config.h"
#include "client_registry.h"
#include "esp_log.h"
#include "esp_now_comm.h"
#include "esp_now_messages.pb.h"
#include "esp_partition.h"
#include "freertos/FreeRTOS.h"
#include "freertos/event_groups.h"
#include "freertos/idf_additions.h"
#include "ota_uart.h"
#include <string.h>

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

#define OTA_ESPNOW_PREPARE_STACK 8192
#define OTA_ESPNOW_PREPARE_PRIO 5

#define OTA_PREPARE_TIMEOUT_MS 120000u
#define OTA_BLOCK_TIMEOUT_MS 30000u
#define OTA_END_TIMEOUT_MS 60000u
#define OTA_PAYLOAD_DELAY_MS 3

#define OTA_ST_PREPARING 1u
#define OTA_ST_READY 2u
#define OTA_ST_BLOCK_ACK 3u
#define OTA_ST_SUCCESS 4u
#define OTA_ST_FAILED 5u

#define OTA_MAX_TARGETS CLIENT_REGISTRY_MAX

static EventGroupHandle_t s_eg;

typedef struct {
  uint8_t count;
  uint8_t mac[OTA_MAX_TARGETS][6];
  uint32_t id[OTA_MAX_TARGETS];
  uint32_t expected_bytes;
} ota_dist_t;

static ota_dist_t s_dist;

static int find_target_index(const uint8_t mac[6]) {
  for (uint8_t i = 0; i < s_dist.count; i++) {
    if (memcmp(s_dist.mac[i], mac, 6) == 0) {
      return (int)i;
    }
  }
  return -1;
}

static uint32_t all_target_bits(void) {
  if (s_dist.count == 0 || s_dist.count > 31) {
    return 0;
  }
  return (1u << s_dist.count) - 1u;
}

static bool wait_target_bits(uint32_t want_bits, uint32_t timeout_ms) {
  if (s_eg == NULL || want_bits == 0) {
    return false;
  }
  EventBits_t got =
      xEventGroupWaitBits(s_eg, want_bits, pdTRUE, pdTRUE,
                          pdMS_TO_TICKS(timeout_ms));
  return (got & want_bits) == want_bits;
}

static void send_slave_status(const uint8_t master_mac[6], uint32_t status,
                              uint32_t bytes_written, uint32_t error) {
  esp_now_comm_send_ota_status(master_mac, status, bytes_written, error);
}

static void ota_slave_prepare_task(void *param) {
  uint32_t total_size = (uint32_t)(uintptr_t)param;
  uint8_t master_mac[6];

  if (!esp_now_comm_get_master_mac(master_mac)) {
    vTaskDelete(NULL);
    return;
  }

  send_slave_status(master_mac, OTA_ST_PREPARING, 0, 0);

  int slot = ota_uart_prepare(total_size);
  if (slot < 0) {
    send_slave_status(master_mac, OTA_ST_FAILED, 0, 1);
    vTaskDelete(NULL);
    return;
  }

  send_slave_status(master_mac, OTA_ST_READY, 0, 0);
  vTaskDelete(NULL);
}

void ota_espnow_slave_on_start(const uint8_t master_mac[6],
                               const alox_EspNowOtaStart *start) {
  if (start == NULL || start->total_size == 0) {
    return;
  }

  ESP_LOGI(TAG, "ESP-NOW OTA_START (%lu bytes)", (unsigned long)start->total_size);

  if (ota_uart_is_active()) {
    send_slave_status(master_mac, OTA_ST_FAILED, 0, 4);
    return;
  }

  if (xTaskCreate(ota_slave_prepare_task, "ota_esp_prep", OTA_ESPNOW_PREPARE_STACK,
                  (void *)(uintptr_t)start->total_size, OTA_ESPNOW_PREPARE_PRIO,
                  NULL) != pdPASS) {
    send_slave_status(master_mac, OTA_ST_FAILED, 0, 5);
  }
}

void ota_espnow_slave_on_payload(const uint8_t master_mac[6],
                                 const alox_EspNowOtaPayload *payload) {
  if (payload == NULL || payload->data.size == 0) {
    send_slave_status(master_mac, OTA_ST_FAILED, 0, 11);
    return;
  }

  if (!ota_uart_is_active()) {
    ESP_LOGW(TAG, "OTA_PAYLOAD seq=%lu but no active session",
             (unsigned long)payload->seq);
    send_slave_status(master_mac, OTA_ST_FAILED, 0, 12);
    return;
  }

  if (payload->seq == 0) {
    ESP_LOGI(TAG, "ESP-NOW OTA payloads started");
  }

  ota_feed_result_t r =
      ota_uart_feed(payload->data.bytes, payload->data.size);
  if (r == OTA_FEED_ERROR) {
    send_slave_status(master_mac, OTA_ST_FAILED, ota_uart_bytes_written(), 13);
    return;
  }
  if (r == OTA_FEED_BLOCK_WRITTEN) {
    uint32_t written = ota_uart_bytes_written();
    ESP_LOGI(TAG, "block written %lu bytes -> ack master", (unsigned long)written);
    send_slave_status(master_mac, OTA_ST_BLOCK_ACK, written, 0);
  }
}

void ota_espnow_slave_on_end(const uint8_t master_mac[6]) {
  ESP_LOGI(TAG, "ESP-NOW OTA_END");
  if (!ota_uart_is_active()) {
    send_slave_status(master_mac, OTA_ST_FAILED, 0, 20);
    return;
  }

  uint32_t written = ota_uart_bytes_written();
  bool success = false;
  esp_err_t err = ota_uart_finish(true, &success);
  if (err != ESP_OK || !success) {
    send_slave_status(master_mac, OTA_ST_FAILED, written, (uint32_t)err);
    return;
  }

  send_slave_status(master_mac, OTA_ST_SUCCESS, written, 0);
  ESP_LOGI(TAG, "slave OTA success (%lu bytes), reboot to run",
           (unsigned long)written);
}

void ota_espnow_master_on_status(const uint8_t slave_mac[6],
                                 const alox_EspNowOtaStatus *status) {
  if (status == NULL || s_eg == NULL) {
    return;
  }

  int idx = find_target_index(slave_mac);
  if (idx < 0) {
    return;
  }

  uint32_t bit = (1u << (unsigned)idx);

  switch (status->status) {
  case OTA_ST_READY:
    xEventGroupSetBits(s_eg, bit);
    break;
  case OTA_ST_BLOCK_ACK:
    if (status->bytes_written >= s_dist.expected_bytes) {
      xEventGroupSetBits(s_eg, bit);
    } else {
      ESP_LOGW(TAG, "slave %lu block ack early (%lu < %lu)",
               (unsigned long)s_dist.id[idx], (unsigned long)status->bytes_written,
               (unsigned long)s_dist.expected_bytes);
    }
    break;
  case OTA_ST_SUCCESS:
    xEventGroupSetBits(s_eg, bit);
    break;
  case OTA_ST_FAILED:
    ESP_LOGW(TAG, "slave %lu OTA failed (err=%lu)",
             (unsigned long)s_dist.id[idx], (unsigned long)status->error);
    break;
  default:
    break;
  }
}

static size_t collect_targets(void) {
  memset(&s_dist, 0, sizeof(s_dist));

  size_t n = client_registry_count();
  for (size_t i = 0; i < n && s_dist.count < OTA_MAX_TARGETS; i++) {
    const client_info_t *c = client_registry_at(i);
    if (c == NULL || !c->available) {
      continue;
    }
    uint8_t slot = s_dist.count;
    memcpy(s_dist.mac[slot], c->mac, 6);
    s_dist.id[slot] = c->id;
    s_dist.count++;
  }
  return s_dist.count;
}

static esp_err_t distribute_image(const esp_partition_t *partition,
                                  uint32_t size) {
  if (s_eg == NULL) {
    s_eg = xEventGroupCreate();
    if (s_eg == NULL) {
      return ESP_ERR_NO_MEM;
    }
  }

  ESP_LOGI(TAG, "distributing %lu bytes from %s to %u slave(s)",
           (unsigned long)size, partition->label, (unsigned)s_dist.count);

  uint32_t target_mask = all_target_bits();
  esp_err_t err;

  xEventGroupClearBits(s_eg, target_mask);
  for (uint8_t i = 0; i < s_dist.count; i++) {
    err = esp_now_comm_send_ota_start(s_dist.mac[i], size);
    if (err != ESP_OK) {
      ESP_LOGW(TAG, "OTA_START to slave %lu failed",
               (unsigned long)s_dist.id[i]);
      return err;
    }
  }

  if (!wait_target_bits(target_mask, OTA_PREPARE_TIMEOUT_MS)) {
    ESP_LOGE(TAG, "timeout waiting for slave OTA ready");
    return ESP_ERR_TIMEOUT;
  }

  uint8_t block_buf[OTA_UART_FLASH_BLOCK_SIZE];
  uint32_t offset = 0;
  uint32_t seq = 0;

  while (offset < size) {
    uint32_t block_len = size - offset;
    if (block_len > OTA_UART_FLASH_BLOCK_SIZE) {
      block_len = OTA_UART_FLASH_BLOCK_SIZE;
    }

    err = esp_partition_read(partition, offset, block_buf, block_len);
    if (err != ESP_OK) {
      ESP_LOGE(TAG, "partition read @%lu failed: %s", (unsigned long)offset,
               esp_err_to_name(err));
      return err;
    }

    uint32_t sent = 0;
    while (sent < block_len) {
      uint32_t chunk = block_len - sent;
      if (chunk > OTA_UART_HOST_CHUNK_SIZE) {
        chunk = OTA_UART_HOST_CHUNK_SIZE;
      }

      for (uint8_t i = 0; i < s_dist.count; i++) {
        err = esp_now_comm_send_ota_payload(s_dist.mac[i], seq,
                                            block_buf + sent, chunk);
        if (err != ESP_OK) {
          return err;
        }
      }
      seq++;
      sent += chunk;
      vTaskDelay(pdMS_TO_TICKS(OTA_PAYLOAD_DELAY_MS));
    }

    const bool full_block = (block_len >= OTA_UART_FLASH_BLOCK_SIZE);
    s_dist.expected_bytes = offset + block_len;

    if (full_block) {
      xEventGroupClearBits(s_eg, target_mask);
      if (!wait_target_bits(target_mask, OTA_BLOCK_TIMEOUT_MS)) {
        ESP_LOGE(TAG, "timeout block ack @%lu bytes",
                 (unsigned long)s_dist.expected_bytes);
        return ESP_ERR_TIMEOUT;
      }
      ESP_LOGI(TAG, "block ack @%lu/%lu (%lu%%)",
               (unsigned long)s_dist.expected_bytes, (unsigned long)size,
               (unsigned long)(s_dist.expected_bytes * 100 / size));
    } else {
      ESP_LOGI(TAG, "final partial block %lu bytes (flush on OTA_END)",
               (unsigned long)block_len);
    }
    offset += block_len;
  }

  xEventGroupClearBits(s_eg, target_mask);
  for (uint8_t i = 0; i < s_dist.count; i++) {
    err = esp_now_comm_send_ota_end(s_dist.mac[i]);
    if (err != ESP_OK) {
      return err;
    }
  }

  if (!wait_target_bits(target_mask, OTA_END_TIMEOUT_MS)) {
    ESP_LOGE(TAG, "timeout waiting for slave OTA success");
    return ESP_ERR_TIMEOUT;
  }

  ESP_LOGI(TAG, "ESP-NOW OTA complete for %u slave(s)", (unsigned)s_dist.count);
  return ESP_OK;
}

esp_err_t ota_espnow_distribute(const esp_partition_t *partition, uint32_t size) {
  if (partition == NULL || size == 0) {
    return ESP_ERR_INVALID_ARG;
  }

  if (collect_targets() == 0) {
    ESP_LOGI(TAG, "no available slaves — skip ESP-NOW OTA");
    return ESP_OK;
  }

  return distribute_image(partition, size);
}