#include "scanner_engine.h"
#include "sd_manager.h"
#include "network_manager.h"
#include <esp_log.h>
#include <freertos/FreeRTOS.h>
#include <freertos/task.h>
#include <freertos/semphr.h>
#include <lwip/ip_addr.h>
#include <lwip/netdb.h>
#include <lwip/sockets.h>
#include <lwip/raw.h>
#include <lwip/ip4.h>
#include <lwip/icmp.h>
#include <lwip/etharp.h>
#include <lwip/dns.h>
#include <lwip/dhcp.h>
#include <lwip/netif.h>
#include <lwip/ethip6.h>
#include <string.h>
#include <stdio.h>
#include <time.h>
#include <errno.h>
#include "esp_netif.h"

static const char *TAG = "SCANNER_ENGINE";

// Scanner state
static bool s_scanner_initialized = false;
static SemaphoreHandle_t s_scanner_mutex = NULL;
static discovered_device_t s_devices[MAX_DEVICES];
static size_t s_device_count = 0;
static scanner_stats_t s_stats = {0};

// DNS monitoring state
static bool s_dns_monitoring = false;
static TaskHandle_t s_dns_task = NULL;

// DHCP monitoring state  
static bool s_dhcp_monitoring = false;
static TaskHandle_t s_dhcp_task = NULL;

// Packet capture state
static bool s_capture_running = false;
static TaskHandle_t s_capture_task = NULL;

// MAC vendor database
static mac_vendor_entry_t s_mac_vendors[MAX_MAC_VENDORS];
static size_t s_mac_vendor_count = 0;

// Function prototypes
static esp_err_t perform_arp_scan(const scan_config_t *config);
static esp_err_t perform_tcp_port_scan(const scan_config_t *config);
static esp_err_t perform_icmp_ping_sweep(const scan_config_t *config);
static esp_err_t add_discovered_device(const ip4_addr_t *ip, const uint8_t *mac, const char *hostname);
static esp_err_t resolve_hostname(const ip4_addr_t *ip, char *hostname, size_t len);

// DNS monitoring task
static void dns_monitor_task(void *arg) {
    ESP_LOGI(TAG, "DNS monitoring task started");

    int sock = socket(AF_INET, SOCK_DGRAM, 0);
    if (sock < 0) {
        ESP_LOGE(TAG, "Failed to create DNS socket");
        vTaskDelete(NULL);
        return;
    }

    struct sockaddr_in addr;
    memset(&addr, 0, sizeof(addr));
    addr.sin_family = AF_INET;
    addr.sin_port = htons(53);
    addr.sin_addr.s_addr = INADDR_ANY;

    if (bind(sock, (struct sockaddr*)&addr, sizeof(addr)) < 0) {
        ESP_LOGE(TAG, "Failed to bind DNS socket");
        close(sock);
        vTaskDelete(NULL);
        return;
    }

    uint8_t buffer[512];
    while (s_dns_monitoring) {
    struct sockaddr_in from;
    socklen_t fromlen = sizeof(from);

        int len = recvfrom(sock, buffer, sizeof(buffer), 0,
                                  (struct sockaddr*)&from, &fromlen);
        
        if (len > 0) {
            // Parse DNS packet and log query
            ESP_LOGI(TAG, "DNS query from %s", inet_ntoa(from.sin_addr));
        }
        
        vTaskDelay(pdMS_TO_TICKS(100));
    }

    close(sock);
    ESP_LOGI(TAG, "DNS monitoring task ended");
    vTaskDelete(NULL);
}

