#include "security_manager.h"
#include "config_manager.h"
#include "log_manager.h"
#include "esp_log.h"
#include "esp_err.h"
#include "esp_random.h"
#include "mbedtls/md.h"
#include "mbedtls/sha256.h"
#include <string.h>
#include <time.h>

static const char *TAG = "SECURITY_MANAGER";

// Session management
#define MAX_SESSIONS 10
#define SESSION_TIMEOUT_SECONDS 3600 // 1 hour

typedef struct {
    char session_id[33];
    char username[32];
    time_t created_time;
    time_t last_activity;
    bool active;
} session_t;

static session_t s_sessions[MAX_SESSIONS];
static bool s_security_manager_initialized = false;
static SemaphoreHandle_t s_security_mutex = NULL;

// Rate limiting
#define MAX_LOGIN_ATTEMPTS 5
#define LOGIN_TIMEOUT_SECONDS 300 // 5 minutes

typedef struct {
    char ip_address[16];
    int failed_attempts;
    time_t last_attempt;
    bool blocked;
} rate_limit_entry_t;

static rate_limit_entry_t s_rate_limit_entries[20];
static int s_rate_limit_count = 0;

// Password hashing constants
#define SALT_LENGTH 16
#define HASH_LENGTH 64 // SHA-256 hex string length

esp_err_t security_manager_init(void)
{
    ESP_LOGI(TAG, "Initializing security manager");
    
    // Create mutex for thread safety
    s_security_mutex = xSemaphoreCreateMutex();
    if (s_security_mutex == NULL) {
        ESP_LOGE(TAG, "Failed to create security mutex");
        return ESP_ERR_NO_MEM;
    }
    
    // Initialize sessions
    memset(s_sessions, 0, sizeof(s_sessions));
    
    // Initialize rate limiting
    memset(s_rate_limit_entries, 0, sizeof(s_rate_limit_entries));
    s_rate_limit_count = 0;
    
    s_security_manager_initialized = true;
    ESP_LOGI(TAG, "Security manager initialized successfully");
    return ESP_OK;
}

esp_err_t security_manager_deinit(void)
{
    if (s_security_mutex != NULL) {
        vSemaphoreDelete(s_security_mutex);
        s_security_mutex = NULL;
    }
    
    s_security_manager_initialized = false;
    ESP_LOGI(TAG, "Security manager deinitialized");
    return ESP_OK;
}

static void generate_session_id(char* session_id, size_t size)
{
    if (session_id == NULL || size < 33) {
        ESP_LOGE(TAG, "Invalid parameters for session ID generation");
        return;
    }
    
    // Generate random session ID
    uint8_t random_bytes[16];
    esp_fill_random(random_bytes, sizeof(random_bytes));
    
    // Convert to hex string with bounds checking
    for (int i = 0; i < 16 && (i * 2 + 2) < size; i++) {
        snprintf(session_id + (i * 2), size - (i * 2), "%02x", random_bytes[i]);
    }
    session_id[32] = '\0';
}

// Generate SHA-256 hash with salt
static esp_err_t hash_password(const char* password, const char* salt, char* hash, size_t hash_size)
{
    if (password == NULL || salt == NULL || hash == NULL || hash_size < HASH_LENGTH + 1) {
        return ESP_ERR_INVALID_ARG;
    }
    
    mbedtls_sha256_context ctx;
    unsigned char digest[32];
    char salted_password[256];
    
    // Create salted password
    snprintf(salted_password, sizeof(salted_password), "%s%s", password, salt);
    
    // Initialize SHA-256
    mbedtls_sha256_init(&ctx);
    mbedtls_sha256_starts(&ctx, 0);
    mbedtls_sha256_update(&ctx, (const unsigned char*)salted_password, strlen(salted_password));
    mbedtls_sha256_finish(&ctx, digest);
    mbedtls_sha256_free(&ctx);
    
    // Convert to hex string
    for (int i = 0; i < 32 && (i * 2 + 2) < hash_size; i++) {
        snprintf(hash + (i * 2), hash_size - (i * 2), "%02x", digest[i]);
    }
    
    return ESP_OK;
}

