#include "ota_espnow.h"
#include "app_config.h"
#include "led_ring.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

/** ESP-NOW OTA receive on slave (blue progress bar). */
#define OTA_LED_ESPNOW_RX_R 0
#define OTA_LED_ESPNOW_RX_G 0
#define OTA_LED_ESPNOW_RX_B 255

#define OTA_MAX_TARGETS CLIENT_REGISTRY_MAX

#define OTA_SLAVE_WORK_QUEUE_LEN 12
#define OTA_SLAVE_WORK_STACK 8192
#define OTA_SLAVE_WORK_PRIO 5

typedef enum {
  OTA_SLAVE_WORK_STATUS = 1,
  OTA_SLAVE_WORK_PAYLOAD,
  OTA_SLAVE_WORK_END,
} ota_slave_work_op_t;

typedef struct {
  ota_slave_work_op_t op;
  uint8_t master_mac[6];
  uint32_t status;
  uint32_t bytes_written;
  uint32_t error;
  alox_EspNowOtaPayload payload;
} ota_slave_work_t;

static EventGroupHandle_t s_eg;
static QueueHandle_t s_slave_work_queue;
static bool s_distribution_active;

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

static ota_dist_t s_dist;

typedef struct {
  uint32_t client_id;
  uint32_t bytes_written;
  uint32_t status;
  uint32_t error;
} ota_prog_entry_t;

static struct {
  bool active;
  uint32_t total_bytes;
  uint32_t aggregate_bytes;
  uint8_t count;
  ota_prog_entry_t entries[OTA_MAX_TARGETS];
} s_prog;

static void prog_begin(uint32_t total_bytes) {
  s_prog.active = true;
  s_prog.total_bytes = total_bytes;
  s_prog.aggregate_bytes = 0;
  s_prog.count = s_dist.count;
  for (uint8_t i = 0; i < s_dist.count; i++) {
    s_prog.entries[i].client_id = s_dist.id[i];
    s_prog.entries[i].bytes_written = 0;
    s_prog.entries[i].status = OTA_ST_PREPARING;
    s_prog.entries[i].error = 0;
  }
}

static void prog_end(void) { s_prog.active = false; }

static void prog_set_aggregate(uint32_t bytes_done) {
  s_prog.aggregate_bytes = bytes_done;
}

static void prog_update_idx(int idx, uint32_t status, uint32_t bytes,
                            uint32_t error) {
  if (idx < 0 || idx >= (int)s_prog.count) {
    return;
  }
  ota_prog_entry_t *e = &s_prog.entries[idx];
  e->status = status;
  if (bytes > e->bytes_written) {
    e->bytes_written = bytes;
  }
  if (error != 0) {
    e->error = error;
  }
}

void ota_espnow_progress_query(uint32_t filter_client_id,
                               alox_OtaSlaveProgressResponse *out) {
  if (out == NULL) {
    return;
  }
  *out = (alox_OtaSlaveProgressResponse)alox_OtaSlaveProgressResponse_init_zero;
  out->active = s_prog.active;
  out->total_bytes = s_prog.total_bytes;
  out->aggregate_bytes = s_prog.aggregate_bytes;
  out->slave_count = s_prog.count;

  for (uint8_t i = 0; i < s_prog.count; i++) {
    const ota_prog_entry_t *e = &s_prog.entries[i];
    if (filter_client_id != 0 && e->client_id != filter_client_id) {
      continue;
    }
    if (out->slaves_count >=
        sizeof(out->slaves) / sizeof(out->slaves[0])) {
      break;
    }
    alox_OtaSlaveProgressEntry *dst = &out->slaves[out->slaves_count++];
    dst->client_id = e->client_id;
    dst->bytes_written = e->bytes_written;
    dst->total_bytes = s_prog.total_bytes;
    dst->status = e->status;
    dst->error = e->error;
  }
}

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

bool ota_espnow_distribution_active(void) { return s_distribution_active; }

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 bool queue_slave_work(const ota_slave_work_t *work) {
  if (work == NULL || s_slave_work_queue == NULL) {
    return false;
  }
  if (xQueueSend(s_slave_work_queue, work, 0) != pdTRUE) {
    ESP_LOGW(TAG, "slave OTA work queue full (op=%d)", (int)work->op);
    return false;
  }
  return true;
}

static void queue_slave_status(const uint8_t master_mac[6], uint32_t status,
                               uint32_t bytes_written, uint32_t error) {
  ota_slave_work_t work = {
      .op = OTA_SLAVE_WORK_STATUS,
      .status = status,
      .bytes_written = bytes_written,
      .error = error,
  };
  memcpy(work.master_mac, master_mac, 6);
  (void)queue_slave_work(&work);
}

static void process_slave_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) {
    led_ring_ota_failed();
    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();
    uint32_t total = ota_uart_total_size();
    ESP_LOGI(TAG, "block written %lu bytes -> ack master", (unsigned long)written);
    led_ring_show_ota_progress(written, total, OTA_LED_ESPNOW_RX_R, OTA_LED_ESPNOW_RX_G,
                               OTA_LED_ESPNOW_RX_B);
    send_slave_status(master_mac, OTA_ST_BLOCK_ACK, written, 0);
    return;
  }

  if (r == OTA_FEED_OK) {
    uint32_t total = ota_uart_total_size();
    if (total > 0) {
      led_ring_show_ota_progress(ota_uart_bytes_received(), total,
                                 OTA_LED_ESPNOW_RX_R, OTA_LED_ESPNOW_RX_G,
                                 OTA_LED_ESPNOW_RX_B);
    }
  }
}