// DHCP monitoring task
static void dhcp_monitor_task(void *arg) {
    ESP_LOGI(TAG, "DHCP monitoring task started");

    int sock = socket(AF_INET, SOCK_DGRAM, 0);
    if (sock < 0) {
        ESP_LOGE(TAG, "Failed to create DHCP socket");
        vTaskDelete(NULL);
        return;
    }

    struct sockaddr_in addr;
    memset(&addr, 0, sizeof(addr));
    addr.sin_family = AF_INET;
    addr.sin_port = htons(67);
    addr.sin_addr.s_addr = INADDR_ANY;

    if (bind(sock, (struct sockaddr*)&addr, sizeof(addr)) < 0) {
        ESP_LOGE(TAG, "Failed to bind DHCP socket");
        close(sock);
        vTaskDelete(NULL);
        return;
    }

    uint8_t buffer[512];
    while (s_dhcp_monitoring) {
    struct sockaddr_in from;
    socklen_t fromlen = sizeof(from);

        int len = recvfrom(sock, buffer, sizeof(buffer), 0,
                                  (struct sockaddr*)&from, &fromlen);
        
        if (len > 0) {
            // Parse DHCP packet and log lease
            ESP_LOGI(TAG, "DHCP packet from %s", inet_ntoa(from.sin_addr));
        }
        
        vTaskDelay(pdMS_TO_TICKS(100));
    }

    close(sock);
    ESP_LOGI(TAG, "DHCP monitoring task ended");
    vTaskDelete(NULL);
}

// Packet capture task (stub implementation)
static void packet_capture_task(void *arg) {
    ESP_LOGI(TAG, "Packet capture not implemented for ESP-IDF");
    // Packet capture not implemented - ESP-IDF compatible APIs required
    // For now, this is a placeholder that simulates capture activity
    while (s_capture_running) {
        vTaskDelay(pdMS_TO_TICKS(100));
    }
    ESP_LOGI(TAG, "Packet capture task ended");
}

// Initialize scanner engine
esp_err_t scanner_engine_init(void) {
    if (s_scanner_initialized) {
    return ESP_OK;
}

    s_scanner_mutex = xSemaphoreCreateMutex();
    if (!s_scanner_mutex) {
        ESP_LOGE(TAG, "Failed to create scanner mutex");
        return ESP_ERR_NO_MEM;
    }

    // Initialize device list
    memset(s_devices, 0, sizeof(s_devices));
    s_device_count = 0;

    // Initialize statistics
    memset(&s_stats, 0, sizeof(s_stats));

    // Load MAC vendor database
    if (load_mac_vendor_database() != ESP_OK) {
        ESP_LOGW(TAG, "Failed to load MAC vendor database");
    }

    s_scanner_initialized = true;
    ESP_LOGI(TAG, "Scanner engine initialized");
    return ESP_OK;
}

// Deinitialize scanner engine
esp_err_t scanner_engine_deinit(void) {
    if (!s_scanner_initialized) {
    return ESP_OK;
}

    // Stop all monitoring tasks
    scanner_stop_dns_monitor();
    scanner_stop_dhcp_monitor();
    scanner_stop_passive_sniff();

    if (s_scanner_mutex) {
        vSemaphoreDelete(s_scanner_mutex);
        s_scanner_mutex = NULL;
    }

    s_scanner_initialized = false;
    ESP_LOGI(TAG, "Scanner engine deinitialized");
    return ESP_OK;
}

// Start ARP scan
esp_err_t scanner_arp_scan(const scan_config_t *config) {
    if (!s_scanner_initialized) {
        return ESP_ERR_INVALID_STATE;
    }

    return perform_arp_scan(config);
}

// Start TCP port scan
esp_err_t scanner_tcp_port_scan(const scan_config_t *config) {
    if (!s_scanner_initialized) {
        return ESP_ERR_INVALID_STATE;
    }

    return perform_tcp_port_scan(config);
}

// Start ICMP ping sweep
esp_err_t scanner_icmp_ping_sweep(const scan_config_t *config) {
    if (!s_scanner_initialized) {
        return ESP_ERR_INVALID_STATE;
    }

    return perform_icmp_ping_sweep(config);
}

// Start DNS monitoring
esp_err_t scanner_start_dns_monitor(uint32_t duration_seconds) {
    if (s_dns_monitoring) {
        return ESP_ERR_INVALID_STATE;
    }

    s_dns_monitoring = true;
    if (xTaskCreate(dns_monitor_task, "dns_monitor", 4096, NULL, 5, &s_dns_task) != pdPASS) {
        s_dns_monitoring = false;
        return ESP_ERR_NO_MEM;
    }

    ESP_LOGI(TAG, "DNS monitoring started for %u seconds", duration_seconds);
    return ESP_OK;
}

