#include "ota_update.h"
#include "driver/uart.h"
#include "esp_err.h"
#include "esp_log.h"
#include "esp_ota_ops.h"
#include "esp_partition.h"
#include "message_builder.h"
#include "message_handler.h"
#include "uart_handler.h"
#include "uart_msg_ids.h"
#include <stddef.h>
#include <stdint.h>
#include <string.h>

#define MAX(a, b) ((a) > (b) ? (a) : (b))
#define MIN(a, b) ((a) < (b) ? (a) : (b))

static uint8_t update_buffer[UPDATE_BUFFER_SIZE];
static uint16_t update_buffer_write_index;
static uint16_t sequenz_counter; // how often the update buffer gets written
static const char *TAG = "ALOX - OTA";
static esp_ota_handle_t update_handle;

void prepare_ota_update() {
  const esp_partition_t *running = esp_ota_get_running_partition();
  ESP_LOGI(TAG, "OTA: Running Partition: %s", running->label);

  char partition_to_update[] = "ota_0";
  if (strcmp(running->label, "ota_0") == 0) {
    strcpy(partition_to_update, "ota_1");
  }

  const esp_partition_t *update_partition = esp_partition_find_first(
      ESP_PARTITION_TYPE_APP, ESP_PARTITION_SUBTYPE_ANY, partition_to_update);

  // Check if the partition was found
  if (update_partition == NULL) {
    ESP_LOGE(TAG, "Failed to find OTA partition: %s", partition_to_update);
    return; // Or handle the error appropriately
  }

  ESP_LOGI(TAG, "Gonna write OTA Update in 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));
    esp_ota_abort(update_handle);
    return;
  }

  ESP_LOGI(TAG, "OTA update started successfully.");
  // Proceed with writing the new firmware to the partition...
}

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

  prepare_ota_update();

  int send_payload_len = 2;
  send_payload_buffer[0] = 0xff;
  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);
}

void payload_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 Payload Uart Command");

  if (update_buffer_write_index < UPDATE_BUFFER_SIZE - UPDATE_PAYLOAD_SIZE) {
    uint32_t write_len = MIN(UPDATE_PAYLOAD_SIZE, payload_len);
    ESP_LOGI(TAG, "Writing Data to Update BUffer Sequence %d, writing Data %d",
             sequenz_counter, write_len);
    memcpy(&update_buffer[update_buffer_write_index], payload, write_len);
    update_buffer_write_index += write_len;
  } else {
    ESP_LOGI(TAG, "Update Buffer full, writing it to OTA Update");

    // write to ota
    esp_err_t err =
        esp_ota_write(update_handle, update_buffer, update_buffer_write_index);
    if (err != ESP_OK) {
      ESP_LOGE(TAG, "GOT ESP ERROR WRITE OTA %d", err);
    }

    update_buffer_write_index = 0;
    sequenz_counter++;
  }

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

  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",
             payload_len, 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 = esp_ota_end(update_handle);
  if (err != ESP_OK) {
    ESP_LOGE(TAG, "GOT ESP ERROR WRITE OTA %d", err);
  }

  int send_payload_len = 0;
  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);
}

void write_ota_update_from_uart_task(void *param) {}

void init_ota() {
  RegisterCallback(UART_OTA_START, start_uart_update);
  RegisterCallback(UART_OTA_PAYLOAD, payload_uart_update);
  RegisterCallback(UART_OTA_END, end_uart_update);
}