#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 "esp_system.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 uint32_t update_size;
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;

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

  char partition_to_update[] = "ota_0";
  if (strcmp(running->label, "ota_0") == 0) {
    strcpy(partition_to_update, "ota_1");
    part = 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 -1; // 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 -2;
  }

  ESP_LOGI(TAG, "OTA update started successfully.");
  return part;
}

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;

  int part = prepare_ota_update();

  // Message:
  // byte partition
  // byte error

  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(uint32_t write_len, uint8_t *payload) {
  if (update_buffer_write_index > UPDATE_BUFFER_SIZE - write_len) {
    // ESP_LOGI(TAG, "Writing Data to Update BUffer Sequence %d, writing Data
    // %d",
    //          sequenz_counter, write_len);
    // write to ota
    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++;
    return err;
  }

  memcpy(&update_buffer[update_buffer_write_index], payload, write_len);
  update_buffer_write_index += write_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");

  uint32_t write_len = MIN(UPDATE_PAYLOAD_SIZE, payload_len);
  update_size += write_len;

  esp_err_t err = write_ota(write_len, payload);

  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",
             payload_len, send_buffer_size, len);
    return;
  }

  uart_write_bytes(MASTER_UART, send_buffer, len);
}

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

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

  ESP_LOGE(TAG, "UPDATE ENDE UPDATGE SIZE SIND %d BYTES", update_size);

  // Hol dir die zuletzt geschriebene Partition
  const esp_partition_t *partition = esp_ota_get_next_update_partition(NULL);
  if (partition == NULL) {
    ESP_LOGE(TAG, "Failed to get updated partition");
    err = ESP_FAIL;
  }

  // Setze sie als Boot-Partition
  ESP_LOGE(TAG, "Setzte nächste Partition auf %s", partition->label);
  err = esp_ota_set_boot_partition(partition);
  if (err != ESP_OK) {
    ESP_LOGE(TAG, "esp_ota_set_boot_partition failed: %s",
             esp_err_to_name(err));
  }
  return err;
}

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

  // message ret esp_err_t
  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 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);
}