// Stop DNS monitoring
esp_err_t scanner_stop_dns_monitor(void) {
    if (!s_dns_monitoring) {
        return ESP_ERR_INVALID_STATE;
    }

    s_dns_monitoring = false;
    if (s_dns_task) {
        vTaskDelete(s_dns_task);
        s_dns_task = NULL;
    }

    ESP_LOGI(TAG, "DNS monitoring stopped");
    return ESP_OK;
}

// Get DNS queries (stub implementation)
esp_err_t scanner_get_dns_queries(dns_query_t *queries, size_t max_queries, size_t *query_count) {
    if (!queries || !query_count) {
        return ESP_ERR_INVALID_ARG;
    }

    *query_count = 0;
    ESP_LOGI(TAG, "DNS queries retrieval not implemented");
    return ESP_OK;
}

// Start DHCP monitoring
esp_err_t scanner_start_dhcp_monitor(uint32_t duration_seconds) {
    if (s_dhcp_monitoring) {
        return ESP_ERR_INVALID_STATE;
    }

    s_dhcp_monitoring = true;
    if (xTaskCreate(dhcp_monitor_task, "dhcp_monitor", 4096, NULL, 5, &s_dhcp_task) != pdPASS) {
        s_dhcp_monitoring = false;
        return ESP_ERR_NO_MEM;
    }

    ESP_LOGI(TAG, "DHCP monitoring started for %u seconds", duration_seconds);
    return ESP_OK;
}

// Stop DHCP monitoring
esp_err_t scanner_stop_dhcp_monitor(void) {
    if (!s_dhcp_monitoring) {
        return ESP_ERR_INVALID_STATE;
    }

    s_dhcp_monitoring = false;
    if (s_dhcp_task) {
        vTaskDelete(s_dhcp_task);
        s_dhcp_task = NULL;
    }

    ESP_LOGI(TAG, "DHCP monitoring stopped");
    return ESP_OK;
}

// Get DHCP leases (stub implementation)
esp_err_t scanner_get_dhcp_leases(dhcp_lease_t *leases, size_t max_leases, size_t *lease_count) {
    if (!leases || !lease_count) {
        return ESP_ERR_INVALID_ARG;
    }

    *lease_count = 0;
    ESP_LOGI(TAG, "DHCP leases retrieval not implemented");
    return ESP_OK;
}

// Start passive sniffing
esp_err_t scanner_start_passive_sniff(uint32_t duration_seconds, const char *pcap_file) {
    if (s_capture_running) {
        return ESP_ERR_INVALID_STATE;
    }

    s_capture_running = true;
    if (xTaskCreate(packet_capture_task, "packet_capture", 4096, NULL, 5, &s_capture_task) != pdPASS) {
        s_capture_running = false;
        return ESP_ERR_NO_MEM;
    }

    ESP_LOGI(TAG, "Passive sniffing started for %u seconds", duration_seconds);
    return ESP_OK;
}

// Stop passive sniffing
esp_err_t scanner_stop_passive_sniff(void) {
    if (!s_capture_running) {
        return ESP_ERR_INVALID_STATE;
    }

    s_capture_running = false;
    if (s_capture_task) {
        vTaskDelete(s_capture_task);
        s_capture_task = NULL;
    }

    ESP_LOGI(TAG, "Passive sniffing stopped");
    return ESP_OK;
}

// Get discovered devices
esp_err_t scanner_engine_get_devices(discovered_device_t *devices, size_t max_devices, size_t *device_count) {
    if (!devices || !device_count) {
        return ESP_ERR_INVALID_ARG;
    }

    if (xSemaphoreTake(s_scanner_mutex, pdMS_TO_TICKS(1000)) != pdTRUE) {
        return ESP_ERR_TIMEOUT;
    }

    size_t count = (s_device_count < max_devices) ? s_device_count : max_devices;
    memcpy(devices, s_devices, count * sizeof(discovered_device_t));
    *device_count = count;

    xSemaphoreGive(s_scanner_mutex);
    return ESP_OK;
}

