#include "ota_update.h"
#include "communication_handler.h"
#include "driver/uart.h"
#include "esp_log.h"
#include "esp_ota_ops.h"
#include "esp_partition.h"
#include "esp_system.h"
#include "freertos/FreeRTOS.h"
#include "freertos/queue.h"
#include "freertos/task.h"
#include "message_builder.h"
#include "message_handler.h"
#include "uart_handler.h"
#include "uart_msg_ids.h"

extern uint32_t g_uart_firmware_total_size;

static const char *TAG = "ALOX - OTA_SLAVE";

static QueueHandle_t ota_task_queue = NULL;

static esp_ota_handle_t update_handle = 0;
static uint8_t update_buffer[UPDATE_BUFFER_SIZE];
static uint32_t chunk_bitmask;
static uint16_t current_block_id;
static uint16_t update_buffer_write_index;
static uint32_t update_size;
static uint16_t sequenz_counter; // how often the update buffer gets written
void ota_task(void *pvParameter) {
  ESP_LOGI(TAG, "ota_task started");
  ota_task_queue_message_t msg;
  while (1) {
    if (xQueueReceive(ota_task_queue, &msg, portMAX_DELAY)) {
      ESP_LOGI(TAG, "ota_task received command: %d", msg.command);
      BaseMessage replyMessage = {};

      switch (msg.command) {
      case OTA_COMMAND_MASTER_READY: {
        int part = prepare_ota_update();
        OTA_READY_TO_RECEIVE_Payload payload = {};
        if (part < 0) {
          payload.status = 1; // ERROR
        } else {
          payload.status = 0; // READY
        }
        replyMessage = MessageBuilder(
            OTA_READY_TO_RECEIVE, *(PayloadUnion *)&payload, sizeof(payload));
        esp_now_send(msg.mac_addr, (uint8_t *)&replyMessage,
                     sizeof(BaseMessage));
        break;
      }
      case OTA_COMMAND_PREPARE: {
        OTA_PREPARE_ACKNOWLEDGED_Payload payload = {};
        replyMessage =
            MessageBuilder(OTA_PREPARE_ACKNOWLEDGED, *(PayloadUnion *)&payload,
                           sizeof(payload));
        esp_now_send(msg.mac_addr, (uint8_t *)&replyMessage,
                     sizeof(BaseMessage));
        chunk_bitmask = 0;
        current_block_id = 0;
        break;
      }
      case OTA_COMMAND_CHUNK: {
        OTA_CHUNK_Payload *payload = &msg.payload.ota_chunk_payload;
        ESP_LOGI(TAG, "OTA_CHUNK: block_id=%d, chunk_id=%d", payload->block_id,
                 payload->chunk_id);
        if (payload->block_id == current_block_id) {
          if (payload->chunk_id < UPDATE_MAX_SEQUENZES) { // UPDATE_MAX_SEQUENZES is 20
            memcpy(&update_buffer[payload->chunk_id * UPDATE_PAYLOAD_SIZE], payload->data, UPDATE_PAYLOAD_SIZE);
            chunk_bitmask |= (1 << payload->chunk_id);
            ESP_LOGD(TAG, "OTA_CHUNK: memcpy done, chunk_bitmask=0x%x",
                     chunk_bitmask);
          } else {
            ESP_LOGE(TAG, "OTA_CHUNK: Invalid chunk_id %d. Max allowed is %d", payload->chunk_id, UPDATE_MAX_SEQUENZES - 1);
            // Optionally, send an error back to master or reset state
          }
        } else {
          ESP_LOGW(TAG, "OTA_CHUNK: Mismatched block_id. Expected %d, got %d",
                   current_block_id, payload->block_id);
        }
        ESP_LOGD(TAG, "OTA_CHUNK: Calling vTaskDelay(10ms)");
        //vTaskDelay(pdMS_TO_TICKS(10));
        break;
      }
      case OTA_COMMAND_REQUEST_STATUS: {
        OTA_REQUEST_BLOCK_STATUS_Payload *payload =
            &msg.payload.ota_request_block_status_payload;
        if (payload->block_id == current_block_id) {
          OTA_BLOCK_STATUS_REPORT_Payload replyPayload = {
              .block_id = current_block_id,
              .chunk_bitmask = chunk_bitmask,
          };
          replyMessage = MessageBuilder(OTA_BLOCK_STATUS_REPORT,
                                        *(PayloadUnion *)&replyPayload,
                                        sizeof(replyPayload));
          esp_now_send(msg.mac_addr, (uint8_t *)&replyMessage,
                       sizeof(BaseMessage));
        }
        break;
      }
      case OTA_COMMAND_COMMIT: {
        OTA_COMMIT_BLOCK_Payload *payload =
            &msg.payload.ota_commit_block_payload;
        ESP_LOGI(TAG, "OTA_COMMIT: block_id=%d", payload->block_id);
        if (payload->block_id == current_block_id) {
          uint32_t write_size = UPDATE_BUFFER_SIZE;
          if ((current_block_id + 1) * UPDATE_BUFFER_SIZE > total_update_size) {
            write_size = total_update_size % UPDATE_BUFFER_SIZE;
          }
          ESP_LOGI(TAG, "OTA_COMMIT: Writing %lu bytes to partition.",
                   write_size);
          esp_err_t err =
              esp_ota_write(update_handle, update_buffer, write_size);
          if (err != ESP_OK) {
            ESP_LOGE(TAG, "OTA_COMMIT: esp_ota_write failed (%s)",
                     esp_err_to_name(err));
          }
          ESP_LOGD(TAG, "OTA_COMMIT: Calling vTaskDelay(10ms)");
          vTaskDelay(pdMS_TO_TICKS(10));

          OTA_BLOCK_COMMITTED_Payload replyPayload = {
              .block_id = current_block_id,
          };
          replyMessage = MessageBuilder(OTA_BLOCK_COMMITTED,
                                        *(PayloadUnion *)&replyPayload,
                                        sizeof(replyPayload));
          esp_now_send(msg.mac_addr, (uint8_t *)&replyMessage,
                       sizeof(BaseMessage));

          current_block_id++;
          chunk_bitmask = 0;
          ESP_LOGI(TAG, "OTA_COMMIT: Block %d committed. Next block_id=%d",
                   payload->block_id, current_block_id);
        } else {
          ESP_LOGW(TAG, "OTA_COMMIT: Mismatched block_id. Expected %d, got %d",
                   current_block_id, payload->block_id);
        }
        break;

      case OTA_COMMAND_FINISH: {
        esp_err_t err = end_ota_update();
        OTA_UPDATE_STATUS_Payload ota_status_payload = {};
        ota_status_payload.status = (err == ESP_OK) ? 0 : 1;
        replyMessage = MessageBuilder(OTA_UPDATE_STATUS,
                                      *(PayloadUnion *)&ota_status_payload,
                                      sizeof(ota_status_payload));
        esp_now_send(msg.mac_addr, (uint8_t *)&replyMessage,
                     sizeof(BaseMessage));
        if (err == ESP_OK) {
          esp_restart();
        }
        break;
      }
      default:
        ESP_LOGW(TAG, "Unknown command received: %d", msg.command);
        break;
      }
      }
    }
  }
}