// Generate salt and hash for new password
esp_err_t security_manager_hash_password(const char* password, char* salt, size_t salt_size, char* hash, size_t hash_size)
{
    if (password == NULL || salt == NULL || hash == NULL || 
        salt_size < SALT_LENGTH + 1 || hash_size < HASH_LENGTH + 1) {
        return ESP_ERR_INVALID_ARG;
    }
    
    // Generate random salt
    uint8_t random_bytes[SALT_LENGTH];
    esp_fill_random(random_bytes, sizeof(random_bytes));
    
    // Convert salt to hex string
    for (int i = 0; i < SALT_LENGTH && (i * 2 + 2) < salt_size; i++) {
        snprintf(salt + (i * 2), salt_size - (i * 2), "%02x", random_bytes[i]);
    }
    salt[SALT_LENGTH * 2] = '\0';
    
    // Hash password with salt
    return hash_password(password, salt, hash, hash_size);
}

bool verify_password_hash(const char* password, const char* stored_hash)
{
    if (password == NULL || stored_hash == NULL) {
        return false;
    }
    
    // For backward compatibility, check if stored hash is plain text
    // In production, all passwords should be hashed
    if (strlen(stored_hash) < HASH_LENGTH) {
        // Plain text password (legacy)
        return (strcmp(password, stored_hash) == 0);
    }
    
    // Extract salt from stored hash (assuming format: salt:hash)
    char salt[SALT_LENGTH * 2 + 1] = {0};
    char hash[HASH_LENGTH + 1] = {0};
    
    char* colon_pos = strchr(stored_hash, ':');
    if (colon_pos == NULL) {
        // No salt separator, assume it's a hash without salt
        strncpy(hash, stored_hash, HASH_LENGTH);
        hash[HASH_LENGTH] = '\0';
    } else {
        // Extract salt and hash
        size_t salt_len = colon_pos - stored_hash;
        if (salt_len >= sizeof(salt)) {
            salt_len = sizeof(salt) - 1;
        }
        strncpy(salt, stored_hash, salt_len);
        salt[salt_len] = '\0';
        
        strncpy(hash, colon_pos + 1, HASH_LENGTH);
        hash[HASH_LENGTH] = '\0';
    }
    
    // Hash the provided password with the salt
    char computed_hash[HASH_LENGTH + 1];
    if (hash_password(password, salt, computed_hash, sizeof(computed_hash)) != ESP_OK) {
        return false;
    }
    
    // Compare hashes
    return (strcmp(computed_hash, hash) == 0);
}

static bool is_ip_rate_limited(const char* ip_address)
{
    time_t now = time(NULL);
    
    // Clean up old entries
    for (int i = 0; i < s_rate_limit_count; i++) {
        if (now - s_rate_limit_entries[i].last_attempt > LOGIN_TIMEOUT_SECONDS) {
            // Remove expired entry
            if (i < s_rate_limit_count - 1) {
                memmove(&s_rate_limit_entries[i], &s_rate_limit_entries[i + 1], 
                       (s_rate_limit_count - i - 1) * sizeof(rate_limit_entry_t));
            }
            s_rate_limit_count--;
            i--; // Recheck this index
        }
    }
    
    // Check if IP is blocked
    for (int i = 0; i < s_rate_limit_count; i++) {
        if (strcmp(s_rate_limit_entries[i].ip_address, ip_address) == 0) {
            if (s_rate_limit_entries[i].blocked) {
                return true; // IP is blocked
            }
            break;
        }
    }
    
    return false;
}