// Clear discovered devices
esp_err_t scanner_engine_clear_devices(void) {
    if (xSemaphoreTake(s_scanner_mutex, pdMS_TO_TICKS(1000)) != pdTRUE) {
        return ESP_ERR_TIMEOUT;
    }

    memset(s_devices, 0, sizeof(s_devices));
    s_device_count = 0;

    xSemaphoreGive(s_scanner_mutex);
    ESP_LOGI(TAG, "Device list cleared");
    return ESP_OK;
}

// Get scanner statistics
esp_err_t scanner_engine_get_stats(scanner_stats_t *stats) {
    if (!stats) {
        return ESP_ERR_INVALID_ARG;
    }

    if (xSemaphoreTake(s_scanner_mutex, pdMS_TO_TICKS(1000)) != pdTRUE) {
        return ESP_ERR_TIMEOUT;
    }

    memcpy(stats, &s_stats, sizeof(scanner_stats_t));

    xSemaphoreGive(s_scanner_mutex);
    return ESP_OK;
}

// Reset scanner statistics
esp_err_t scanner_engine_reset_stats(void) {
    if (xSemaphoreTake(s_scanner_mutex, pdMS_TO_TICKS(1000)) != pdTRUE) {
        return ESP_ERR_TIMEOUT;
    }

    memset(&s_stats, 0, sizeof(s_stats));

    xSemaphoreGive(s_scanner_mutex);
    ESP_LOGI(TAG, "Scanner statistics reset");
    return ESP_OK;
}

// Load MAC vendor database
esp_err_t load_mac_vendor_database(void) {
    FILE *f = fopen("/sdcard/config/mac_vendors.csv", "r");
    if (f == NULL) {
        ESP_LOGE(TAG, "Failed to open MAC vendor database file");
        return ESP_FAIL;
    }

    char line[128];
    while (fgets(line, sizeof(line), f)) {
        if (s_mac_vendor_count >= MAX_MAC_VENDORS) {
            ESP_LOGW(TAG, "MAC vendor database is full");
            break;
        }

        char *prefix_str = strtok(line, ",");
        char *vendor_str = strtok(NULL, "\n");

        if (prefix_str && vendor_str) {
            sscanf(prefix_str, "%hhx:%hhx:%hhx",
                   &s_mac_vendors[s_mac_vendor_count].prefix[0],
                   &s_mac_vendors[s_mac_vendor_count].prefix[1],
                   &s_mac_vendors[s_mac_vendor_count].prefix[2]);
            strncpy(s_mac_vendors[s_mac_vendor_count].vendor, vendor_str,
                    sizeof(s_mac_vendors[s_mac_vendor_count].vendor) - 1);
            s_mac_vendor_count++;
        }
    }

    fclose(f);
    ESP_LOGI(TAG, "Loaded %d MAC vendor entries", s_mac_vendor_count);
    return ESP_OK;
}

// Perform ARP scan
static esp_err_t perform_arp_scan(const scan_config_t *config) {
    if (!config) {
        return ESP_ERR_INVALID_ARG;
    }

    if (xSemaphoreTake(s_scanner_mutex, pdMS_TO_TICKS(1000)) != pdTRUE) {
        return ESP_ERR_TIMEOUT;
    }

    ESP_LOGI(TAG, "Starting ARP scan from " IPSTR " to " IPSTR, IP2STR(&config->start_ip), IP2STR(&config->end_ip));

    s_stats.total_scans++;
    s_stats.scan_time_ms = esp_log_timestamp();

    ip4_addr_t start_ip = config->start_ip;
    ip4_addr_t end_ip = config->end_ip;

    for (uint32_t ip = ntohl(start_ip.addr); ip <= ntohl(end_ip.addr); ip++) {
        ip4_addr_t target_ip;
        target_ip.addr = htonl(ip);

        // Simplified ARP scan - just add the IP without MAC resolution
        // In a real implementation, you would use raw sockets to capture ARP responses
        add_discovered_device(&target_ip, NULL, NULL);

        vTaskDelay(pdMS_TO_TICKS(config->delay_ms));
    }

    s_stats.scan_time_ms = esp_log_timestamp() - s_stats.scan_time_ms;
    s_stats.successful_scans++;

    xSemaphoreGive(s_scanner_mutex);

    ESP_LOGI(TAG, "ARP scan completed in %lu ms", s_stats.scan_time_ms);
    return ESP_OK;
}

