#include "ota_uart.h"
#include "esp_log.h"
#include "esp_ota_ops.h"
#include <string.h>

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

typedef struct {
  bool active;
  esp_ota_handle_t handle;
  const esp_partition_t *update_partition;
  uint32_t total_size;
  uint32_t received;
  uint32_t written;
  uint32_t expected_seq;
  int target_slot;
  uint8_t block_buf[OTA_UART_FLASH_BLOCK_SIZE];
  size_t block_len;
} ota_uart_state_t;

static ota_uart_state_t s_ota;

static struct {
  bool valid;
  const esp_partition_t *partition;
  uint32_t size;
  int target_slot;
} s_staged;

static int partition_slot(const esp_partition_t *part) {
  if (part == NULL) {
    return -1;
  }
  if (part->subtype == ESP_PARTITION_SUBTYPE_APP_OTA_0) {
    return 0;
  }
  if (part->subtype == ESP_PARTITION_SUBTYPE_APP_OTA_1) {
    return 1;
  }
  return -1;
}

bool ota_uart_is_active(void) { return s_ota.active; }

int ota_uart_target_slot(void) {
  return s_ota.active ? s_ota.target_slot : -1;
}

void ota_uart_abort(void) {
  if (!s_ota.active) {
    return;
  }
  esp_ota_abort(s_ota.handle);
  memset(&s_ota, 0, sizeof(s_ota));
}

static esp_err_t flush_block(void) {
  if (s_ota.block_len == 0) {
    return ESP_OK;
  }
  esp_err_t err =
      esp_ota_write(s_ota.handle, s_ota.block_buf, s_ota.block_len);
  if (err != ESP_OK) {
    ESP_LOGE(TAG, "esp_ota_write %u bytes failed: %s", (unsigned)s_ota.block_len,
             esp_err_to_name(err));
    return err;
  }
  s_ota.written += (uint32_t)s_ota.block_len;
  s_ota.block_len = 0;
  return ESP_OK;
}

int ota_uart_prepare(uint32_t total_size) {
  if (s_ota.active) {
    ESP_LOGW(TAG, "OTA already active");
    return -1;
  }

  const esp_partition_t *running = esp_ota_get_running_partition();
  const esp_partition_t *update_partition =
      esp_ota_get_next_update_partition(NULL);
  if (update_partition == NULL) {
    ESP_LOGE(TAG, "no OTA update partition");
    return -1;
  }

  ESP_LOGI(TAG, "running=%s, update=%s, image_size=%lu",
           running != NULL ? running->label : "?",
           update_partition->label, (unsigned long)total_size);

  if (total_size > 0 && total_size > update_partition->size) {
    ESP_LOGE(TAG, "image too large (%lu > %lu)", (unsigned long)total_size,
             (unsigned long)update_partition->size);
    return -1;
  }

  esp_ota_handle_t handle;
  esp_err_t err = esp_ota_begin(update_partition, OTA_SIZE_UNKNOWN, &handle);
  if (err != ESP_OK) {
    ESP_LOGE(TAG, "esp_ota_begin failed: %s", esp_err_to_name(err));
    return -1;
  }

  memset(&s_ota, 0, sizeof(s_ota));
  s_ota.active = true;
  s_ota.handle = handle;
  s_ota.update_partition = update_partition;
  s_ota.total_size = total_size;
  s_ota.target_slot = partition_slot(update_partition);

  ESP_LOGI(TAG, "OTA prepared, target slot %d (%s) — send 4 KiB chunks",
           s_ota.target_slot, update_partition->label);
  return s_ota.target_slot;
}

bool ota_uart_block_ready_for_reack(void) {
  if (!s_ota.active) {
    return false;
  }
  return s_ota.written > 0 &&
         (s_ota.written % OTA_UART_FLASH_BLOCK_SIZE) == 0 &&
         s_ota.block_len == 0;
}