static uint32_t total_update_size = 0;

void start_uart_update(uint8_t msgid, const uint8_t *payload,
                       size_t payload_len, uint8_t *send_payload_buffer,
                       size_t send_payload_buffer_size, uint8_t *send_buffer,
                       size_t send_buffer_size) {
  ESP_LOGI(TAG, "OTA Update Start Uart Command");

  vTaskPrioritySet(NULL, 2);

  update_size = 0;
  update_buffer_write_index = 0;
  sequenz_counter = 0;

  // Assuming payload contains total_size as uint32_2 at the beginning
  if (payload_len >= sizeof(uint32_t)) {
    g_uart_firmware_total_size = *(uint32_t *)payload;
    ESP_LOGI(TAG, "UART OTA: Received total firmware size: %lu bytes", g_uart_firmware_total_size);
  } else {
    ESP_LOGE(TAG, "UART OTA: Payload too short for total firmware size.");
    g_uart_firmware_total_size = 0; // Reset or handle error appropriately
  }

  int part = prepare_ota_update();

  if (part < 0) {
    send_payload_buffer[1] = (part * -1) & 0xff;
  } else {
    send_payload_buffer[0] = part & 0xff;
  }

  int send_payload_len = 2;
  int len = build_message(UART_OTA_START, send_payload_buffer, send_payload_len,
                          send_buffer, send_buffer_size);
  if (len < 0) {
    ESP_LOGE(TAG,
             "Error Building UART Message: payload_len, %d, sendbuffer_size: "
             "%d, mes_len(error): %d",
             payload_len, send_buffer_size, len);
    return;
  }

  uart_write_bytes(MASTER_UART, send_buffer, len);
}

esp_err_t write_ota_update(uint32_t write_len, const uint8_t *payload) {
  if (update_buffer_write_index + write_len > UPDATE_BUFFER_SIZE) {
    esp_err_t err =
        esp_ota_write(update_handle, update_buffer, update_buffer_write_index);
    if (err != ESP_OK) {
      return err;
    }
    update_buffer_write_index = 0;
    sequenz_counter++;
  }

  memcpy(&update_buffer[update_buffer_write_index], payload, write_len);
  update_buffer_write_index += write_len;
  return ESP_OK;
}

