/**
 * OTA Manager - Firmware Update Management
 * Handles Over-The-Air firmware updates via Wi-Fi
 * 
 * This implementation provides secure OTA updates for the ESP32-P4-NANO
 * network auditing tool, supporting both HTTP and HTTPS updates.
 */

#include "ota_manager.h"
#include "log_manager.h"
#include "esp_log.h"
#include "esp_err.h"
#include "esp_ota_ops.h"
#include "esp_https_ota.h"
#include "esp_http_client.h"
#include "esp_system.h"
#include "esp_app_desc.h"
#include <string.h>
#include <stdio.h>

static const char *TAG = "OTA_MANAGER";

// OTA progress structure (using types from header)
static ota_progress_t s_ota_progress = {0};
static bool s_ota_manager_initialized = false;
static SemaphoreHandle_t s_ota_mutex = NULL;

// Forward declaration
static void ota_manager_set_error(const char* error_message);

// OTA URLs (configured via web interface)
static char s_ota_server_url[256] = "http://192.168.4.1:8080/firmware/";
static char s_ota_version_url[256] = "http://192.168.4.1:8080/version.json";

esp_err_t ota_manager_init(void)
{
    ESP_LOGI(TAG, "Initializing OTA manager");
    
    // Create mutex for thread safety
    s_ota_mutex = xSemaphoreCreateMutex();
    if (s_ota_mutex == NULL) {
        ESP_LOGE(TAG, "Failed to create OTA mutex");
        return ESP_ERR_NO_MEM;
    }
    
    // Initialize progress structure
    memset(&s_ota_progress, 0, sizeof(ota_progress_t));
    s_ota_progress.state = OTA_STATE_IDLE;
    
    // Get current version
    const esp_app_desc_t* app_desc = esp_app_get_description();
    if (app_desc) {
        strncpy(s_ota_progress.current_version, app_desc->version, sizeof(s_ota_progress.current_version) - 1);
    }
    
    s_ota_manager_initialized = true;
    ESP_LOGI(TAG, "OTA manager initialized successfully");
    ESP_LOGI(TAG, "Current firmware version: %s", s_ota_progress.current_version);
    
    return ESP_OK;
}

esp_err_t ota_manager_deinit(void)
{
    if (s_ota_mutex != NULL) {
        vSemaphoreDelete(s_ota_mutex);
        s_ota_mutex = NULL;
    }
    
    s_ota_manager_initialized = false;
    ESP_LOGI(TAG, "OTA manager deinitialized");
    return ESP_OK;
}

esp_err_t ota_manager_set_server_url(const char* server_url)
{
    if (!s_ota_manager_initialized) {
        return ESP_ERR_INVALID_STATE;
    }
    
    if (server_url == NULL) {
        return ESP_ERR_INVALID_ARG;
    }
    
    if (xSemaphoreTake(s_ota_mutex, pdMS_TO_TICKS(1000)) != pdTRUE) {
        return ESP_ERR_TIMEOUT;
    }
    
    strncpy(s_ota_server_url, server_url, sizeof(s_ota_server_url) - 1);
    s_ota_server_url[sizeof(s_ota_server_url) - 1] = '\0';
    
    // Update version URL
    snprintf(s_ota_version_url, sizeof(s_ota_version_url), "%sversion.json", server_url);
    
    xSemaphoreGive(s_ota_mutex);
    
    log_manager_log_info(TAG, "OTA server URL updated to: %s", server_url);
    return ESP_OK;
}