ota_feed_result_t ota_uart_feed_chunk(uint32_t seq, const uint8_t *data,
                                      size_t len) {
  if (!s_ota.active || data == NULL || len == 0) {
    return OTA_FEED_ERROR;
  }
  if (seq < s_ota.expected_seq) {
    return OTA_FEED_SEQ_DUP;
  }
  if (seq > s_ota.expected_seq) {
    ESP_LOGW(TAG, "seq gap: got %lu expected %lu", (unsigned long)seq,
             (unsigned long)s_ota.expected_seq);
    return OTA_FEED_SEQ_GAP;
  }
  s_ota.expected_seq++;

  if (len > OTA_UART_HOST_CHUNK_SIZE) {
    ESP_LOGW(TAG, "chunk %u > %u, truncating", (unsigned)len,
             OTA_UART_HOST_CHUNK_SIZE);
    len = OTA_UART_HOST_CHUNK_SIZE;
  }

  bool block_written = false;
  size_t offset = 0;
  while (offset < len) {
    size_t space = OTA_UART_FLASH_BLOCK_SIZE - s_ota.block_len;
    size_t n = len - offset;
    if (n > space) {
      n = space;
    }
    memcpy(s_ota.block_buf + s_ota.block_len, data + offset, n);
    s_ota.block_len += n;
    s_ota.received += (uint32_t)n;
    offset += n;

    if (s_ota.block_len < OTA_UART_FLASH_BLOCK_SIZE) {
      continue;
    }

    if (flush_block() != ESP_OK) {
      ota_uart_abort();
      return OTA_FEED_ERROR;
    }
    block_written = true;
  }

  return block_written ? OTA_FEED_BLOCK_WRITTEN : OTA_FEED_OK;
}

uint32_t ota_uart_bytes_written(void) { return s_ota.written; }

uint32_t ota_uart_bytes_received(void) { return s_ota.received; }

uint32_t ota_uart_total_size(void) { return s_ota.active ? s_ota.total_size : 0; }

bool ota_uart_get_staged_image(const esp_partition_t **partition_out,
                               uint32_t *size_out) {
  if (!s_staged.valid || s_staged.partition == NULL) {
    return false;
  }
  if (partition_out != NULL) {
    *partition_out = s_staged.partition;
  }
  if (size_out != NULL) {
    *size_out = s_staged.size;
  }
  return true;
}

void ota_uart_clear_staged(void) { memset(&s_staged, 0, sizeof(s_staged)); }

esp_err_t ota_uart_apply_boot(void) {
  if (!s_staged.valid || s_staged.partition == NULL) {
    return ESP_ERR_INVALID_STATE;
  }
  esp_err_t err = esp_ota_set_boot_partition(s_staged.partition);
  if (err != ESP_OK) {
    ESP_LOGE(TAG, "esp_ota_set_boot_partition failed: %s", esp_err_to_name(err));
    return err;
  }
  ESP_LOGI(TAG, "boot partition set to %s (slot %d)",
           s_staged.partition->label, s_staged.target_slot);
  ota_uart_clear_staged();
  return ESP_OK;
}

esp_err_t ota_uart_finish(bool set_boot, bool *success_out) {
  if (success_out != NULL) {
    *success_out = false;
  }
  if (!s_ota.active) {
    return ESP_ERR_INVALID_STATE;
  }

  esp_err_t err = flush_block();
  if (err != ESP_OK) {
    ota_uart_abort();
    return err;
  }

  if (s_ota.total_size > 0 && s_ota.received != s_ota.total_size) {
    ESP_LOGE(TAG, "size mismatch: received=%lu expected=%lu",
             (unsigned long)s_ota.received, (unsigned long)s_ota.total_size);
    ota_uart_abort();
    return ESP_ERR_INVALID_SIZE;
  }

  err = esp_ota_end(s_ota.handle);
  if (err != ESP_OK) {
    ESP_LOGE(TAG, "esp_ota_end failed: %s", esp_err_to_name(err));
    ota_uart_abort();
    return err;
  }

  ota_uart_clear_staged();
  s_staged.valid = true;
  s_staged.partition = s_ota.update_partition;
  s_staged.size = s_ota.written;
  s_staged.target_slot = s_ota.target_slot;

  ESP_LOGI(TAG, "OTA image staged: %lu bytes on %s (slot %d)",
           (unsigned long)s_staged.size, s_staged.partition->label,
           s_staged.target_slot);

  memset(&s_ota, 0, sizeof(s_ota));

  if (set_boot) {
    err = ota_uart_apply_boot();
    if (err != ESP_OK) {
      return err;
    }
  }

  if (success_out != NULL) {
    *success_out = true;
  }
  return ESP_OK;
}