void payload_uart_update(uint8_t msgid, const uint8_t *payload_data_from_uart,
                         size_t total_payload_len_from_uart,
                         uint8_t *send_payload_buffer,
                         size_t send_payload_buffer_size, uint8_t *send_buffer,
                         size_t send_buffer_size) {

  const uint8_t *actual_firmware_data = payload_data_from_uart;

  uint32_t write_len = total_payload_len_from_uart;

  if (update_size == 0) {
    ESP_LOGI(TAG, "First chunk received. write_len: %d", write_len);
  }

  update_size += write_len;

  esp_err_t err = write_ota_update(write_len, actual_firmware_data);

  if (err != ESP_OK) {
    ESP_LOGE(TAG, "GOT ESP ERROR WRITE OTA %d", err);
  }

  size_t send_payload_len = 4;
  memcpy(send_payload_buffer, &sequenz_counter, 2);
  memcpy(&send_payload_buffer[2], &update_buffer_write_index, 2);
  send_payload_buffer[4] = 0x00; // error

  int len = build_message(UART_OTA_PAYLOAD, send_payload_buffer,
                          send_payload_len, send_buffer, send_buffer_size);
  if (len < 0) {
    ESP_LOGE(TAG,
             "Error Building UART Message: payload_len, %d, sendbuffer_size: "
             "%d, mes_len(error): %d",
             total_payload_len_from_uart, send_buffer_size, len);
    return;
  }

  uart_write_bytes(MASTER_UART, send_buffer, len);
}

void end_uart_update(uint8_t msgid, const uint8_t *payload, size_t payload_len,
                     uint8_t *send_payload_buffer,
                     size_t send_payload_buffer_size, uint8_t *send_buffer,
                     size_t send_buffer_size) {
  ESP_LOGI(TAG, "OTA Update End Uart Command");

  esp_err_t err = end_ota_update();

  int send_payload_len = 1;
  send_payload_buffer[0] = err & 0xff;
  int len = build_message(UART_OTA_END, send_payload_buffer, send_payload_len,
                          send_buffer, send_buffer_size);
  if (len < 0) {
    ESP_LOGE(TAG,
             "Error Building UART Message: payload_len, %d, sendbuffer_size: "
             "%d, mes_len(error): %d",
             payload_len, send_buffer_size, len);
    return;
  }

  uart_write_bytes(MASTER_UART, send_buffer, len);

  vTaskPrioritySet(NULL, 1);
}

void init_ota() {
  ota_task_queue = xQueueCreate(50, sizeof(ota_task_queue_message_t));
  xTaskCreate(ota_task, "ota_task", 8192, NULL, 4, NULL);
  RegisterCallback(UART_OTA_START, start_uart_update);
  RegisterCallback(UART_OTA_PAYLOAD, payload_uart_update);
  RegisterCallback(UART_OTA_END, end_uart_update);
}

int prepare_ota_update() {
  const esp_partition_t *running = esp_ota_get_running_partition();
  ESP_LOGI(TAG, "Running Partition: %s", running->label);
  const esp_partition_t *update_partition =
      esp_ota_get_next_update_partition(NULL);
  if (update_partition == NULL) {
    ESP_LOGE(TAG, "Failed to find OTA partition.");
    return -1;
  }

  ESP_LOGI(TAG, "Writing OTA Update to Partition: %s", update_partition->label);

  esp_err_t err =
      esp_ota_begin(update_partition, OTA_SIZE_UNKNOWN, &update_handle);
  if (err != ESP_OK) {
    ESP_LOGE(TAG, "esp_ota_begin failed (%s)", esp_err_to_name(err));
    return -2;
  }
  return 0;
}

esp_err_t end_ota_update() {
  if (update_buffer_write_index > 0) {
    esp_err_t err =
        esp_ota_write(update_handle, update_buffer, update_buffer_write_index);
    vTaskDelay(1);
    if (err != ESP_OK) {
      ESP_LOGE(TAG, "Error writing remaining data to partition: %s",
               esp_err_to_name(err));
      return err;
    }
  }

  g_uart_firmware_total_size = update_size;
  esp_err_t err = esp_ota_end(update_handle);
  if (err != ESP_OK) {
    ESP_LOGE(TAG, "esp_ota_end failed: %s", esp_err_to_name(err));
    ESP_LOGI(TAG, "Total blocks written: %u, Last partial block size: %u", sequenz_counter, update_buffer_write_index);
    return err;
  }

  const esp_partition_t *update_partition =
      esp_ota_get_next_update_partition(NULL);
  err = esp_ota_set_boot_partition(update_partition);
  if (err != ESP_OK) {
    ESP_LOGE(TAG, "esp_ota_set_boot_partition failed: %s",
             esp_err_to_name(err));
  }
  return err;
}

