#include "cmd_ota.h"
#include "led_ring.h"
#include "ota_espnow.h"
#include "ota_uart.h"
#include "uart_cmd.h"
#include "esp_log.h"
#include "freertos/FreeRTOS.h"
#include "freertos/idf_additions.h"
#include <stdlib.h>
#include <string.h>

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

#define OTA_PREPARE_STACK 8192
#define OTA_PREPARE_PRIO 5
#define OTA_DIST_STACK 8192
#define OTA_DIST_PRIO 5

/** UART OTA upload to this node (master). */
#define OTA_LED_UART_R 0
#define OTA_LED_UART_G 0
#define OTA_LED_UART_B 255
/** ESP-NOW distribution from master to slaves. */
#define OTA_LED_ESPNOW_TX_R 0
#define OTA_LED_ESPNOW_TX_G 255
#define OTA_LED_ESPNOW_TX_B 0

typedef struct {
  uint32_t written;
  int slot;
} ota_dist_job_t;
static void send_ota_status(ota_uart_status_t status, uint32_t err_code) {
  alox_UartMessage response;
  uart_cmd_init_response(&response, alox_MessageType_OTA_STATUS,
                         alox_UartMessage_ota_status_tag);
  response.payload.ota_status.status = (uint32_t)status;
  response.payload.ota_status.bytes_written = ota_uart_bytes_written();
  int slot = ota_uart_target_slot();
  response.payload.ota_status.target_slot =
      slot >= 0 ? (uint32_t)slot : 0;
  response.payload.ota_status.error = err_code;
  uart_cmd_send(&response, TAG);
}

static void send_ota_failed(uint32_t err_code) {
  led_ring_ota_failed();
  send_ota_status(OTA_UART_ST_FAILED, err_code);
}

static void send_ota_distributing(uint32_t kind, uint32_t bytes_done,
                                 uint32_t target_slot) {
  alox_UartMessage response;
  uart_cmd_init_response(&response, alox_MessageType_OTA_STATUS,
                         alox_UartMessage_ota_status_tag);
  response.payload.ota_status.status = (uint32_t)OTA_UART_ST_DISTRIBUTING;
  response.payload.ota_status.bytes_written = bytes_done;
  response.payload.ota_status.target_slot = target_slot;
  response.payload.ota_status.error = kind;
  uart_cmd_send(&response, TAG);
}

static void ota_dist_aggregate(uint32_t bytes_done, uint32_t total_bytes,
                               uint8_t slave_count) {
  (void)slave_count;
  led_ring_show_ota_progress(bytes_done, total_bytes, OTA_LED_ESPNOW_TX_R,
                             OTA_LED_ESPNOW_TX_G, OTA_LED_ESPNOW_TX_B);
  send_ota_distributing(OTA_DIST_AGGREGATE, bytes_done, (uint32_t)slave_count);
}

static void ota_dist_per_slave(uint32_t slave_id, uint32_t bytes_done,
                             uint32_t total_bytes) {
  (void)total_bytes;
  send_ota_distributing(OTA_DIST_PER_SLAVE, bytes_done, slave_id);
}

static const ota_espnow_progress_cbs_t s_dist_progress = {
    .aggregate = ota_dist_aggregate,
    .per_slave = ota_dist_per_slave,
};

static void ota_prepare_task(void *param) {
  uint32_t total_size = (uint32_t)(uintptr_t)param;

  int slot = ota_uart_prepare(total_size);
  if (slot < 0) {
    send_ota_failed(1);
    vTaskDelete(NULL);
    return;
  }

  alox_UartMessage response;
  uart_cmd_init_response(&response, alox_MessageType_OTA_STATUS,
                         alox_UartMessage_ota_status_tag);
  response.payload.ota_status.status = (uint32_t)OTA_UART_ST_READY;
  response.payload.ota_status.bytes_written = 0;
  response.payload.ota_status.target_slot = (uint32_t)slot;
  response.payload.ota_status.error = 0;
  uart_cmd_send(&response, TAG);

  led_ring_show_ota_progress(0, total_size, OTA_LED_UART_R, OTA_LED_UART_G,
                             OTA_LED_UART_B);

  vTaskDelete(NULL);
}