esp_err_t ota_manager_check_for_updates(void)
{
    if (!s_ota_manager_initialized || s_ota_mutex == NULL) {
        return ESP_ERR_INVALID_STATE;
    }
    
    if (xSemaphoreTake(s_ota_mutex, pdMS_TO_TICKS(5000)) != pdTRUE) {
        return ESP_ERR_TIMEOUT;
    }
    
    // Update state
    s_ota_progress.state = OTA_STATE_CHECKING;
    s_ota_progress.progress_percent = 0;
    strcpy(s_ota_progress.status_message, "Checking for updates...");
    
    xSemaphoreGive(s_ota_mutex);
    
    // Create HTTP client for version check
    esp_http_client_config_t config = {
        .url = s_ota_version_url,
        .timeout_ms = 10000,
    };
    
    esp_http_client_handle_t client = esp_http_client_init(&config);
    if (client == NULL) {
        log_manager_log_error(TAG, "Failed to create HTTP client for version check");
        ota_manager_set_error("Failed to create HTTP client");
        return ESP_ERR_NO_MEM;
    }
    
    esp_err_t ret = esp_http_client_perform(client);
    if (ret != ESP_OK) {
        log_manager_log_error(TAG, "HTTP request failed: %s", esp_err_to_name(ret));
        ota_manager_set_error("Failed to connect to update server");
        esp_http_client_cleanup(client);
        return ret;
    }
    
    int status_code = esp_http_client_get_status_code(client);
    if (status_code != 200) {
        log_manager_log_error(TAG, "HTTP request failed with status: %d", status_code);
        ota_manager_set_error("Server returned error status");
        esp_http_client_cleanup(client);
        return ESP_ERR_INVALID_RESPONSE;
    }
    
    int content_length = esp_http_client_get_content_length(client);
    if (content_length <= 0 || content_length > 1024) {
        log_manager_log_error(TAG, "Invalid content length: %d", content_length);
        ota_manager_set_error("Invalid server response");
        esp_http_client_cleanup(client);
        return ESP_ERR_INVALID_SIZE;
    }
    
    // Read version information
    char* buffer = malloc(content_length + 1);
    if (buffer == NULL) {
        log_manager_log_error(TAG, "Failed to allocate memory for version check");
        ota_manager_set_error("Memory allocation failed");
        esp_http_client_cleanup(client);
        return ESP_ERR_NO_MEM;
    }
    
    int bytes_read = esp_http_client_read(client, buffer, content_length);
    buffer[bytes_read] = '\0';
    
    esp_http_client_cleanup(client);
    
    // Parse version JSON (simplified - in production, use proper JSON parsing)
    char* version_start = strstr(buffer, "\"version\":");
    if (version_start) {
        version_start += 11; // Skip "version":
        char* version_end = strchr(version_start, '"');
        if (version_end) {
            int version_len = version_end - version_start;
            if (version_len < sizeof(s_ota_progress.new_version)) {
                strncpy(s_ota_progress.new_version, version_start, version_len);
                s_ota_progress.new_version[version_len] = '\0';
                
                // Check if new version is available
                if (strcmp(s_ota_progress.new_version, s_ota_progress.current_version) != 0) {
                    s_ota_progress.update_available = true;
                    log_manager_log_info(TAG, "New firmware version available: %s", s_ota_progress.new_version);
                    
                    xSemaphoreTake(s_ota_mutex, pdMS_TO_TICKS(1000));
                    s_ota_progress.state = OTA_STATE_IDLE;
                    s_ota_progress.progress_percent = 100;
                    strcpy(s_ota_progress.status_message, "Update available");
                    xSemaphoreGive(s_ota_mutex);
                } else {
                    s_ota_progress.update_available = false;
                    log_manager_log_info(TAG, "Firmware is up to date");
                    
                    xSemaphoreTake(s_ota_mutex, pdMS_TO_TICKS(1000));
                    s_ota_progress.state = OTA_STATE_IDLE;
                    s_ota_progress.progress_percent = 100;
                    strcpy(s_ota_progress.status_message, "Firmware is up to date");
                    xSemaphoreGive(s_ota_mutex);
                }
            }
        }
    }
    
    free(buffer);
    return ESP_OK;
}

esp_err_t ota_manager_start_update(const char* url)
{
    if (!s_ota_manager_initialized || s_ota_mutex == NULL) {
        return ESP_ERR_INVALID_STATE;
    }
    
    if (!s_ota_progress.update_available) {
        return ESP_ERR_INVALID_STATE;
    }
    
    if (xSemaphoreTake(s_ota_mutex, pdMS_TO_TICKS(5000)) != pdTRUE) {
        return ESP_ERR_TIMEOUT;
    }
    
    // Update state
    s_ota_progress.state = OTA_STATE_DOWNLOADING;
    s_ota_progress.progress_percent = 0;
    s_ota_progress.bytes_downloaded = 0;
    s_ota_progress.total_bytes = 0;
    strcpy(s_ota_progress.status_message, "Starting firmware download...");
    
    xSemaphoreGive(s_ota_mutex);
    
    // Use provided URL or create default URL
    char ota_url[512];
    if (url != NULL && strlen(url) > 0) {
        strncpy(ota_url, url, sizeof(ota_url) - 1);
        ota_url[sizeof(ota_url) - 1] = '\0';
    } else {
        snprintf(ota_url, sizeof(ota_url), "%sfirmware_%s.bin", s_ota_server_url, s_ota_progress.new_version);
    }
    
    log_manager_log_info(TAG, "Starting OTA update from: %s", ota_url);
    
    // Configure OTA
    esp_http_client_config_t config = {
        .url = ota_url,
        .timeout_ms = 30000,
    };
    
    esp_https_ota_config_t ota_config = {
        .http_config = &config,
        .http_client_init_cb = NULL,
        .bulk_flash_erase = true,
        .partial_http_download = true,
    };
    
    esp_err_t ret = esp_https_ota(&ota_config);
    if (ret != ESP_OK) {
        log_manager_log_error(TAG, "OTA HTTPS failed: %s", esp_err_to_name(ret));
        ota_manager_set_error("Failed to start OTA update");
        return ret;
    }
    
    // Update progress
    if (xSemaphoreTake(s_ota_mutex, pdMS_TO_TICKS(1000)) == pdTRUE) {
        s_ota_progress.state = OTA_STATE_READY_TO_UPDATE;
        s_ota_progress.progress_percent = 100;
        strcpy(s_ota_progress.status_message, "Update ready to install");
        xSemaphoreGive(s_ota_mutex);
    }
    
    log_manager_log_info(TAG, "OTA update completed successfully");
    return ESP_OK;
}