void slave_Prep_Upgrade_Callback(const esp_now_recv_info_t *esp_now_info,
                                 const uint8_t *data, int data_len) {
  const BaseMessage *message = (const BaseMessage *)data;
  const OTA_PREPARE_FOR_UPDATE_Payload *payload =
      &message->payload.ota_prepare_for_update_payload;
  total_update_size = payload->total_size;

  ota_task_queue_message_t msg = {.command = OTA_COMMAND_PREPARE};
  memcpy(msg.mac_addr, esp_now_info->src_addr, ESP_NOW_ETH_ALEN);
  if (xQueueSend(ota_task_queue, &msg, pdMS_TO_TICKS(10)) != pdTRUE) {
    ESP_LOGE(TAG, "Failed to send prepare command to OTA task");
  }
}

void slave_ota_chunk_callback(const esp_now_recv_info_t *esp_now_info,
                              const uint8_t *data, int data_len) {
  const BaseMessage *message = (const BaseMessage *)data;
  ota_task_queue_message_t msg = {.command = OTA_COMMAND_CHUNK};
  memcpy(msg.mac_addr, esp_now_info->src_addr, ESP_NOW_ETH_ALEN);
  memcpy(&msg.payload.ota_chunk_payload, &message->payload.ota_chunk_payload,
         sizeof(OTA_CHUNK_Payload));
  if (xQueueSend(ota_task_queue, &msg, pdMS_TO_TICKS(10)) != pdTRUE) {
    ESP_LOGE(TAG, "Failed to send chunk to OTA task");
  }
}

void slave_request_block_status_callback(
    const esp_now_recv_info_t *esp_now_info, const uint8_t *data,
    int data_len) {
  const BaseMessage *message = (const BaseMessage *)data;
  ota_task_queue_message_t msg = {.command = OTA_COMMAND_REQUEST_STATUS};
  memcpy(msg.mac_addr, esp_now_info->src_addr, ESP_NOW_ETH_ALEN);
  memcpy(&msg.payload.ota_request_block_status_payload,
         &message->payload.ota_request_block_status_payload,
         sizeof(OTA_REQUEST_BLOCK_STATUS_Payload));
  if (xQueueSend(ota_task_queue, &msg, pdMS_TO_TICKS(10)) != pdTRUE) {
    ESP_LOGE(TAG, "Failed to send block status request to OTA task");
  }
}

void slave_commit_block_callback(const esp_now_recv_info_t *esp_now_info,
                                 const uint8_t *data, int data_len) {
  const BaseMessage *message = (const BaseMessage *)data;
  ota_task_queue_message_t msg = {.command = OTA_COMMAND_COMMIT};
  memcpy(msg.mac_addr, esp_now_info->src_addr, ESP_NOW_ETH_ALEN);
  memcpy(&msg.payload.ota_commit_block_payload,
         &message->payload.ota_commit_block_payload,
         sizeof(OTA_COMMIT_BLOCK_Payload));
  if (xQueueSend(ota_task_queue, &msg, pdMS_TO_TICKS(10)) != pdTRUE) {
    ESP_LOGE(TAG, "Failed to send commit command to OTA task");
  }
}

void slave_finish_update_callback(const esp_now_recv_info_t *esp_now_info,
                                  const uint8_t *data, int data_len) {
  ota_task_queue_message_t msg = {.command = OTA_COMMAND_FINISH};
  memcpy(msg.mac_addr, esp_now_info->src_addr, ESP_NOW_ETH_ALEN);
  if (xQueueSend(ota_task_queue, &msg, pdMS_TO_TICKS(10)) != pdTRUE) {
    ESP_LOGE(TAG, "Failed to send finish command to OTA task");
  }
}

void slave_master_ready_to_send_chunks_callback(
    const esp_now_recv_info_t *esp_now_info, const uint8_t *data,
    int data_len) {
  ota_task_queue_message_t msg = {.command = OTA_COMMAND_MASTER_READY};
  memcpy(msg.mac_addr, esp_now_info->src_addr, ESP_NOW_ETH_ALEN);
  if (xQueueSend(ota_task_queue, &msg, pdMS_TO_TICKS(10)) != pdTRUE) {
    ESP_LOGE(TAG, "Failed to send master ready command to OTA task");
  }
}