#include "state_manager.h"
#include "config.h"
#include "motor_control.h"
#include "esp_log.h"
#include "esp_system.h"
#include "nvs_flash.h"
#include "nvs.h"
#include <string.h>
#include <stdio.h>

// Private functions
static esp_err_t validate_and_clamp_values(uint8_t* mode, uint8_t* speed, uint8_t* last_mode, uint8_t* last_speed);
static void log_reset_analysis(void);
static esp_err_t should_restore_state(bool* should_restore, motor_mode_t stored_mode, int stored_speed);

// Validate and clamp loaded values to safe ranges
static esp_err_t validate_and_clamp_values(uint8_t* mode, uint8_t* speed, uint8_t* last_mode, uint8_t* last_speed) {
    // Validate and clamp motor mode
    if (*mode > MOTOR_INTAKE) {
        ESP_LOGW(SYSTEM_TAG, "Invalid stored mode %d, clamping to OFF", *mode);
        *mode = MOTOR_OFF;
    }
    
    // Validate and clamp speed
    if (!IS_VALID_SPEED(*speed)) {
        ESP_LOGW(SYSTEM_TAG, "Invalid stored speed %d, clamping to 0", *speed);
        *speed = 0;
    }
    
    // Validate and clamp last ON mode
    if (*last_mode < MOTOR_EXHAUST || *last_mode > MOTOR_INTAKE) {
        ESP_LOGW(SYSTEM_TAG, "Invalid stored last mode %d, defaulting to EXHAUST", *last_mode);
        *last_mode = MOTOR_EXHAUST;
    }
    
    // Validate and clamp last ON speed
    if (!IS_VALID_SPEED(*last_speed) || *last_speed == 0) {
        ESP_LOGW(SYSTEM_TAG, "Invalid stored last speed %d, defaulting to 50", *last_speed);
        *last_speed = 50;
    }
    
    return ESP_OK;
}

// Log detailed reset analysis
static void log_reset_analysis(void) {
    esp_reset_reason_t reset_reason = esp_reset_reason();
    bool was_watchdog = state_manager_was_watchdog_reset();
    
    ESP_LOGI(SYSTEM_TAG, "=== RESET ANALYSIS ===");
    ESP_LOGI(SYSTEM_TAG, "Reset reason code: %d", reset_reason);
    ESP_LOGI(SYSTEM_TAG, "Reset reason: %s", state_manager_get_reset_reason_string());
    ESP_LOGI(SYSTEM_TAG, "Watchdog reset: %s", was_watchdog ? "YES" : "NO");
    ESP_LOGI(SYSTEM_TAG, "====================");
}

// Determine if state should be restored based on reset reason and user preferences
static esp_err_t should_restore_state(bool* should_restore, motor_mode_t stored_mode, int stored_speed) {
    if (!should_restore) return ESP_ERR_INVALID_ARG;
    
    *should_restore = false;
    
    bool was_watchdog = state_manager_was_watchdog_reset();
    bool user_turned_off = motor_get_user_turned_off();
    
    if (was_watchdog) {
        ESP_LOGI(SYSTEM_TAG, "⚠️  TRUE watchdog reset detected - starting in OFF state for safety");
        return ESP_OK;
    }
    
    if (user_turned_off) {
        ESP_LOGI(SYSTEM_TAG, "🔒 User had turned off manually - staying OFF");
        return ESP_OK;
    }
    
    if (stored_mode != MOTOR_OFF && stored_speed > 0) {
        ESP_LOGI(SYSTEM_TAG, "🔋 Power restored - will resume previous state: %s @ %d%%", 
                 motor_mode_to_string(stored_mode), stored_speed);
        *should_restore = true;
        return ESP_OK;
    }
    
    ESP_LOGI(SYSTEM_TAG, "❌ No valid state to restore (mode=%s, speed=%d)", 
             motor_mode_to_string(stored_mode), stored_speed);
    return ESP_OK;
}

// Public API Implementation

esp_err_t state_manager_init(void) {
    ESP_LOGI(SYSTEM_TAG, "Initializing state manager...");
    
    esp_err_t ret = nvs_flash_init();
    if (ret == ESP_ERR_NVS_NO_FREE_PAGES || ret == ESP_ERR_NVS_NEW_VERSION_FOUND) {
        ESP_LOGW(SYSTEM_TAG, "NVS partition was truncated, erasing...");
        ESP_ERROR_CHECK(nvs_flash_erase());
        ret = nvs_flash_init();
    }
    
    if (ret != ESP_OK) {
        ESP_LOGE(SYSTEM_TAG, "Failed to initialize NVS: %s", esp_err_to_name(ret));
        return ret;
    }
    
    ESP_LOGI(SYSTEM_TAG, "State manager initialized successfully");
    return ESP_OK;
}

