#include "scanner_engine.h"
#include "network_manager.h"
#include <esp_log.h>
#include <esp_timer.h>
#include <lwip/sockets.h>
#include <lwip/netdb.h>
#include <lwip/ip4_addr.h>
#include <lwip/icmp.h>
#include <lwip/inet.h>
#include <lwip/raw.h>
#include <fcntl.h>
#include <errno.h>
#include <freertos/FreeRTOS.h>
#include <freertos/task.h>
#include <freertos/queue.h>
#include <string.h>

static const char *TAG = "SCANNER";

// Scanner state
static bool s_initialized = false;
static QueueHandle_t s_scan_queue = NULL;
static scan_request_t s_active_scans[MAX_CONCURRENT_SCANS];
static discovered_device_t s_devices[MAX_DISCOVERED_DEVICES];
static uint32_t s_device_count = 0;
static scanner_stats_t s_stats = {0};

// Task handles
static TaskHandle_t s_scanner_task = NULL;

// Helper function to check if IP is in range
static bool ip_in_range(ip4_addr_t ip, ip4_addr_t start, ip4_addr_t end) {
    return (ip.addr >= start.addr && ip.addr <= end.addr);
}

// ARP scan implementation
static esp_err_t perform_arp_scan(scan_config_t *config, uint32_t scan_id) {
    ESP_LOGI(TAG, "Starting ARP scan from %s to %s", 
             ip4addr_ntoa(&config->start_ip), 
             ip4addr_ntoa(&config->end_ip));
    
    // Get Ethernet interface
    esp_netif_t *eth_netif = network_manager_get_ethernet_netif();
    if (!eth_netif) {
        ESP_LOGE(TAG, "No Ethernet interface available");
        return ESP_FAIL;
    }
    
    // Create raw socket for ARP
    int sock = socket(AF_INET, SOCK_DGRAM, 0);
    if (sock < 0) {
        ESP_LOGE(TAG, "Failed to create socket");
        return ESP_FAIL;
    }
    
    // Set socket timeout
    struct timeval tv;
    tv.tv_sec = config->timeout_ms / 1000;
    tv.tv_usec = (config->timeout_ms % 1000) * 1000;
    setsockopt(sock, SOL_SOCKET, SO_RCVTIMEO, &tv, sizeof(tv));
    
    uint32_t total_ips = ntohl(config->end_ip.addr) - ntohl(config->start_ip.addr) + 1;
    uint32_t scanned = 0;
    
    // Scan each IP in range
    for (uint32_t ip = ntohl(config->start_ip.addr); ip <= ntohl(config->end_ip.addr); ip++) {
        ip4_addr_t target;
        target.addr = htonl(ip);
        
        // Update progress
        scanned++;
        for (int i = 0; i < MAX_CONCURRENT_SCANS; i++) {
            if (s_active_scans[i].scan_id == scan_id) {
                s_active_scans[i].progress_percent = (scanned * 100) / total_ips;
                break;
            }
        }
        
        // Try to reach the target (this will trigger ARP)
        struct sockaddr_in dest;
        dest.sin_family = AF_INET;
        dest.sin_port = htons(7); // Echo port
        dest.sin_addr.s_addr = target.addr;
        
        char dummy = 0;
        sendto(sock, &dummy, 1, 0, (struct sockaddr*)&dest, sizeof(dest));
        
        // Small delay between requests
        vTaskDelay(pdMS_TO_TICKS(config->delay_ms));
        
        // Check if scan was cancelled
        bool cancelled = false;
        for (int i = 0; i < MAX_CONCURRENT_SCANS; i++) {
            if (s_active_scans[i].scan_id == scan_id && 
                s_active_scans[i].status == SCAN_STATUS_CANCELLED) {
                cancelled = true;
                break;
            }
        }
        if (cancelled) break;
    }
    
    close(sock);
    
    // Give time for ARP responses
    vTaskDelay(pdMS_TO_TICKS(1000));
    
    // Read ARP table to get results
    // Note: ESP-IDF doesn't provide direct ARP table access
    // In a real implementation, you'd need to use netif functions
    
    ESP_LOGI(TAG, "ARP scan completed. Scanned %lu IPs", scanned);
    s_stats.total_scans++;
    s_stats.successful_scans++;
    
    return ESP_OK;
}