// Perform TCP port scan
static esp_err_t perform_tcp_port_scan(const scan_config_t *config) {
    if (!config) {
        return ESP_ERR_INVALID_ARG;
    }

    if (xSemaphoreTake(s_scanner_mutex, pdMS_TO_TICKS(1000)) != pdTRUE) {
        return ESP_ERR_TIMEOUT;
    }

    ESP_LOGI(TAG, "Starting TCP port scan from port %d to %d",
             config->port_start, config->port_end);

    s_stats.total_scans++;
    s_stats.scan_time_ms = esp_log_timestamp();

    for (size_t i = 0; i < s_device_count; i++) {
        if (!s_devices[i].is_active) {
            continue;
        }

        for (uint16_t port = config->port_start; port <= config->port_end; port++) {
            int sock = socket(AF_INET, SOCK_STREAM, 0);
            if (sock < 0) {
                ESP_LOGE(TAG, "Failed to create socket");
                continue;
            }

            struct sockaddr_in dest_addr;
            dest_addr.sin_addr.s_addr = s_devices[i].ip_addr.addr;
            dest_addr.sin_family = AF_INET;
            dest_addr.sin_port = htons(port);

            struct timeval timeout;
            timeout.tv_sec = config->timeout_ms / 1000;
            timeout.tv_usec = (config->timeout_ms % 1000) * 1000;
            setsockopt(sock, SOL_SOCKET, SO_RCVTIMEO, &timeout, sizeof(timeout));

            if (connect(sock, (struct sockaddr *)&dest_addr, sizeof(dest_addr)) == 0) {
                if (s_devices[i].open_port_count < 255) {
                    s_devices[i].open_ports[s_devices[i].open_port_count++] = port;
                    ESP_LOGI(TAG, "Port %d is open on " IPSTR, port, IP2STR(&s_devices[i].ip_addr));
                }
            }

            close(sock);
            s_stats.ports_scanned++;
            vTaskDelay(pdMS_TO_TICKS(config->delay_ms));
        }
    }

    s_stats.scan_time_ms = esp_log_timestamp() - s_stats.scan_time_ms;
    s_stats.successful_scans++;

    xSemaphoreGive(s_scanner_mutex);

    ESP_LOGI(TAG, "TCP port scan completed in %lu ms", s_stats.scan_time_ms);
    return ESP_OK;
}

// Perform ICMP ping sweep
static esp_err_t perform_icmp_ping_sweep(const scan_config_t *config) {
    if (!config) {
        return ESP_ERR_INVALID_ARG;
    }

    if (xSemaphoreTake(s_scanner_mutex, pdMS_TO_TICKS(1000)) != pdTRUE) {
        return ESP_ERR_TIMEOUT;
    }

    ESP_LOGI(TAG, "Starting ICMP ping sweep from " IPSTR " to " IPSTR, IP2STR(&config->start_ip), IP2STR(&config->end_ip));

    s_stats.total_scans++;
    s_stats.scan_time_ms = esp_log_timestamp();

    ip4_addr_t start_ip = config->start_ip;
    ip4_addr_t end_ip = config->end_ip;

    for (uint32_t ip = ntohl(start_ip.addr); ip <= ntohl(end_ip.addr); ip++) {
        ip4_addr_t target_ip;
        target_ip.addr = htonl(ip);

        // This is a simplified ICMP ping. For a more robust implementation,
        // you would use raw sockets and construct the ICMP packets manually.
        // For now, we'll just log the attempt.
        ESP_LOGI(TAG, "Pinging " IPSTR, IP2STR(&target_ip));
        s_stats.packets_sent++;

        vTaskDelay(pdMS_TO_TICKS(config->delay_ms));
    }

    s_stats.scan_time_ms = esp_log_timestamp() - s_stats.scan_time_ms;
    s_stats.successful_scans++;

    xSemaphoreGive(s_scanner_mutex);

    ESP_LOGI(TAG, "ICMP ping sweep completed in %lu ms", s_stats.scan_time_ms);
    return ESP_OK;
}

