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

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

#define BUNDLE_SCRATCH_BUFSIZE 4096

typedef struct
{
  char magic[4];
  uint32_t fw_size;
  uint32_t www_size;
} bundle_header_t;

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

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

  if (req->content_len < sizeof(bundle_header_t))
  {
    ESP_LOGE("OTA_BUNDLE", "Request content too short");
    httpd_resp_send_err(req, HTTPD_400_BAD_REQUEST, "Content too short");
    return ESP_FAIL;
  }

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

  // 1. Read Header
  bundle_header_t header;
  int overhead = httpd_req_recv(req, (char *)&header, sizeof(bundle_header_t));
  if (overhead <= 0)
  {
    free(buf);
    httpd_resp_send_err(req, HTTPD_500_INTERNAL_SERVER_ERROR,
                        "Header receive failed");
    return ESP_FAIL;
  }

  if (memcmp(header.magic, "BNDL", 4) != 0)
  {
    free(buf);
    ESP_LOGE("OTA_BUNDLE", "Invalid magic: %.4s", header.magic);
    httpd_resp_send_err(req, HTTPD_400_BAD_REQUEST, "Invalid bundle magic");
    return ESP_FAIL;
  }

  ESP_LOGI("OTA_BUNDLE",
           "Starting Universal Update: FW %lu bytes, WWW %lu bytes",
           header.fw_size, header.www_size);

  // 2. Prepare Firmware Update
  const esp_partition_t *fw_part = esp_ota_get_next_update_partition(NULL);
  esp_ota_handle_t fw_handle = 0;
  esp_err_t err = esp_ota_begin(fw_part, header.fw_size, &fw_handle);
  if (err != ESP_OK)
  {
    free(buf);
    ESP_LOGE("OTA_BUNDLE", "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;
  }

  // 3. Stream Firmware
  uint32_t fw_remaining = header.fw_size;
  bool fw_first_chunk = true;
  while (fw_remaining > 0)
  {
    int recv_len =
        httpd_req_recv(req, buf, MIN(fw_remaining, BUNDLE_SCRATCH_BUFSIZE));
    if (recv_len <= 0)
    {
      if (recv_len == HTTPD_SOCK_ERR_TIMEOUT)
        continue;
      esp_ota_abort(fw_handle);
      free(buf);
      return ESP_FAIL;
    }

    if (fw_first_chunk && recv_len > 0)
    {
      if ((uint8_t)buf[0] != 0xE9)
      {
        ESP_LOGE("OTA_BUNDLE", "Invalid FW magic in bundle: %02X",
                 (uint8_t)buf[0]);
        esp_ota_abort(fw_handle);
        free(buf);
        httpd_resp_send_err(
            req, HTTPD_400_BAD_REQUEST,
            "Invalid Bundle: Firmware part is corrupted or invalid.");
        return ESP_FAIL;
      }
      fw_first_chunk = false;
    }

    esp_ota_write(fw_handle, buf, recv_len);
    fw_remaining -= recv_len;
  }
  esp_ota_end(fw_handle);

  // 4. Prepare WWW Update
  uint8_t target_www_slot = g_Active_WWW_Partition == 0 ? 1 : 0;
  const char *www_label = target_www_slot == 0 ? "www_0" : "www_1";
  const esp_partition_t *www_part = esp_partition_find_first(
      ESP_PARTITION_TYPE_DATA, ESP_PARTITION_SUBTYPE_DATA_LITTLEFS, www_label);

  esp_partition_erase_range(www_part, 0, www_part->size);

  // 5. Stream WWW
  uint32_t www_remaining = header.www_size;
  uint32_t www_written = 0;
  while (www_remaining > 0)
  {
    int recv_len =
        httpd_req_recv(req, buf, MIN(www_remaining, BUNDLE_SCRATCH_BUFSIZE));
    if (recv_len <= 0)
    {
      if (recv_len == HTTPD_SOCK_ERR_TIMEOUT)
        continue;
      free(buf);
      return ESP_FAIL;
    }
    esp_partition_write(www_part, www_written, buf, recv_len);
    www_written += recv_len;
    www_remaining -= recv_len;
  }

  free(buf);

  // 6. Commit Updates
  esp_ota_set_boot_partition(fw_part);

  nvs_handle_t nvs_h;
  if (nvs_open("storage", NVS_READWRITE, &nvs_h) == ESP_OK)
  {
    nvs_set_u8(nvs_h, "www_part", target_www_slot);
    nvs_commit(nvs_h);
    nvs_close(nvs_h);
  }

  ESP_LOGI("OTA_BUNDLE", "Universal Update Complete! Rebooting...");

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

  // Reboot
  esp_timer_create_args_t tmr_args = {};
  tmr_args.callback = &bundle_ota_restart_timer_callback;
  tmr_args.name = "bundle_reboot";
  esp_timer_handle_t tmr;
  esp_timer_create(&tmr_args, &tmr);
  esp_timer_start_once(tmr, 1'000'000);

  return ESP_OK;
}

internal const httpd_uri_t api_ota_bundle_uri = {.uri = "/api/ota/bundle",
                                                 .method = HTTP_POST,
                                                 .handler =
                                                     api_ota_bundle_handler,
                                                 .user_ctx = NULL};