esp_err_t state_manager_save(void) {
    nvs_handle_t nvs_handle;
    esp_err_t err;
    
    err = nvs_open(NVS_NAMESPACE, NVS_READWRITE, &nvs_handle);
    if (err != ESP_OK) {
        ESP_LOGE(SYSTEM_TAG, "Error opening NVS handle: %s", esp_err_to_name(err));
        return err;
    }
    
    // Get current motor state
    const motor_state_t* state = motor_get_state();
    motor_mode_t last_on_mode;
    int last_on_speed;
    motor_get_last_on_state(&last_on_mode, &last_on_speed);
    bool user_turned_off = motor_get_user_turned_off();
    
    ESP_LOGI(SYSTEM_TAG, "=== SAVING STATE TO NVS ===");
    ESP_LOGI(SYSTEM_TAG, "Mode: %s, Speed: %d%%, Last ON: %s@%d%%, User OFF: %s",
             motor_mode_to_string(state->mode), state->target_speed, 
             motor_mode_to_string(last_on_mode), last_on_speed,
             user_turned_off ? "YES" : "NO");
    
    // Save current motor state
    err = nvs_set_u8(nvs_handle, NVS_KEY_MODE, (uint8_t)state->mode);
    if (err == ESP_OK) {
        err = nvs_set_u8(nvs_handle, NVS_KEY_SPEED, (uint8_t)state->target_speed);
    }
    
    // Save last ON state
    if (err == ESP_OK) {
        err = nvs_set_u8(nvs_handle, NVS_KEY_LAST_ON_MODE, (uint8_t)last_on_mode);
    }
    if (err == ESP_OK) {
        err = nvs_set_u8(nvs_handle, NVS_KEY_LAST_ON_SPEED, (uint8_t)last_on_speed);
    }
    
    // Save power state (whether user turned off manually)
    if (err == ESP_OK) {
        err = nvs_set_u8(nvs_handle, NVS_KEY_POWER_STATE, user_turned_off ? 1 : 0);
    }
    
    if (err == ESP_OK) {
        err = nvs_commit(nvs_handle);
        if (err == ESP_OK) {
            ESP_LOGI(SYSTEM_TAG, "✓ Motor state successfully saved to NVS");
        } else {
            ESP_LOGE(SYSTEM_TAG, "✗ NVS commit failed: %s", esp_err_to_name(err));
        }
    } else {
        ESP_LOGE(SYSTEM_TAG, "✗ Error saving to NVS: %s", esp_err_to_name(err));
    }
    ESP_LOGI(SYSTEM_TAG, "===========================");
    
    nvs_close(nvs_handle);
    return err;
}

esp_err_t state_manager_load_and_restore(void) {
    nvs_handle_t nvs_handle;
    esp_err_t err;
    
    err = nvs_open(NVS_NAMESPACE, NVS_READONLY, &nvs_handle);
    if (err != ESP_OK) {
        ESP_LOGI(SYSTEM_TAG, "NVS not found, using default state");
        return ESP_ERR_NVS_NOT_FOUND;
    }
    
    uint8_t stored_mode = 0;
    uint8_t stored_speed = 0;
    uint8_t stored_last_mode = MOTOR_EXHAUST;  // Default to MOTOR_EXHAUST
    uint8_t stored_last_speed = 50;
    uint8_t stored_power_state = 0;
    
    // Load current motor state
    err = nvs_get_u8(nvs_handle, NVS_KEY_MODE, &stored_mode);
    if (err == ESP_OK) {
        nvs_get_u8(nvs_handle, NVS_KEY_SPEED, &stored_speed);
        nvs_get_u8(nvs_handle, NVS_KEY_LAST_ON_MODE, &stored_last_mode);
        nvs_get_u8(nvs_handle, NVS_KEY_LAST_ON_SPEED, &stored_last_speed);
        nvs_get_u8(nvs_handle, NVS_KEY_POWER_STATE, &stored_power_state);
        
        // Validate and clamp all values
        validate_and_clamp_values(&stored_mode, &stored_speed, &stored_last_mode, &stored_last_speed);
        
        // Set the loaded state in motor control module
        motor_set_last_on_state((motor_mode_t)stored_last_mode, stored_last_speed);
        motor_set_user_turned_off(stored_power_state == 1);
        
        ESP_LOGI(SYSTEM_TAG, "Loaded state from NVS - Mode: %s, Speed: %d%%, Last ON: %s@%d%%, User OFF: %s",
                 motor_mode_to_string((motor_mode_t)stored_mode), stored_speed, 
                 motor_mode_to_string((motor_mode_t)stored_last_mode), stored_last_speed,
                 stored_power_state ? "YES" : "NO");
        
        // Log reset analysis
        log_reset_analysis();
        
        // Determine if we should restore the motor state
        bool should_restore = false;
        should_restore_state(&should_restore, (motor_mode_t)stored_mode, stored_speed);
        
        if (should_restore) {
            // Restore the motor to its previous state
            motor_set_speed((motor_mode_t)stored_mode, stored_speed);
        }
        // Note: If should_restore is false, motor remains in default OFF state
        
    } else {
        ESP_LOGI(SYSTEM_TAG, "No saved state found, using defaults");
        err = ESP_ERR_NVS_NOT_FOUND;
    }
    
    nvs_close(nvs_handle);
    return err;
}