static void record_failed_login(const char* ip_address)
{
    time_t now = time(NULL);
    
    // Find existing entry
    for (int i = 0; i < s_rate_limit_count; i++) {
        if (strcmp(s_rate_limit_entries[i].ip_address, ip_address) == 0) {
            s_rate_limit_entries[i].failed_attempts++;
            s_rate_limit_entries[i].last_attempt = now;
            
            if (s_rate_limit_entries[i].failed_attempts >= MAX_LOGIN_ATTEMPTS) {
                s_rate_limit_entries[i].blocked = true;
                log_manager_log_warn(TAG, "IP %s blocked due to too many failed login attempts", ip_address);
            }
            return;
        }
    }
    
    // Add new entry
    if (s_rate_limit_count < 20) {
        strncpy(s_rate_limit_entries[s_rate_limit_count].ip_address, ip_address, 15);
        s_rate_limit_entries[s_rate_limit_count].failed_attempts = 1;
        s_rate_limit_entries[s_rate_limit_count].last_attempt = now;
        s_rate_limit_entries[s_rate_limit_count].blocked = false;
        s_rate_limit_count++;
    }
}

esp_err_t security_manager_authenticate(const char* username, const char* password, const char* ip_address, char* session_id, size_t session_id_size)
{
    if (!s_security_manager_initialized || s_security_mutex == NULL) {
        return ESP_ERR_INVALID_STATE;
    }
    
    if (username == NULL || password == NULL || ip_address == NULL || session_id == NULL) {
        return ESP_ERR_INVALID_ARG;
    }
    
    if (xSemaphoreTake(s_security_mutex, pdMS_TO_TICKS(1000)) != pdTRUE) {
        return ESP_ERR_TIMEOUT;
    }
    
    esp_err_t ret = ESP_ERR_INVALID_ARG;
    
    // Check rate limiting
    if (is_ip_rate_limited(ip_address)) {
        log_manager_log_warn(TAG, "Login attempt from blocked IP: %s", ip_address);
        xSemaphoreGive(s_security_mutex);
        return ESP_ERR_INVALID_STATE;
    }
    
    // Get configuration
    const config_t* config = config_manager_get_config();
    if (config == NULL) {
        xSemaphoreGive(s_security_mutex);
        return ESP_ERR_INVALID_STATE;
    }
    
    // Verify credentials
    if (strcmp(username, config->web_server.username) == 0 &&
        verify_password_hash(password, config->web_server.password)) {
        
        // Generate session ID
        generate_session_id(session_id, session_id_size);
        
        // Create session
        time_t now = time(NULL);
        for (int i = 0; i < MAX_SESSIONS; i++) {
            if (!s_sessions[i].active) {
                strncpy(s_sessions[i].session_id, session_id, 32);
                strncpy(s_sessions[i].username, username, 31);
                s_sessions[i].created_time = now;
                s_sessions[i].last_activity = now;
                s_sessions[i].active = true;
                
                log_manager_log_info(TAG, "User %s authenticated successfully from IP %s", username, ip_address);
                ret = ESP_OK;
                break;
            }
        }
        
        if (ret != ESP_OK) {
            log_manager_log_warn(TAG, "No available session slots for user %s", username);
            ret = ESP_ERR_NO_MEM;
        }
    } else {
        // Record failed attempt
        record_failed_login(ip_address);
        log_manager_log_warn(TAG, "Failed login attempt for user %s from IP %s", username, ip_address);
        ret = ESP_ERR_INVALID_ARG;
    }
    
    xSemaphoreGive(s_security_mutex);
    return ret;
}

