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

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

#define OTA_PREPARE_STACK 8192
#define OTA_PREPARE_PRIO 5

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 ota_prepare_task(void *param) {
  uint32_t total_size = (uint32_t)(uintptr_t)param;

  send_ota_status(OTA_UART_ST_PREPARING, 0);

  int slot = ota_uart_prepare(total_size);
  if (slot < 0) {
    send_ota_status(OTA_UART_ST_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);

  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_status(OTA_UART_ST_FAILED, 2);
    return;
  }

  const alox_OtaStartPayload *req_ptr =
      UART_CMD_REQ(&uart_msg, alox_UartMessage_ota_start_tag, ota_start);
  if (req_ptr != NULL) {
    req = *req_ptr;
  }

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

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

  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_status(OTA_UART_ST_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_status(OTA_UART_ST_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_status(OTA_UART_ST_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_status(OTA_UART_ST_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_status(OTA_UART_ST_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_status(OTA_UART_ST_FAILED, 13);
    return;
  }
  if (r == OTA_FEED_BLOCK_WRITTEN) {
    ESP_LOGI(TAG, "OTA block ack (%lu bytes in flash)",
             (unsigned long)ota_uart_bytes_written());
    send_ota_status(OTA_UART_ST_BLOCK_ACK, 0);
  }
}

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

  if (!ota_uart_is_active()) {
    send_ota_status(OTA_UART_ST_FAILED, 20);
    return;
  }

  uint32_t written = ota_uart_bytes_written();
  int slot = ota_uart_target_slot();
  bool success = false;
  esp_err_t err = ota_uart_finish(&success);
  if (err != ESP_OK || !success) {
    send_ota_status(OTA_UART_ST_FAILED, (uint32_t)err);
    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 = written;
  response.payload.ota_status.target_slot = slot >= 0 ? (uint32_t)slot : 0;
  response.payload.ota_status.error = 0;
  uart_cmd_send(&response, TAG);
}

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