bool state_manager_was_watchdog_reset(void) {
    esp_reset_reason_t reset_reason = esp_reset_reason();
    
    // Only consider TASK_WDT and INT_WDT as true watchdog resets
    // ESP_RST_WDT can be triggered by power disconnection, so we exclude it
    return (reset_reason == ESP_RST_TASK_WDT || 
            reset_reason == ESP_RST_INT_WDT);
}

const char* state_manager_get_reset_reason_string(void) {
    esp_reset_reason_t reset_reason = esp_reset_reason();
    
    switch (reset_reason) {
        case ESP_RST_POWERON: return "POWER_ON";
        case ESP_RST_EXT: return "EXTERNAL";
        case ESP_RST_SW: return "SOFTWARE";
        case ESP_RST_PANIC: return "PANIC";
        case ESP_RST_INT_WDT: return "INT_WDT";
        case ESP_RST_TASK_WDT: return "TASK_WDT";
        case ESP_RST_WDT: return "WDT";
        case ESP_RST_DEEPSLEEP: return "DEEPSLEEP";
        case ESP_RST_BROWNOUT: return "BROWNOUT";
        case ESP_RST_SDIO: return "SDIO";
        default: return "UNKNOWN";
    }
}

esp_err_t state_manager_clear_all(void) {
    nvs_handle_t nvs_handle;
    esp_err_t err;
    
    err = nvs_open(NVS_NAMESPACE, NVS_READWRITE, &nvs_handle);
    if (err != ESP_OK) {
        ESP_LOGE(SYSTEM_TAG, "Error opening NVS handle for clear: %s", esp_err_to_name(err));
        return err;
    }
    
    ESP_LOGI(SYSTEM_TAG, "Clearing all saved state from NVS...");
    
    err = nvs_erase_all(nvs_handle);
    if (err == ESP_OK) {
        err = nvs_commit(nvs_handle);
        if (err == ESP_OK) {
            ESP_LOGI(SYSTEM_TAG, "✓ All state cleared from NVS");
        } else {
            ESP_LOGE(SYSTEM_TAG, "✗ Failed to commit NVS clear: %s", esp_err_to_name(err));
        }
    } else {
        ESP_LOGE(SYSTEM_TAG, "✗ Failed to clear NVS: %s", esp_err_to_name(err));
    }
    
    nvs_close(nvs_handle);
    return err;
}

esp_err_t state_manager_save_last_on_state(motor_mode_t mode, int speed) {
    if (mode == MOTOR_OFF || !IS_VALID_SPEED(speed) || speed == 0) {
        ESP_LOGW(SYSTEM_TAG, "Invalid last ON state: mode=%s, speed=%d", 
                 motor_mode_to_string(mode), speed);
        return ESP_ERR_INVALID_ARG;
    }
    
    nvs_handle_t nvs_handle;
    esp_err_t err;
    
    err = nvs_open(NVS_NAMESPACE, NVS_READWRITE, &nvs_handle);
    if (err != ESP_OK) {
        ESP_LOGE(SYSTEM_TAG, "Error opening NVS handle: %s", esp_err_to_name(err));
        return err;
    }
    
    // Save only the last ON state
    err = nvs_set_u8(nvs_handle, NVS_KEY_LAST_ON_MODE, (uint8_t)mode);
    if (err == ESP_OK) {
        err = nvs_set_u8(nvs_handle, NVS_KEY_LAST_ON_SPEED, (uint8_t)speed);
    }
    
    if (err == ESP_OK) {
        err = nvs_commit(nvs_handle);
    }
    
    nvs_close(nvs_handle);
    
    if (err == ESP_OK) {
        ESP_LOGI(SYSTEM_TAG, "Last ON state saved: %s @ %d%%", 
                 motor_mode_to_string(mode), speed);
    } else {
        ESP_LOGE(SYSTEM_TAG, "Failed to save last ON state: %s", esp_err_to_name(err));
    }
    
    return err;
}