// ICMP ping scan
static esp_err_t perform_ping_scan(scan_config_t *config, uint32_t scan_id) {
    ESP_LOGI(TAG, "Starting ICMP ping scan");
    
    int sock = socket(AF_INET, SOCK_RAW, IPPROTO_ICMP);
    if (sock < 0) {
        ESP_LOGE(TAG, "Failed to create raw socket (needs root/admin)");
        return ESP_FAIL;
    }
    
    // Set timeout
    struct timeval tv;
    tv.tv_sec = config->timeout_ms / 1000;
    tv.tv_usec = (config->timeout_ms % 1000) * 1000;
    setsockopt(sock, SOL_SOCKET, SO_RCVTIMEO, &tv, sizeof(tv));
    
    uint32_t total_ips = ntohl(config->end_ip.addr) - ntohl(config->start_ip.addr) + 1;
    uint32_t scanned = 0;
    
    // ICMP packet structure
    struct {
        uint8_t type;
        uint8_t code;
        uint16_t checksum;
        uint16_t id;
        uint16_t sequence;
    } icmp_packet;
    
    // Scan each IP
    for (uint32_t ip = ntohl(config->start_ip.addr); ip <= ntohl(config->end_ip.addr); ip++) {
        ip4_addr_t target;
        target.addr = htonl(ip);
        
        scanned++;
        
        // Update progress
        for (int i = 0; i < MAX_CONCURRENT_SCANS; i++) {
            if (s_active_scans[i].scan_id == scan_id) {
                s_active_scans[i].progress_percent = (scanned * 100) / total_ips;
                break;
            }
        }
        
        // Build ICMP echo request
        icmp_packet.type = 8; // Echo request
        icmp_packet.code = 0;
        icmp_packet.checksum = 0;
        icmp_packet.id = scan_id & 0xFFFF;
        icmp_packet.sequence = scanned & 0xFFFF;
        
        // Calculate checksum (simplified)
        uint32_t sum = 0;
        uint16_t *ptr = (uint16_t*)&icmp_packet;
        for (int i = 0; i < sizeof(icmp_packet)/2; i++) {
            sum += ptr[i];
        }
        icmp_packet.checksum = ~(sum + (sum >> 16));
        
        struct sockaddr_in dest;
        dest.sin_family = AF_INET;
        dest.sin_addr.s_addr = target.addr;
        
        // Send ICMP packet
        int64_t send_time = esp_timer_get_time();
        sendto(sock, &icmp_packet, sizeof(icmp_packet), 0, 
               (struct sockaddr*)&dest, sizeof(dest));
        
        // Try to receive reply
        char recv_buf[128];
        struct sockaddr_in from;
        socklen_t fromlen = sizeof(from);
        
        int n = recvfrom(sock, recv_buf, sizeof(recv_buf), 0,
                        (struct sockaddr*)&from, &fromlen);
        
        if (n > 0) {
            int64_t recv_time = esp_timer_get_time();
            uint32_t rtt_ms = (recv_time - send_time) / 1000;
            
            // Add to discovered devices
            if (s_device_count < MAX_DISCOVERED_DEVICES) {
                discovered_device_t *dev = &s_devices[s_device_count++];
                dev->ip_addr = target;
                dev->response_time = rtt_ms;
                dev->is_active = true;
                dev->last_seen = xTaskGetTickCount();
                
                ESP_LOGI(TAG, "Host %s is alive (RTT: %lu ms)", 
                        ip4addr_ntoa(&target), rtt_ms);
                s_stats.devices_discovered++;
            }
        }
        
        vTaskDelay(pdMS_TO_TICKS(config->delay_ms));
        
        // Check if cancelled
        bool cancelled = false;
        for (int i = 0; i < MAX_CONCURRENT_SCANS; i++) {
            if (s_active_scans[i].scan_id == scan_id && 
                s_active_scans[i].status == SCAN_STATUS_CANCELLED) {
                cancelled = true;
                break;
            }
        }
        if (cancelled) break;
    }
    
    close(sock);
    
    ESP_LOGI(TAG, "Ping scan completed. Found %lu active hosts", s_device_count);
    return ESP_OK;
}