// Add discovered device
static esp_err_t add_discovered_device(const ip4_addr_t *ip, const uint8_t *mac, const char *hostname) {
    if (!ip || !mac) {
        return ESP_ERR_INVALID_ARG;
    }

    if (s_device_count >= MAX_DEVICES) {
        return ESP_ERR_NO_MEM;
    }

    // Check if device already exists
    for (size_t i = 0; i < s_device_count; i++) {
        if (memcmp(s_devices[i].mac_addr, mac, 6) == 0) {
            s_devices[i].last_seen = esp_log_timestamp();
            return ESP_OK;
        }
    }

    // Add new device
    discovered_device_t *device = &s_devices[s_device_count];
    device->ip_addr = *ip;
    memcpy(device->mac_addr, mac, 6);
    device->last_seen = esp_log_timestamp();
    device->is_active = true;

    if (hostname) {
        strncpy(device->hostname, hostname, sizeof(device->hostname) - 1);
    } else {
        resolve_hostname(ip, device->hostname, sizeof(device->hostname));
    }

    // Lookup MAC vendor
    for (size_t i = 0; i < s_mac_vendor_count; i++) {
        if (memcmp(s_mac_vendors[i].prefix, mac, 3) == 0) {
            strncpy(device->vendor, s_mac_vendors[i].vendor, sizeof(device->vendor) - 1);
            break;
        }
    }

    s_device_count++;
    s_stats.devices_discovered++;

    ESP_LOGI(TAG, "Discovered device: " IPSTR " [%02X:%02X:%02X:%02X:%02X:%02X]", IP2STR(ip), mac[0], mac[1], mac[2], mac[3], mac[4], mac[5]);

    return ESP_OK;
}

// Resolve hostname
static esp_err_t resolve_hostname(const ip4_addr_t *ip, char *hostname, size_t len) {
    if (!ip || !hostname || len == 0) {
        return ESP_ERR_INVALID_ARG;
    }

    // Simplified hostname resolution - just use IP address as hostname
    // In a real implementation, you would use DNS reverse lookup
    snprintf(hostname, len, IPSTR, IP2STR(ip));
    return ESP_OK;
}

// Scanner engine queue processing (stub implementation)
void scanner_engine_process_queue(void) {
    ESP_LOGI(TAG, "Scanner queue processing not implemented");
}

// Start scan (stub implementation)
esp_err_t scanner_engine_start_scan(const scan_config_t *config, uint32_t *scan_id) {
    if (!config || !scan_id) {
        return ESP_ERR_INVALID_ARG;
    }
    
    ESP_LOGI(TAG, "Scanner engine start scan not implemented");
    *scan_id = 1; // Dummy scan ID
    return ESP_OK;
}

// Stop scan (stub implementation)
esp_err_t scanner_engine_stop_scan(uint32_t scan_id) {
    ESP_LOGI(TAG, "Scanner engine stop scan not implemented");
        return ESP_OK;
}

// Get scan status (stub implementation)
esp_err_t scanner_engine_get_scan_status(uint32_t scan_id, scan_request_t *request) {
    if (!request) {
        return ESP_ERR_INVALID_ARG;
    }
    
    ESP_LOGI(TAG, "Scanner engine get scan status not implemented");
    request->scan_id = scan_id;
    request->status = SCAN_STATUS_COMPLETED;
    return ESP_OK;
}