esp_err_t security_manager_validate_session(const char* session_id, char* username, size_t username_size)
{
    if (!s_security_manager_initialized || s_security_mutex == NULL) {
        return ESP_ERR_INVALID_STATE;
    }
    
    if (session_id == NULL || username == NULL) {
        return ESP_ERR_INVALID_ARG;
    }
    
    if (xSemaphoreTake(s_security_mutex, pdMS_TO_TICKS(1000)) != pdTRUE) {
        return ESP_ERR_TIMEOUT;
    }
    
    esp_err_t ret = ESP_ERR_INVALID_ARG;
    time_t now = time(NULL);
    
    // Find and validate session
    for (int i = 0; i < MAX_SESSIONS; i++) {
        if (s_sessions[i].active && strcmp(s_sessions[i].session_id, session_id) == 0) {
            // Check session timeout
            if (now - s_sessions[i].last_activity > SESSION_TIMEOUT_SECONDS) {
                // Session expired
                s_sessions[i].active = false;
                log_manager_log_info(TAG, "Session expired for user %s", s_sessions[i].username);
                ret = ESP_ERR_INVALID_STATE;
            } else {
                // Update last activity
                s_sessions[i].last_activity = now;
                strncpy(username, s_sessions[i].username, username_size - 1);
                username[username_size - 1] = '\0';
                ret = ESP_OK;
            }
            break;
        }
    }
    
    xSemaphoreGive(s_security_mutex);
    return ret;
}

esp_err_t security_manager_logout(const char* session_id)
{
    if (!s_security_manager_initialized || s_security_mutex == NULL) {
        return ESP_ERR_INVALID_STATE;
    }
    
    if (session_id == NULL) {
        return ESP_ERR_INVALID_ARG;
    }
    
    if (xSemaphoreTake(s_security_mutex, pdMS_TO_TICKS(1000)) != pdTRUE) {
        return ESP_ERR_TIMEOUT;
    }
    
    esp_err_t ret = ESP_ERR_INVALID_ARG;
    
    // Find and invalidate session
    for (int i = 0; i < MAX_SESSIONS; i++) {
        if (s_sessions[i].active && strcmp(s_sessions[i].session_id, session_id) == 0) {
            log_manager_log_info(TAG, "User %s logged out", s_sessions[i].username);
            s_sessions[i].active = false;
            ret = ESP_OK;
            break;
        }
    }
    
    xSemaphoreGive(s_security_mutex);
    return ret;
}

esp_err_t security_manager_cleanup_expired_sessions(void)
{
    if (!s_security_manager_initialized || s_security_mutex == NULL) {
        return ESP_ERR_INVALID_STATE;
    }
    
    if (xSemaphoreTake(s_security_mutex, pdMS_TO_TICKS(1000)) != pdTRUE) {
        return ESP_ERR_TIMEOUT;
    }
    
    time_t now = time(NULL);
    int expired_count = 0;
    
    // Clean up expired sessions
    for (int i = 0; i < MAX_SESSIONS; i++) {
        if (s_sessions[i].active && 
            (now - s_sessions[i].last_activity > SESSION_TIMEOUT_SECONDS)) {
            s_sessions[i].active = false;
            expired_count++;
        }
    }
    
    if (expired_count > 0) {
        log_manager_log_info(TAG, "Cleaned up %d expired sessions", expired_count);
    }
    
    xSemaphoreGive(s_security_mutex);
    return ESP_OK;
}

esp_err_t security_manager_get_active_sessions(int* count)
{
    if (!s_security_manager_initialized || s_security_mutex == NULL) {
        return ESP_ERR_INVALID_STATE;
    }
    
    if (count == NULL) {
        return ESP_ERR_INVALID_ARG;
    }
    
    if (xSemaphoreTake(s_security_mutex, pdMS_TO_TICKS(1000)) != pdTRUE) {
        return ESP_ERR_TIMEOUT;
    }
    
    *count = 0;
    for (int i = 0; i < MAX_SESSIONS; i++) {
        if (s_sessions[i].active) {
            (*count)++;
        }
    }
    
    xSemaphoreGive(s_security_mutex);
    return ESP_OK;
}

esp_err_t security_manager_is_authentication_required(bool* required)
{
    if (!s_security_manager_initialized) {
        return ESP_ERR_INVALID_STATE;
    }
    
    if (required == NULL) {
        return ESP_ERR_INVALID_ARG;
    }
    
    const config_t* config = config_manager_get_config();
    if (config == NULL) {
        return ESP_ERR_INVALID_STATE;
    }
    
    *required = config->web_server.auth_enabled;
    return ESP_OK;
} 