static void process_slave_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) {
    led_ring_ota_failed();
    send_slave_status(master_mac, OTA_ST_FAILED, written, (uint32_t)err);
    return;
  }

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

static void ota_slave_work_task(void *param) {
  (void)param;
  ota_slave_work_t work;

  while (1) {
    if (xQueueReceive(s_slave_work_queue, &work, portMAX_DELAY) != pdTRUE) {
      continue;
    }

    switch (work.op) {
    case OTA_SLAVE_WORK_STATUS:
      send_slave_status(work.master_mac, work.status, work.bytes_written,
                        work.error);
      break;
    case OTA_SLAVE_WORK_PAYLOAD:
      process_slave_payload(work.master_mac, &work.payload);
      break;
    case OTA_SLAVE_WORK_END:
      process_slave_end(work.master_mac);
      break;
    default:
      break;
    }
  }
}

void ota_espnow_slave_init(void) {
  if (s_slave_work_queue != NULL) {
    return;
  }
  s_slave_work_queue = xQueueCreate(OTA_SLAVE_WORK_QUEUE_LEN, sizeof(ota_slave_work_t));
  if (s_slave_work_queue == NULL) {
    ESP_LOGE(TAG, "failed to create slave OTA work queue");
    return;
  }
  if (xTaskCreate(ota_slave_work_task, "ota_slave_wrk", OTA_SLAVE_WORK_STACK, NULL,
                  OTA_SLAVE_WORK_PRIO, NULL) != pdPASS) {
    ESP_LOGE(TAG, "failed to create slave OTA work task");
  }
}

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);
  led_ring_show_ota_progress(0, total_size, OTA_LED_ESPNOW_RX_R, OTA_LED_ESPNOW_RX_G,
                             OTA_LED_ESPNOW_RX_B);
  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()) {
    queue_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) {
    queue_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) {
    queue_slave_status(master_mac, OTA_ST_FAILED, 0, 11);
    return;
  }

  ota_slave_work_t work = {.op = OTA_SLAVE_WORK_PAYLOAD, .payload = *payload};
  memcpy(work.master_mac, master_mac, 6);
  if (!queue_slave_work(&work)) {
    queue_slave_status(master_mac, OTA_ST_FAILED, 0, 14);
  }
}

void ota_espnow_slave_on_end(const uint8_t master_mac[6]) {
  ota_slave_work_t work = {.op = OTA_SLAVE_WORK_END};
  memcpy(work.master_mac, master_mac, 6);
  if (!queue_slave_work(&work)) {
    queue_slave_status(master_mac, OTA_ST_FAILED, 0, 15);
  }
}

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:
    prog_update_idx(idx, OTA_ST_READY, 0, 0);
    xEventGroupSetBits(s_eg, bit);
    break;
  case OTA_ST_BLOCK_ACK:
    prog_update_idx(idx, OTA_ST_BLOCK_ACK, status->bytes_written, 0);
    if (s_dist.progress.per_slave != NULL) {
      s_dist.progress.per_slave(s_dist.id[idx], status->bytes_written,
                                s_dist.total_bytes);
    }
    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:
    prog_update_idx(idx, OTA_ST_SUCCESS, status->bytes_written, 0);
    xEventGroupSetBits(s_eg, bit);
    break;
  case OTA_ST_FAILED:
    prog_update_idx(idx, OTA_ST_FAILED, status->bytes_written,
                    status->error);
    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,
                                  const ota_espnow_progress_cbs_t *progress) {
  if (s_eg == NULL) {
    s_eg = xEventGroupCreate();
    if (s_eg == NULL) {
      return ESP_ERR_NO_MEM;
    }
  }

  s_distribution_active = true;

  memset(&s_dist.progress, 0, sizeof(s_dist.progress));
  if (progress != NULL) {
    s_dist.progress = *progress;
  }
  s_dist.total_bytes = size;
  prog_begin(size);

  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]);
      prog_end();
      s_distribution_active = false;
      return err;
    }
  }

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

  prog_set_aggregate(0);
  if (s_dist.progress.aggregate != NULL) {
    s_dist.progress.aggregate(0, size, s_dist.count);
  }

  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));
      prog_end();
      s_distribution_active = false;
      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) {
          prog_end();
          s_distribution_active = false;
          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);
        prog_end();
        s_distribution_active = false;
        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;
    prog_set_aggregate(offset);
    if (s_dist.progress.aggregate != NULL) {
      s_dist.progress.aggregate(offset, size, s_dist.count);
    }
  }

  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) {
      prog_end();
      s_distribution_active = false;
      return err;
    }
  }

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

  prog_set_aggregate(size);
  prog_end();
  s_distribution_active = false;
  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,
                                const ota_espnow_progress_cbs_t *progress) {
  if (partition == NULL || size == 0) {
    return ESP_ERR_INVALID_ARG;
  }

  if (collect_targets() == 0) {
    ESP_LOGI(TAG, "no available slaves — skip ESP-NOW OTA");
    memset(&s_prog, 0, sizeof(s_prog));
    s_prog.total_bytes = size;
    if (progress != NULL && progress->aggregate != NULL) {
      progress->aggregate(size, size, 0);
    }
    return ESP_OK;
  }

  return distribute_image(partition, size, progress);
}