// SDK
#include "cJSON.h"
#include "esp_http_server.h"
#include "esp_log.h"
#include "esp_partition.h"
#include "esp_system.h"
#include "esp_timer.h"
#include "nvs.h"
#include "nvs_flash.h"

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

#include <sys/param.h>

#define OTA_SCRATCH_BUFSIZE 4096

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

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

  uint8_t target_slot = g_Active_WWW_Partition == 0 ? 1 : 0;
  const char *target_label = target_slot == 0 ? "www_0" : "www_1";

  const esp_partition_t *partition = esp_partition_find_first(
      ESP_PARTITION_TYPE_DATA, ESP_PARTITION_SUBTYPE_DATA_LITTLEFS,
      target_label);
  if (!partition)
  {
    ESP_LOGE("OTA", "Could not find partition %s", target_label);
    httpd_resp_send_err(req, HTTPD_500_INTERNAL_SERVER_ERROR,
                        "Partition not found");
    return ESP_FAIL;
  }

  ESP_LOGI("OTA", "Starting OTA to partition %s (size %ld)", target_label,
           partition->size);

  esp_err_t err = esp_partition_erase_range(partition, 0, partition->size);
  if (err != ESP_OK)
  {
    ESP_LOGE("OTA", "Failed to erase partition: %s", esp_err_to_name(err));
    httpd_resp_send_err(req, HTTPD_500_INTERNAL_SERVER_ERROR,
                        "Failed to erase partition");
    return ESP_FAIL;
  }

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

  int total_read = 0;
  int remaining = req->content_len;
  bool first_chunk = true;

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

    if (first_chunk)
    {
      if ((uint8_t)buf[0] == 0xE9)
      {
        ESP_LOGE("OTA", "Magic 0xE9 detected. This looks like a FIRMWARE bin, "
                        "but you are uploading to FRONTEND slot!");
        free(buf);
        httpd_resp_send_err(
            req, HTTPD_400_BAD_REQUEST,
            "Invalid file: This is a Firmware binary, not a UI binary.");
        return ESP_FAIL;
      }
      first_chunk = false;
    }

    err = esp_partition_write(partition, total_read, buf, recv_len);
    if (err != ESP_OK)
    {
      ESP_LOGE("OTA", "Failed to write to partition: %s", esp_err_to_name(err));
      free(buf);
      httpd_resp_send_err(req, HTTPD_500_INTERNAL_SERVER_ERROR,
                          "Flash write failed");
      return ESP_FAIL;
    }

    total_read += recv_len;
    remaining -= recv_len;
  }

  free(buf);
  ESP_LOGI("OTA", "OTA complete. Written %d bytes. Updating NVS...",
           total_read);

  nvs_handle_t my_handle;
  if (nvs_open("storage", NVS_READWRITE, &my_handle) == ESP_OK)
  {
    err = nvs_set_u8(my_handle, "www_part", target_slot);
    if (err == ESP_OK)
    {
      nvs_commit(my_handle);
    }
    nvs_close(my_handle);
  }
  else
  {
    ESP_LOGE("OTA", "Failed to open NVS to update partition index");
    httpd_resp_send_err(req, HTTPD_500_INTERNAL_SERVER_ERROR,
                        "NVS update failed");
    return ESP_FAIL;
  }

  httpd_resp_set_type(req, "application/json");
  cJSON *root = cJSON_CreateObject();
  cJSON_AddStringToObject(root, "status", "success");
  cJSON_AddStringToObject(root, "message", "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
  const esp_timer_create_args_t restart_timer_args = {
      .callback = &ota_restart_timer_callback,
      .arg = (void *)0,
      .dispatch_method = ESP_TIMER_TASK,
      .name = "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_frontend_uri = {.uri = "/api/ota/frontend",
                                                   .method = HTTP_POST,
                                                   .handler =
                                                       api_ota_frontend_handler,
                                                   .user_ctx = NULL};