// SDK
#include "cJSON.h"
#include "esp_http_server.h"
#include "esp_log.h"
#include "esp_ota_ops.h"
#include "esp_system.h"
#include "esp_timer.h"
#include <sys/param.h>

// Project
#include "appstate.hpp"
#include "types.hpp"

#define OTA_FIRMWARE_SCRATCH_BUFSIZE 4096

internal void firmware_ota_restart_timer_callback(void *arg) { esp_restart(); }

internal esp_err_t api_ota_firmware_handler(httpd_req_t *req)
{
  httpd_resp_set_hdr(req, "Access-Control-Allow-Origin", "*");

  const esp_partition_t *update_partition =
      esp_ota_get_next_update_partition(NULL);
  if (update_partition == NULL)
  {
    ESP_LOGE("OTA_FW", "Passive OTA partition not found");
    httpd_resp_send_err(req, HTTPD_500_INTERNAL_SERVER_ERROR,
                        "OTA partition not found");
    return ESP_FAIL;
  }

  ESP_LOGI("OTA_FW", "Writing to partition subtype %d at offset 0x%lx",
           update_partition->subtype, update_partition->address);

  esp_ota_handle_t update_handle = 0;
  esp_err_t err =
      esp_ota_begin(update_partition, OTA_SIZE_UNKNOWN, &update_handle);
  if (err != ESP_OK)
  {
    ESP_LOGE("OTA_FW", "esp_ota_begin failed (%s)", esp_err_to_name(err));
    httpd_resp_send_err(req, HTTPD_500_INTERNAL_SERVER_ERROR,
                        "OTA begin failed");
    return ESP_FAIL;
  }

  char *buf = (char *)malloc(OTA_FIRMWARE_SCRATCH_BUFSIZE);
  if (!buf)
  {
    ESP_LOGE("OTA_FW", "Failed to allocate buffer");
    httpd_resp_send_err(req, HTTPD_500_INTERNAL_SERVER_ERROR, "Out of memory");
    return ESP_FAIL;
  }

  int binary_file_len = 0;
  int remaining = req->content_len;

  while (remaining > 0)
  {
    int recv_len =
        httpd_req_recv(req, buf, MIN(remaining, OTA_FIRMWARE_SCRATCH_BUFSIZE));
    if (recv_len <= 0)
    {
      if (recv_len == HTTPD_SOCK_ERR_TIMEOUT)
      {
        continue;
      }
      ESP_LOGE("OTA_FW", "Receive failed");
      esp_ota_abort(update_handle);
      free(buf);
      httpd_resp_send_err(req, HTTPD_500_INTERNAL_SERVER_ERROR,
                          "Receive failed");
      return ESP_FAIL;
    }

    err = esp_ota_write(update_handle, (const void *)buf, recv_len);
    if (err != ESP_OK)
    {
      ESP_LOGE("OTA_FW", "esp_ota_write failed (%s)", esp_err_to_name(err));
      esp_ota_abort(update_handle);
      free(buf);
      httpd_resp_send_err(req, HTTPD_500_INTERNAL_SERVER_ERROR,
                          "Flash write failed");
      return ESP_FAIL;
    }

    binary_file_len += recv_len;
    remaining -= recv_len;
  }

  free(buf);
  ESP_LOGI("OTA_FW", "Total binary data written: %d", binary_file_len);

  err = esp_ota_end(update_handle);
  if (err != ESP_OK)
  {
    if (err == ESP_ERR_OTA_VALIDATE_FAILED)
    {
      ESP_LOGE("OTA_FW", "Image validation failed, image is corrupted");
    }
    else
    {
      ESP_LOGE("OTA_FW", "esp_ota_end failed (%s)!", esp_err_to_name(err));
    }
    httpd_resp_send_err(req, HTTPD_500_INTERNAL_SERVER_ERROR,
                        "OTA validation/end failed");
    return ESP_FAIL;
  }

  err = esp_ota_set_boot_partition(update_partition);
  if (err != ESP_OK)
  {
    ESP_LOGE("OTA_FW", "esp_ota_set_boot_partition failed (%s)!",
             esp_err_to_name(err));
    httpd_resp_send_err(req, HTTPD_500_INTERNAL_SERVER_ERROR,
                        "Failed to set boot partition");
    return ESP_FAIL;
  }

  ESP_LOGI("OTA_FW", "OTA successful, rebooting...");

  httpd_resp_set_type(req, "application/json");
  cJSON *root = cJSON_CreateObject();
  cJSON_AddStringToObject(root, "status", "success");
  cJSON_AddStringToObject(root, "message",
                          "Firmware update successful, rebooting...");
  const char *response_text = cJSON_Print(root);
  httpd_resp_sendstr(req, response_text);
  free((void *)response_text);
  cJSON_Delete(root);

  // Trigger reboot with 1s delay
  const esp_timer_create_args_t restart_timer_args = {
      .callback = &firmware_ota_restart_timer_callback,
      .arg = (void *)0,
      .dispatch_method = ESP_TIMER_TASK,
      .name = "fw_ota_restart_timer",
      .skip_unhandled_events = false};
  esp_timer_handle_t restart_timer;
  esp_timer_create(&restart_timer_args, &restart_timer);
  esp_timer_start_once(restart_timer, 1'000'000);

  return ESP_OK;
}

internal const httpd_uri_t api_ota_firmware_uri = {.uri = "/api/ota/firmware",
                                                   .method = HTTP_POST,
                                                   .handler =
                                                       api_ota_firmware_handler,
                                                   .user_ctx = NULL};