// TCP port scan
static esp_err_t perform_port_scan(scan_config_t *config, uint32_t scan_id) {
    ESP_LOGI(TAG, "Starting TCP port scan");
    
    uint32_t total_operations = 
        (ntohl(config->end_ip.addr) - ntohl(config->start_ip.addr) + 1) * 
        (config->port_end - config->port_start + 1);
    uint32_t completed = 0;
    
    // Scan each IP
    for (uint32_t ip = ntohl(config->start_ip.addr); ip <= ntohl(config->end_ip.addr); ip++) {
        ip4_addr_t target;
        target.addr = htonl(ip);
        
        // Scan each port
        for (uint16_t port = config->port_start; port <= config->port_end; port++) {
            int sock = socket(AF_INET, SOCK_STREAM, 0);
            if (sock < 0) continue;
            
            // Set non-blocking mode
            int flags = fcntl(sock, F_GETFL, 0);
            fcntl(sock, F_SETFL, flags | O_NONBLOCK);
            
            // Set timeout
            struct timeval tv;
            tv.tv_sec = config->timeout_ms / 1000;
            tv.tv_usec = (config->timeout_ms % 1000) * 1000;
            setsockopt(sock, SOL_SOCKET, SO_SNDTIMEO, &tv, sizeof(tv));
            
            struct sockaddr_in dest;
            dest.sin_family = AF_INET;
            dest.sin_port = htons(port);
            dest.sin_addr.s_addr = target.addr;
            
            // Try to connect
            int result = connect(sock, (struct sockaddr*)&dest, sizeof(dest));
            
            if (result == 0 || (result < 0 && errno == EINPROGRESS)) {
                // Wait for connection or timeout
                fd_set write_fds;
                FD_ZERO(&write_fds);
                FD_SET(sock, &write_fds);
                
                result = select(sock + 1, NULL, &write_fds, NULL, &tv);
                
                if (result > 0 && FD_ISSET(sock, &write_fds)) {
                    // Check if really connected
                    int error = 0;
                    socklen_t len = sizeof(error);
                    getsockopt(sock, SOL_SOCKET, SO_ERROR, &error, &len);
                    
                    if (error == 0) {
                        ESP_LOGI(TAG, "Port %d open on %s", port, ip4addr_ntoa(&target));
                        
                        // Find or add device
                        for (int i = 0; i < s_device_count; i++) {
                            if (s_devices[i].ip_addr.addr == target.addr) {
                                if (s_devices[i].open_port_count < 255) {
                                    s_devices[i].open_ports[s_devices[i].open_port_count++] = port;
                                }
                                break;
                            }
                        }
                        s_stats.ports_scanned++;
                    }
                }
            }
            
            close(sock);
            
            completed++;
            
            // Update progress
            for (int i = 0; i < MAX_CONCURRENT_SCANS; i++) {
                if (s_active_scans[i].scan_id == scan_id) {
                    s_active_scans[i].progress_percent = (completed * 100) / total_operations;
                    break;
                }
            }
            
            vTaskDelay(pdMS_TO_TICKS(config->delay_ms));
            
            // Check if cancelled
            bool cancelled = false;
            for (int i = 0; i < MAX_CONCURRENT_SCANS; i++) {
                if (s_active_scans[i].scan_id == scan_id && 
                    s_active_scans[i].status == SCAN_STATUS_CANCELLED) {
                    cancelled = true;
                    break;
                }
            }
            if (cancelled) goto scan_end;
        }
    }
    
scan_end:
    ESP_LOGI(TAG, "Port scan completed");
    return ESP_OK;
}

// Scanner task
static void scanner_task(void *pvParameters) {
    scan_request_t request;
    
    while (1) {
        if (xQueueReceive(s_scan_queue, &request, portMAX_DELAY) == pdTRUE) {
            ESP_LOGI(TAG, "Processing scan request %lu", request.scan_id);
            
            // Mark scan as running
            for (int i = 0; i < MAX_CONCURRENT_SCANS; i++) {
                if (s_active_scans[i].scan_id == request.scan_id) {
                    s_active_scans[i].status = SCAN_STATUS_RUNNING;
                    s_active_scans[i].start_time = xTaskGetTickCount();
                    break;
                }
            }
            
            // Perform scan based on type
            esp_err_t result = ESP_OK;
            switch (request.config.type) {
                case SCAN_TYPE_ARP:
                    result = perform_arp_scan(&request.config, request.scan_id);
                    break;
                case SCAN_TYPE_ICMP_PING:
                    result = perform_ping_scan(&request.config, request.scan_id);
                    break;
                case SCAN_TYPE_PORT_TCP:
                    result = perform_port_scan(&request.config, request.scan_id);
                    break;
                default:
                    ESP_LOGW(TAG, "Unsupported scan type: %d", request.config.type);
                    result = ESP_ERR_NOT_SUPPORTED;
            }
            
            // Mark scan as completed
            for (int i = 0; i < MAX_CONCURRENT_SCANS; i++) {
                if (s_active_scans[i].scan_id == request.scan_id) {
                    s_active_scans[i].status = (result == ESP_OK) ? 
                        SCAN_STATUS_COMPLETED : SCAN_STATUS_ERROR;
                    s_active_scans[i].end_time = xTaskGetTickCount();
                    s_active_scans[i].devices_found = s_device_count;
                    break;
                }
            }
        }
    }
}

