#include "ota_handler.h"
#include "config.h"
#include "esp_log.h"
#include "esp_ota_ops.h"
#include "esp_http_client.h"
#include "esp_https_ota.h"
#include "freertos/FreeRTOS.h"
#include "freertos/task.h"
#include "freertos/semphr.h"
#include <string.h>

static const char *TAG = "OTA";

// OTA state
static bool ota_initialized = false;
static bool ota_updating = false;
static int ota_progress = 0;
static SemaphoreHandle_t ota_mutex = NULL;
static TaskHandle_t ota_task_handle = NULL;
static esp_ota_handle_t ota_upload_handle = 0;
static const esp_partition_t *ota_update_partition = NULL;
static size_t ota_bytes_written = 0;

// OTA configuration
#define OTA_TIMEOUT_MS 30000
#define OTA_BUFFER_SIZE 4096

// ============================================================================
// OTA Update Task
// ============================================================================

typedef struct {
	char url[256];
} ota_update_params_t;

static void ota_update_task(void *arg) {
	ota_update_params_t *params = (ota_update_params_t *)arg;

	if (!params) {
		ESP_LOGE(TAG, "Invalid OTA parameters");
		ota_updating = false;
		vTaskDelete(NULL);
		return;
	}

	ESP_LOGI(TAG, "Starting OTA update from: %s", params->url);

	esp_http_client_config_t http_config = {
	    .url = params->url,
	    .timeout_ms = OTA_TIMEOUT_MS,
	    .buffer_size = OTA_BUFFER_SIZE,
	    .skip_cert_common_name_check = true,
	};

	esp_https_ota_config_t ota_config = {
	    .http_config = &http_config,
	};

	esp_https_ota_handle_t ota_handle = NULL;
	esp_err_t ret = esp_https_ota_begin(&ota_config, &ota_handle);
	if (ret != ESP_OK) {
		ESP_LOGE(TAG, "OTA begin failed: %s", esp_err_to_name(ret));
		free(params);
		ota_updating = false;
		vTaskDelete(NULL);
		return;
	}

	int last_progress = -1;

	while (1) {
		ret = esp_https_ota_perform(ota_handle);
		if (ret == ESP_ERR_HTTPS_OTA_IN_PROGRESS) {
			// Update progress
			int downloaded = esp_https_ota_get_image_len_read(ota_handle);
			int total = esp_https_ota_get_image_size(ota_handle);

			if (total > 0) {
				ota_progress = (downloaded * 100) / total;
				if (ota_progress != last_progress) {
					ESP_LOGI(TAG, "OTA progress: %d%% (%d/%d bytes)", ota_progress, downloaded, total);
					last_progress = ota_progress;
				}
			}
		} else if (ret == ESP_OK) {
			// OTA complete
			ESP_LOGI(TAG, "OTA download complete");
			break;
		} else {
			ESP_LOGE(TAG, "OTA perform failed: %s", esp_err_to_name(ret));
			esp_https_ota_abort(ota_handle);
			free(params);
			ota_updating = false;
			ota_progress = 0;
			vTaskDelete(NULL);
			return;
		}
	}

	ret = esp_https_ota_finish(ota_handle);
	if (ret == ESP_OK) {
		ESP_LOGI(TAG, "OTA update successful, restarting...");
		free(params);
		ota_progress = 100;
		vTaskDelay(pdMS_TO_TICKS(1000));
		esp_restart();
	} else {
		ESP_LOGE(TAG, "OTA finish failed: %s", esp_err_to_name(ret));
		free(params);
		ota_updating = false;
		ota_progress = 0;
	}

	vTaskDelete(NULL);
}

// ============================================================================
// Public Functions
// ============================================================================

esp_err_t ota_handler_init(void) {
	if (ota_initialized) {
		ESP_LOGW(TAG, "OTA handler already initialized");
		return ESP_OK;
	}

	ESP_LOGI(TAG, "Initializing OTA handler");

	ota_mutex = xSemaphoreCreateMutex();
	if (!ota_mutex) {
		ESP_LOGE(TAG, "Failed to create OTA mutex");
		return ESP_ERR_NO_MEM;
	}

	ota_initialized = true;
	ota_updating = false;
	ota_progress = 0;

	// Print current partition info
	const esp_partition_t *running = esp_ota_get_running_partition();
	if (running) {
		ESP_LOGI(TAG, "Running from partition: %s (offset: 0x%lx)", running->label, running->address);
	}

	return ESP_OK;
}

esp_err_t ota_handler_deinit(void) {
	if (!ota_initialized) {
		return ESP_OK;
	}

	ESP_LOGI(TAG, "Deinitializing OTA handler");

	if (ota_mutex) {
		vSemaphoreDelete(ota_mutex);
		ota_mutex = NULL;
	}

	ota_initialized = false;
	return ESP_OK;
}

const char *ota_get_running_partition(void) {
	const esp_partition_t *running = esp_ota_get_running_partition();
	if (running) {
		return running->label;
	}
	return "unknown";
}

bool ota_is_updating(void) {
	return ota_updating;
}

int ota_get_progress(void) {
	if (ota_updating) {
		return ota_progress;
	}
	return -1;
}

// ============================================================================
// OTA Start Function (called from web server)
// ============================================================================