static void handle_ota_start(const uint8_t *data, size_t len) {
  alox_UartMessage uart_msg;
  alox_OtaStartPayload req = alox_OtaStartPayload_init_zero;

  if (uart_cmd_decode(data, len, &uart_msg) != ESP_OK) {
    send_ota_failed( 2);
    return;
  }

  const alox_OtaStartPayload *req_ptr =
      UART_CMD_REQ(&uart_msg, alox_UartMessage_ota_start_tag, ota_start);
  if (req_ptr == NULL) {
    ESP_LOGW(TAG, "OTA_START: missing ota_start payload");
    send_ota_failed(3);
    return;
  }
  req = *req_ptr;

  if (req.total_size == 0) {
    ESP_LOGW(TAG, "OTA_START: total_size required");
    send_ota_failed(3);
    return;
  }

  if (ota_uart_is_active()) {
    ESP_LOGW(TAG, "OTA_START while session active");
    send_ota_failed(4);
    return;
  }

  send_ota_status(OTA_UART_ST_PREPARING, 0);

  if (xTaskCreate(ota_prepare_task, "ota_prepare", OTA_PREPARE_STACK,
                  (void *)(uintptr_t)req.total_size, OTA_PREPARE_PRIO,
                  NULL) != pdPASS) {
    ESP_LOGE(TAG, "failed to create ota_prepare task");
    send_ota_failed(5);
  }
}

static void handle_ota_payload(const uint8_t *data, size_t len) {
  alox_UartMessage uart_msg;

  if (uart_cmd_decode(data, len, &uart_msg) != ESP_OK) {
    ESP_LOGW(TAG, "OTA_PAYLOAD decode failed");
    send_ota_failed( 10);
    return;
  }

  const alox_OtaPayload *req_ptr =
      UART_CMD_REQ(&uart_msg, alox_UartMessage_ota_payload_tag, ota_payload);
  if (req_ptr == NULL) {
    ESP_LOGW(TAG, "OTA_PAYLOAD: missing ota_payload (which=%u)",
             (unsigned)uart_msg.which_payload);
    send_ota_failed( 11);
    return;
  }

  if (req_ptr->data.size == 0) {
    ESP_LOGW(TAG, "OTA_PAYLOAD: empty data (seq=%lu)",
             (unsigned long)req_ptr->seq);
    send_ota_failed( 11);
    return;
  }

  if (!ota_uart_is_active()) {
    ESP_LOGW(TAG, "OTA_PAYLOAD without active session (seq=%lu)",
             (unsigned long)req_ptr->seq);
    send_ota_failed( 12);
    return;
  }

  ota_feed_result_t r =
      ota_uart_feed(req_ptr->data.bytes, req_ptr->data.size);
  if (r == OTA_FEED_ERROR) {
    send_ota_failed( 13);
    return;
  }
  if (r == OTA_FEED_BLOCK_WRITTEN) {
    uint32_t total = ota_uart_total_size();
    uint32_t done = ota_uart_bytes_written();
    ESP_LOGI(TAG, "OTA block ack (%lu bytes in flash)",
             (unsigned long)done);
    led_ring_show_ota_progress(done, total, OTA_LED_UART_R, OTA_LED_UART_G,
                               OTA_LED_UART_B);
    send_ota_status(OTA_UART_ST_BLOCK_ACK, 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_UART_R, OTA_LED_UART_G, OTA_LED_UART_B);
    }
  }
}

static void ota_distribute_task(void *param) {
  ota_dist_job_t *job = (ota_dist_job_t *)param;
  if (job == NULL) {
    vTaskDelete(NULL);
    return;
  }

  const esp_partition_t *part = NULL;
  uint32_t image_size = 0;
  if (!ota_uart_get_staged_image(&part, &image_size)) {
    send_ota_failed( 30);
    free(job);
    vTaskDelete(NULL);
    return;
  }

  led_ring_show_ota_progress(0, image_size, OTA_LED_ESPNOW_TX_R, OTA_LED_ESPNOW_TX_G,
                             OTA_LED_ESPNOW_TX_B);
  send_ota_distributing(OTA_DIST_AGGREGATE, 0, 0);

  esp_err_t err = ota_espnow_distribute(part, image_size, &s_dist_progress);
  if (err != ESP_OK) {
    ESP_LOGE(TAG, "slave OTA distribution failed: %s", esp_err_to_name(err));
    ota_uart_clear_staged();
    send_ota_failed(31);
    free(job);
    vTaskDelete(NULL);
    return;
  }

  err = ota_uart_apply_boot();
  if (err != ESP_OK) {
    send_ota_failed((uint32_t)err);
    free(job);
    vTaskDelete(NULL);
    return;
  }

  led_ring_show_ota_progress(image_size, image_size, OTA_LED_ESPNOW_TX_R,
                             OTA_LED_ESPNOW_TX_G, OTA_LED_ESPNOW_TX_B);

  alox_UartMessage response;
  uart_cmd_init_response(&response, alox_MessageType_OTA_STATUS,
                         alox_UartMessage_ota_status_tag);
  response.payload.ota_status.status = (uint32_t)OTA_UART_ST_SUCCESS;
  response.payload.ota_status.bytes_written = job->written;
  response.payload.ota_status.target_slot =
      job->slot >= 0 ? (uint32_t)job->slot : 0;
  response.payload.ota_status.error = 0;
  uart_cmd_send(&response, TAG);

  led_ring_ota_success();
  free(job);
  vTaskDelete(NULL);
}