// Public functions
esp_err_t scanner_engine_init(void) {
    if (s_initialized) {
        return ESP_OK;
    }
    
    ESP_LOGI(TAG, "Initializing scanner engine");
    
    // Create scan queue
    s_scan_queue = xQueueCreate(10, sizeof(scan_request_t));
    if (!s_scan_queue) {
        ESP_LOGE(TAG, "Failed to create scan queue");
        return ESP_ERR_NO_MEM;
    }
    
    // Initialize scan slots
    memset(s_active_scans, 0, sizeof(s_active_scans));
    memset(s_devices, 0, sizeof(s_devices));
    
    // Create scanner task
    BaseType_t ret = xTaskCreatePinnedToCore(
        scanner_task,
        "scanner",
        8192,
        NULL,
        5,
        &s_scanner_task,
        0  // Core 0 for network operations
    );
    
    if (ret != pdPASS) {
        ESP_LOGE(TAG, "Failed to create scanner task");
        vQueueDelete(s_scan_queue);
        return ESP_FAIL;
    }
    
    s_initialized = true;
    ESP_LOGI(TAG, "Scanner engine initialized successfully");
    return ESP_OK;
}

esp_err_t scanner_engine_deinit(void) {
    if (!s_initialized) {
        return ESP_OK;
    }
    
    // Delete task
    if (s_scanner_task) {
        vTaskDelete(s_scanner_task);
        s_scanner_task = NULL;
    }
    
    // Delete queue
    if (s_scan_queue) {
        vQueueDelete(s_scan_queue);
        s_scan_queue = NULL;
    }
    
    s_initialized = false;
    return ESP_OK;
}

esp_err_t scanner_engine_start_scan(const scan_config_t *config, uint32_t *scan_id) {
    if (!s_initialized || !config || !scan_id) {
        return ESP_ERR_INVALID_ARG;
    }
    
    // Find free slot
    int slot = -1;
    for (int i = 0; i < MAX_CONCURRENT_SCANS; i++) {
        if (s_active_scans[i].status == SCAN_STATUS_IDLE) {
            slot = i;
            break;
        }
    }
    
    if (slot < 0) {
        ESP_LOGE(TAG, "No free scan slots");
        return ESP_ERR_NO_MEM;
    }
    
    // Generate scan ID
    static uint32_t next_scan_id = 1;
    *scan_id = next_scan_id++;
    
    // Prepare scan request
    scan_request_t *request = &s_active_scans[slot];
    request->scan_id = *scan_id;
    memcpy(&request->config, config, sizeof(scan_config_t));
    request->status = SCAN_STATUS_IDLE;
    request->progress_percent = 0;
    request->devices_found = 0;
    
    // Queue scan request
    scan_request_t queue_item = *request;
    if (xQueueSend(s_scan_queue, &queue_item, 0) != pdTRUE) {
        ESP_LOGE(TAG, "Failed to queue scan request");
        request->status = SCAN_STATUS_IDLE;
        return ESP_FAIL;
    }
    
    ESP_LOGI(TAG, "Scan %lu queued successfully", *scan_id);
    return ESP_OK;
}

esp_err_t scanner_engine_stop_scan(uint32_t scan_id) {
    if (!s_initialized) {
        return ESP_ERR_INVALID_STATE;
    }
    
    // Mark scan as cancelled
    for (int i = 0; i < MAX_CONCURRENT_SCANS; i++) {
        if (s_active_scans[i].scan_id == scan_id) {
            s_active_scans[i].status = SCAN_STATUS_CANCELLED;
            ESP_LOGI(TAG, "Scan %lu cancelled", scan_id);
            return ESP_OK;
        }
    }
    
    return ESP_ERR_NOT_FOUND;
}

void scanner_engine_process_queue(void) {
    // Queue processing is handled by scanner task
    // This function is kept for API compatibility
}

esp_err_t scanner_engine_get_results(discovered_device_t *devices, size_t *count) {
    if (!devices || !count) {
        return ESP_ERR_INVALID_ARG;
    }
    
    size_t copy_count = (*count < s_device_count) ? *count : s_device_count;
    memcpy(devices, s_devices, copy_count * sizeof(discovered_device_t));
    *count = copy_count;
    
    return ESP_OK;
}

esp_err_t scanner_engine_get_scan_status(uint32_t scan_id, scan_request_t *status) {
    if (!status) {
        return ESP_ERR_INVALID_ARG;
    }
    
    for (int i = 0; i < MAX_CONCURRENT_SCANS; i++) {
        if (s_active_scans[i].scan_id == scan_id) {
            memcpy(status, &s_active_scans[i], sizeof(scan_request_t));
            return ESP_OK;
        }
    }
    
    return ESP_ERR_NOT_FOUND;
}

esp_err_t scanner_engine_get_stats(scanner_stats_t *stats) {
    if (!stats) {
        return ESP_ERR_INVALID_ARG;
    }
    
    memcpy(stats, &s_stats, sizeof(scanner_stats_t));
    return ESP_OK;
} 