esp_err_t ota_start_update(const char *url) {
	if (!ota_initialized) {
		ESP_LOGE(TAG, "OTA handler not initialized");
		return ESP_ERR_INVALID_STATE;
	}

	if (ota_updating) {
		ESP_LOGW(TAG, "OTA update already in progress");
		return ESP_ERR_INVALID_STATE;
	}

	if (!url || strlen(url) == 0) {
		ESP_LOGE(TAG, "Invalid URL");
		return ESP_ERR_INVALID_ARG;
	}

	// Allocate parameters
	ota_update_params_t *params = calloc(1, sizeof(ota_update_params_t));
	if (!params) {
		ESP_LOGE(TAG, "Failed to allocate OTA params");
		return ESP_ERR_NO_MEM;
	}

	strncpy(params->url, url, sizeof(params->url) - 1);
	params->url[sizeof(params->url) - 1] = '\0';

	ota_updating = true;
	ota_progress = 0;

	// Create OTA task
	BaseType_t ret = xTaskCreate(ota_update_task, "ota_update", 8192, params, 5, &ota_task_handle);

	if (ret != pdPASS) {
		ESP_LOGE(TAG, "Failed to create OTA task");
		free(params);
		ota_updating = false;
		ota_progress = 0;
		return ESP_ERR_NO_MEM;
	}

	return ESP_OK;
}

// ============================================================================
// OTA Upload Functions (for direct file upload)
// ============================================================================

esp_err_t ota_start_upload(void) {
	if (!ota_initialized) {
		ESP_LOGE(TAG, "OTA handler not initialized");
		return ESP_ERR_INVALID_STATE;
	}

	if (ota_updating) {
		ESP_LOGW(TAG, "OTA update already in progress");
		return ESP_ERR_INVALID_STATE;
	}

	// Get next update partition
	ota_update_partition = esp_ota_get_next_update_partition(NULL);
	if (!ota_update_partition) {
		ESP_LOGE(TAG, "No OTA partition found");
		return ESP_ERR_NOT_FOUND;
	}

	ESP_LOGI(TAG, "Starting OTA upload to partition: %s", ota_update_partition->label);

	// Begin OTA
	esp_err_t ret = esp_ota_begin(ota_update_partition, OTA_SIZE_UNKNOWN, &ota_upload_handle);
	if (ret != ESP_OK) {
		ESP_LOGE(TAG, "esp_ota_begin failed: %s", esp_err_to_name(ret));
		ota_update_partition = NULL;
		return ret;
	}

	ota_updating = true;
	ota_progress = 0;
	ota_bytes_written = 0;

	return ESP_OK;
}

esp_err_t ota_write_data(const uint8_t *data, size_t len) {
	if (!ota_updating || !ota_update_partition) {
		ESP_LOGE(TAG, "OTA upload not started");
		return ESP_ERR_INVALID_STATE;
	}

	esp_err_t ret = esp_ota_write(ota_upload_handle, data, len);
	if (ret != ESP_OK) {
		ESP_LOGE(TAG, "esp_ota_write failed: %s", esp_err_to_name(ret));
		return ret;
	}

	ota_bytes_written += len;

	// Update progress (we don't know total size, so show bytes written)
	ESP_LOGD(TAG, "Written %zu bytes", ota_bytes_written);

	return ESP_OK;
}

esp_err_t ota_finish_upload(void) {
	if (!ota_updating || !ota_update_partition) {
		ESP_LOGE(TAG, "OTA upload not started");
		return ESP_ERR_INVALID_STATE;
	}

	ESP_LOGI(TAG, "Finishing OTA upload, total bytes: %zu", ota_bytes_written);

	esp_err_t ret = esp_ota_end(ota_upload_handle);
	if (ret != ESP_OK) {
		ESP_LOGE(TAG, "esp_ota_end failed: %s", esp_err_to_name(ret));
		ota_updating = false;
		ota_upload_handle = 0;
		ota_update_partition = NULL;
		ota_bytes_written = 0;
		return ret;
	}

	// Set boot partition
	ret = esp_ota_set_boot_partition(ota_update_partition);
	if (ret != ESP_OK) {
		ESP_LOGE(TAG, "esp_ota_set_boot_partition failed: %s", esp_err_to_name(ret));
		ota_updating = false;
		ota_upload_handle = 0;
		ota_update_partition = NULL;
		ota_bytes_written = 0;
		return ret;
	}

	ESP_LOGI(TAG, "OTA upload complete, restarting...");
	ota_progress = 100;
	ota_updating = false;
	ota_upload_handle = 0;
	ota_update_partition = NULL;
	ota_bytes_written = 0;

	// Restart after short delay
	vTaskDelay(pdMS_TO_TICKS(1000));
	esp_restart();

	return ESP_OK;
}

esp_err_t ota_abort_upload(void) {
	if (!ota_updating) {
		return ESP_OK;
	}

	ESP_LOGI(TAG, "Aborting OTA upload");

	if (ota_upload_handle) {
		esp_ota_abort(ota_upload_handle);
		ota_upload_handle = 0;
	}

	ota_updating = false;
	ota_progress = 0;
	ota_update_partition = NULL;
	ota_bytes_written = 0;

	return ESP_OK;
}