static void handle_ota_end(const uint8_t *data, size_t len) {
  (void)data;
  (void)len;

  if (!ota_uart_is_active()) {
    send_ota_failed( 20);
    return;
  }

  ota_dist_job_t *job = calloc(1, sizeof(*job));
  if (job == NULL) {
    send_ota_failed( 21);
    return;
  }
  job->written = ota_uart_bytes_written();
  job->slot = ota_uart_target_slot();
  uint32_t uart_total = ota_uart_total_size();

  bool success = false;
  esp_err_t err = ota_uart_finish(false, &success);
  if (err != ESP_OK || !success) {
    send_ota_failed((uint32_t)err);
    free(job);
    return;
  }

  if (uart_total > 0) {
    led_ring_show_ota_progress(job->written, uart_total, OTA_LED_UART_R,
                               OTA_LED_UART_G, OTA_LED_UART_B);
  }

  if (xTaskCreate(ota_distribute_task, "ota_dist", OTA_DIST_STACK, job,
                  OTA_DIST_PRIO, NULL) != pdPASS) {
    ESP_LOGE(TAG, "failed to create ota_dist task");
    send_ota_failed( 22);
    free(job);
  }
}

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

  const esp_partition_t *part = NULL;
  uint32_t image_size = 0;
  if (!ota_uart_get_staged_image(&part, &image_size)) {
    send_ota_failed(41);
    vTaskDelete(NULL);
    return;
  }

  esp_err_t err = ota_espnow_distribute(part, image_size, &s_dist_progress);
  if (err != ESP_OK) {
    send_ota_failed(42);
    vTaskDelete(NULL);
    return;
  }

  err = ota_uart_apply_boot();
  if (err != ESP_OK) {
    send_ota_failed((uint32_t)err);
    vTaskDelete(NULL);
    return;
  }

  alox_UartMessage response;
  uart_cmd_init_response(&response, alox_MessageType_OTA_STATUS,
                         alox_UartMessage_ota_status_tag);
  response.payload.ota_status.status = (uint32_t)OTA_UART_ST_SUCCESS;
  response.payload.ota_status.bytes_written = image_size;
  response.payload.ota_status.error = 0;
  uart_cmd_send(&response, TAG);
  led_ring_ota_success();
  vTaskDelete(NULL);
}

static void handle_ota_start_espnow(const uint8_t *data, size_t len) {
  (void)data;
  (void)len;

  if (ota_uart_is_active()) {
    send_ota_failed(40);
    return;
  }

  const esp_partition_t *part = NULL;
  uint32_t image_size = 0;
  if (!ota_uart_get_staged_image(&part, &image_size)) {
    send_ota_failed(41);
    return;
  }

  if (xTaskCreate(ota_start_espnow_task, "ota_espnow", OTA_DIST_STACK, NULL,
                  OTA_DIST_PRIO, NULL) != pdPASS) {
    ESP_LOGE(TAG, "failed to create ota_start_espnow task");
    send_ota_failed(43);
  }
}

void cmd_ota_register(void) {
  uart_cmd_register(alox_MessageType_OTA_START, handle_ota_start);
  uart_cmd_register(alox_MessageType_OTA_PAYLOAD, handle_ota_payload);
  uart_cmd_register(alox_MessageType_OTA_END, handle_ota_end);
  uart_cmd_register(alox_MessageType_OTA_START_ESPNOW, handle_ota_start_espnow);
}