esp_err_t ota_manager_install_update(void)
{
    if (!s_ota_manager_initialized) {
        return ESP_ERR_INVALID_STATE;
    }
    
    if (xSemaphoreTake(s_ota_mutex, pdMS_TO_TICKS(1000)) != pdTRUE) {
        return ESP_ERR_TIMEOUT;
    }
    
    if (s_ota_progress.state != OTA_STATE_READY_TO_UPDATE) {
        xSemaphoreGive(s_ota_mutex);
        return ESP_ERR_INVALID_STATE;
    }
    
    s_ota_progress.state = OTA_STATE_UPDATING;
    strcpy(s_ota_progress.status_message, "Installing firmware update...");
    
    xSemaphoreGive(s_ota_mutex);
    
    log_manager_log_info(TAG, "Installing firmware update");
    
    // Set boot partition to OTA partition
    const esp_partition_t* running = esp_ota_get_running_partition();
    esp_ota_img_states_t ota_state;
    if (esp_ota_get_state_partition(running, &ota_state) == ESP_OK) {
        if (ota_state == ESP_OTA_IMG_PENDING_VERIFY) {
            esp_ota_mark_app_valid_cancel_rollback();
        }
    }
    
    // Restart to apply update
    esp_restart();
    
    return ESP_OK;
}

esp_err_t ota_manager_get_progress(ota_progress_t* progress)
{
    if (!s_ota_manager_initialized) {
        return ESP_ERR_INVALID_STATE;
    }
    
    if (progress == NULL) {
        return ESP_ERR_INVALID_ARG;
    }
    
    if (xSemaphoreTake(s_ota_mutex, pdMS_TO_TICKS(1000)) != pdTRUE) {
        return ESP_ERR_TIMEOUT;
    }
    
    memcpy(progress, &s_ota_progress, sizeof(ota_progress_t));
    
    xSemaphoreGive(s_ota_mutex);
    return ESP_OK;
}

esp_err_t ota_manager_cancel_update(void)
{
    if (!s_ota_manager_initialized) {
        return ESP_ERR_INVALID_STATE;
    }
    
    if (xSemaphoreTake(s_ota_mutex, pdMS_TO_TICKS(1000)) != pdTRUE) {
        return ESP_ERR_TIMEOUT;
    }
    
    s_ota_progress.state = OTA_STATE_IDLE;
    s_ota_progress.progress_percent = 0;
    strcpy(s_ota_progress.status_message, "Update cancelled");
    
    xSemaphoreGive(s_ota_mutex);
    
    log_manager_log_info(TAG, "OTA update cancelled by user");
    return ESP_OK;
}

esp_err_t ota_manager_get_current_version(char* version)
{
    if (!s_ota_manager_initialized) {
        return ESP_ERR_INVALID_STATE;
    }
    
    if (version == NULL) {
        return ESP_ERR_INVALID_ARG;
    }
    
    // Get current firmware version from app description
    const esp_app_desc_t* app_desc = esp_app_get_description();
    if (app_desc != NULL) {
        strncpy(version, app_desc->version, 31);
        version[31] = '\0';
    } else {
        strcpy(version, "unknown");
    }
    
    return ESP_OK;
}

static void ota_manager_set_error(const char* error_message)
{
    if (xSemaphoreTake(s_ota_mutex, pdMS_TO_TICKS(1000)) == pdTRUE) {
        s_ota_progress.state = OTA_STATE_ERROR;
        strncpy(s_ota_progress.error_message, error_message, sizeof(s_ota_progress.error_message) - 1);
        s_ota_progress.error_message[sizeof(s_ota_progress.error_message) - 1] = '\0';
        xSemaphoreGive(s_ota_mutex);
    }
} 