esp_err_t state_manager_load_last_on_state(motor_mode_t* mode, int* speed) {
    if (!mode || !speed) return ESP_ERR_INVALID_ARG;
    
    nvs_handle_t nvs_handle;
    esp_err_t err;
    
    err = nvs_open(NVS_NAMESPACE, NVS_READONLY, &nvs_handle);
    if (err != ESP_OK) {
        return ESP_ERR_NVS_NOT_FOUND;
    }
    
    uint8_t stored_mode = MOTOR_EXHAUST;
    uint8_t stored_speed = 50;
    
    err = nvs_get_u8(nvs_handle, NVS_KEY_LAST_ON_MODE, &stored_mode);
    if (err == ESP_OK) {
        nvs_get_u8(nvs_handle, NVS_KEY_LAST_ON_SPEED, &stored_speed);
        
        // Validate values
        if (stored_mode < MOTOR_EXHAUST || stored_mode > MOTOR_INTAKE) {
            stored_mode = MOTOR_EXHAUST;
        }
        if (!IS_VALID_SPEED(stored_speed) || stored_speed == 0) {
            stored_speed = 50;
        }
        
        *mode = (motor_mode_t)stored_mode;
        *speed = stored_speed;
    }
    
    nvs_close(nvs_handle);
    return err;
}

bool state_manager_get_user_turned_off(void) {
    return motor_get_user_turned_off();
}

esp_err_t state_manager_set_user_turned_off(bool turned_off) {
    motor_set_user_turned_off(turned_off);
    
    // Save just this flag to NVS
    nvs_handle_t nvs_handle;
    esp_err_t err = nvs_open(NVS_NAMESPACE, NVS_READWRITE, &nvs_handle);
    if (err != ESP_OK) {
        return err;
    }
    
    err = nvs_set_u8(nvs_handle, NVS_KEY_POWER_STATE, turned_off ? 1 : 0);
    if (err == ESP_OK) {
        err = nvs_commit(nvs_handle);
    }
    
    nvs_close(nvs_handle);
    return err;
}

esp_err_t state_manager_get_debug_info(char* info_buffer, size_t buffer_size) {
    if (!info_buffer || buffer_size == 0) return ESP_ERR_INVALID_ARG;
    
    const motor_state_t* state = motor_get_state();
    motor_mode_t last_on_mode;
    int last_on_speed;
    motor_get_last_on_state(&last_on_mode, &last_on_speed);
    
    int written = snprintf(info_buffer, buffer_size,
        "=== STATE MANAGER DEBUG INFO ===\n"
        "Reset Reason: %s\n"
        "Watchdog Reset: %s\n"
        "Current Mode: %s\n"
        "Current Speed: %d%%\n"
        "Target Speed: %d%%\n"
        "Motor State: %s\n"
        "Last ON: %s @ %d%%\n"
        "User Turned Off: %s\n"
        "Has Saved State: %s\n"
        "===============================",
        state_manager_get_reset_reason_string(),
        state_manager_was_watchdog_reset() ? "YES" : "NO",
        motor_mode_to_string(state->mode),
        state->current_speed,
        state->target_speed,
        motor_state_to_string(state->state),
        motor_mode_to_string(last_on_mode),
        last_on_speed,
        motor_get_user_turned_off() ? "YES" : "NO",
        state_manager_has_saved_state() ? "YES" : "NO"
    );
    
    return (written < buffer_size) ? ESP_OK : ESP_ERR_INVALID_SIZE;
}

bool state_manager_has_saved_state(void) {
    nvs_handle_t nvs_handle;
    esp_err_t err = nvs_open(NVS_NAMESPACE, NVS_READONLY, &nvs_handle);
    if (err != ESP_OK) {
        return false;
    }
    
    uint8_t dummy;
    err = nvs_get_u8(nvs_handle, NVS_KEY_MODE, &dummy);
    nvs_close(nvs_handle);
    
    return (err == ESP_OK);
}