#include "wifi_monitor.h"
#include "esp_log.h"
#include "esp_wifi.h"
#include "string.h"

static const char *TAG = "WiFi_Monitor";

// Module state
static bool monitor_running = false;
static wifi_monitor_cb_t user_callback = NULL;
static wifi_collapse_stats_t stats = {0};

// Tunable thresholds (accessible via GDB for runtime adjustment)
uint32_t threshold_high_nav_us = 5000;           // NAV threshold for "high" classification
uint32_t threshold_duration_mismatch_us = 10000; // Log mismatches when NAV exceeds this
uint32_t threshold_phy_rate_fallback_mbps = 100; // PHY rate below this = fallback
float threshold_retry_rate_percent = 20.0f;      // Retry rate for collapse detection
uint32_t threshold_avg_nav_collapse_us = 10000;  // Avg NAV threshold for collapse
float threshold_collision_percent = 10.0f;       // Collision event percentage
float threshold_mismatch_percent = 5.0f;         // Duration mismatch percentage
uint32_t threshold_duration_multiplier = 2;      // NAV > expected * this = mismatch

// Logging control
uint32_t log_every_n_mismatches = 1;             // Log every Nth mismatch (1 = all, 10 = every 10th)
static uint32_t s_mismatch_log_counter = 0;

// Forward declarations
static void wifi_promiscuous_rx_cb(void *buf, wifi_promiscuous_pkt_type_t type);

/**
 * @brief Parse 802.11 MAC header
 */
esp_err_t wifi_parse_frame(const uint8_t *payload, uint16_t len, wifi_frame_info_t *frame_info) {
    if (!payload || !frame_info || len < 24) {
        return ESP_ERR_INVALID_ARG;
    }

    memset(frame_info, 0, sizeof(wifi_frame_info_t));

    // Parse Frame Control (bytes 0-1)
    frame_info->frame_control = (payload[1] << 8) | payload[0];
    frame_info->protocol_version = frame_info->frame_control & 0x03;
    frame_info->type = (frame_info->frame_control >> 2) & 0x03;
    frame_info->subtype = (frame_info->frame_control >> 4) & 0x0F;
    frame_info->to_ds = (frame_info->frame_control & FRAME_CTRL_TO_DS) != 0;
    frame_info->from_ds = (frame_info->frame_control & FRAME_CTRL_FROM_DS) != 0;
    frame_info->more_frag = (frame_info->frame_control & FRAME_CTRL_MORE_FRAG) != 0;
    frame_info->retry = (frame_info->frame_control & FRAME_CTRL_RETRY) != 0;
    frame_info->pwr_mgmt = (frame_info->frame_control & FRAME_CTRL_PWR_MGMT) != 0;
    frame_info->more_data = (frame_info->frame_control & FRAME_CTRL_MORE_DATA) != 0;
    frame_info->protected_frame = (frame_info->frame_control & FRAME_CTRL_PROTECTED) != 0;
    frame_info->order = (frame_info->frame_control & FRAME_CTRL_ORDER) != 0;

    // Parse Duration/ID (bytes 2-3) - THIS IS THE NAV FIELD!
    frame_info->duration_id = (payload[3] << 8) | payload[2];

    // Parse Address 1 (Receiver) (bytes 4-9)
    memcpy(frame_info->addr1, &payload[4], 6);

    // Parse Address 2 (Transmitter) (bytes 10-15)
    memcpy(frame_info->addr2, &payload[10], 6);

    // Parse Address 3 (BSSID/SA/DA) (bytes 16-21)
    memcpy(frame_info->addr3, &payload[16], 6);

    // Parse Sequence Control (bytes 22-23)
    frame_info->seq_ctrl = (payload[23] << 8) | payload[22];
    frame_info->fragment_num = frame_info->seq_ctrl & 0x0F;
    frame_info->sequence_num = (frame_info->seq_ctrl >> 4) & 0x0FFF;

    // Check for Address 4 (only present if To DS and From DS both set)
    frame_info->has_addr4 = frame_info->to_ds && frame_info->from_ds;
    if (frame_info->has_addr4 && len >= 30) {
        memcpy(frame_info->addr4, &payload[24], 6);
    }

    frame_info->frame_len = len;

    return ESP_OK;
}

/**
 * @brief Promiscuous mode RX callback
 */
static void wifi_promiscuous_rx_cb(void *buf, wifi_promiscuous_pkt_type_t type) {
    if (!buf) return;

    wifi_promiscuous_pkt_t *pkt = (wifi_promiscuous_pkt_t *)buf;
    wifi_pkt_rx_ctrl_t *rx_ctrl = &pkt->rx_ctrl;
    const uint8_t *payload = pkt->payload;
    uint16_t len = rx_ctrl->sig_len;

    // Parse frame header
    wifi_frame_info_t frame_info;
    if (wifi_parse_frame(payload, len, &frame_info) != ESP_OK) {
        return;
    }

    // Add RX metadata
    frame_info.rssi = rx_ctrl->rssi;
    frame_info.channel = rx_ctrl->channel;
    frame_info.timestamp = rx_ctrl->timestamp;

    // Extract PHY rate info from RX control
    frame_info.mcs = 0;
    frame_info.rate = rx_ctrl->rate;  // This is the rate index
    frame_info.sig_mode = 0;
    frame_info.sgi = false;
    frame_info.bandwidth = 0;

    // Estimate PHY rate from rate index (rough approximation)
    static const uint32_t rate_table[] = {
        1000, 2000, 5500, 11000,  // 1, 2, 5.5, 11 Mbps (DSSS)
        6000, 9000, 12000, 18000, 24000, 36000, 48000, 54000,  // OFDM rates
        65000, 130000, 195000, 260000  // Rough HT estimates
    };

    if (rx_ctrl->rate < sizeof(rate_table) / sizeof(rate_table[0])) {
        frame_info.phy_rate_kbps = rate_table[rx_ctrl->rate];
    } else {
        frame_info.phy_rate_kbps = 100000;  // Assume 100 Mbps default
    }

    // Update statistics
    stats.total_frames++;

    if (frame_info.retry) {
        stats.retry_frames++;
    }

    if (frame_info.duration_id > threshold_high_nav_us) {
        stats.high_nav_frames++;
    }

    if (frame_info.duration_id > stats.max_nav) {
        stats.max_nav = frame_info.duration_id;
    }

    // Track PHY rate statistics
    uint16_t phy_rate_mbps = frame_info.phy_rate_kbps / 1000;
    if (phy_rate_mbps > 0) {
        if (stats.min_phy_rate_mbps == 0 || phy_rate_mbps < stats.min_phy_rate_mbps) {
            stats.min_phy_rate_mbps = phy_rate_mbps;
        }
        if (phy_rate_mbps > stats.max_phy_rate_mbps) {
            stats.max_phy_rate_mbps = phy_rate_mbps;
        }

        stats.avg_phy_rate_mbps = (stats.avg_phy_rate_mbps * (stats.total_frames - 1) + phy_rate_mbps) / stats.total_frames;

        if (frame_info.channel >= 36 && phy_rate_mbps < threshold_phy_rate_fallback_mbps) {
            stats.rate_fallback_frames++;
        }

        // Calculate expected duration
        uint32_t tx_time_us = (frame_info.frame_len * 8000) / frame_info.phy_rate_kbps;
        uint32_t overhead_us = 44;
        if (frame_info.sig_mode == 0) {
            overhead_us = 24;
        } else if (frame_info.bandwidth == 2) {
            overhead_us = 52;
        }
        uint32_t expected_duration = tx_time_us + overhead_us;

        // ---------------------------------------------------------
        // DURATION MISMATCH CHECK (Attacker Identification Added)
        // ---------------------------------------------------------
        if (frame_info.duration_id > expected_duration * threshold_duration_multiplier) {
            stats.duration_mismatch_frames++;

            if (frame_info.duration_id > threshold_duration_mismatch_us) {
                s_mismatch_log_counter++;
                if ((s_mismatch_log_counter % log_every_n_mismatches) == 0) {
                    ESP_LOGW("MONITOR", "Duration mismatch: %s frame, %u bytes @ %u Mbps",
                             wifi_frame_type_str(frame_info.type, frame_info.subtype),
                             frame_info.frame_len, phy_rate_mbps);

                    // NEW: Log the Source MAC (Addr2)
                    ESP_LOGW("MONITOR", "  Source MAC: %02x:%02x:%02x:%02x:%02x:%02x",
                             frame_info.addr2[0], frame_info.addr2[1], frame_info.addr2[2],
                             frame_info.addr2[3], frame_info.addr2[4], frame_info.addr2[5]);

                    ESP_LOGW("MONITOR", "  Expected: %lu us, Actual NAV: %u us (+%ld us)",
                             expected_duration, frame_info.duration_id,
                             frame_info.duration_id - expected_duration);
                    ESP_LOGW("MONITOR", "  Retry: %s, RSSI: %d dBm",
                             frame_info.retry ? "YES" : "no", frame_info.rssi);
                }
            }
        }

        // ---------------------------------------------------------
        // COLLISION CHECK (Attacker Identification Added)
        // ---------------------------------------------------------
        if (frame_info.retry && frame_info.duration_id > threshold_high_nav_us &&
            phy_rate_mbps < threshold_phy_rate_fallback_mbps) {
            ESP_LOGW("MONITOR", "⚠⚠⚠ COLLISION DETECTED!");

            // NEW: Log the Attacker MAC
            ESP_LOGW("MONITOR", "  Attacker MAC: %02x:%02x:%02x:%02x:%02x:%02x",
                     frame_info.addr2[0], frame_info.addr2[1], frame_info.addr2[2],
                     frame_info.addr2[3], frame_info.addr2[4], frame_info.addr2[5]);

            ESP_LOGW("MONITOR", "  Type: %s, Size: %u bytes, Rate: %u Mbps",
                     wifi_frame_type_str(frame_info.type, frame_info.subtype),
                     frame_info.frame_len, phy_rate_mbps);
            ESP_LOGW("MONITOR", "  NAV: %u us (expected %lu us), Retry: YES",
                     frame_info.duration_id, expected_duration);
        }
    }

    // Count frame types
    switch (frame_info.type) {
        case FRAME_TYPE_MANAGEMENT:
            stats.mgmt_frames++;
            break;

        case FRAME_TYPE_CONTROL:
            if (frame_info.subtype == CTRL_RTS) {
                stats.rts_frames++;
            } else if (frame_info.subtype == CTRL_CTS) {
                stats.cts_frames++;
            } else if (frame_info.subtype == CTRL_ACK) {
                stats.ack_frames++;
            }
            break;

        case FRAME_TYPE_DATA:
            stats.data_frames++;
            break;
    }

    if (frame_info.retry && frame_info.duration_id > threshold_high_nav_us) {
        stats.collision_events++;
    }

    if (stats.total_frames > 0) {
        stats.retry_rate = (float)stats.retry_frames / stats.total_frames * 100.0f;
        stats.avg_nav = (stats.avg_nav * (stats.total_frames - 1) + frame_info.duration_id) / stats.total_frames;
    }

    if (user_callback) {
        user_callback(&frame_info, payload, len);
    }
}

/**
 * @brief Initialize WiFi monitor mode
 */
esp_err_t wifi_monitor_init(uint8_t channel, wifi_monitor_cb_t callback) {
    ESP_LOGI(TAG, "Initializing WiFi monitor mode on channel %d", channel);

    user_callback = callback;
    monitor_running = false;

    // Initialize WiFi in NULL mode (no STA or AP)
    wifi_init_config_t cfg = WIFI_INIT_CONFIG_DEFAULT();
    esp_err_t ret = esp_wifi_init(&cfg);
    if (ret != ESP_OK && ret != ESP_ERR_WIFI_NOT_INIT) {
        ESP_LOGE(TAG, "WiFi init failed: %s", esp_err_to_name(ret));
        return ret;
    }

    // Set WiFi mode to NULL (required for promiscuous mode)
    ret = esp_wifi_set_mode(WIFI_MODE_NULL);
    if (ret != ESP_OK) {
        ESP_LOGE(TAG, "Set mode failed: %s", esp_err_to_name(ret));
        return ret;
    }

    // Start WiFi
    ret = esp_wifi_start();
    if (ret != ESP_OK) {
        ESP_LOGE(TAG, "WiFi start failed: %s", esp_err_to_name(ret));
        return ret;
    }

    // Set channel
    ret = wifi_monitor_set_channel(channel);
    if (ret != ESP_OK) {
        ESP_LOGE(TAG, "Set channel failed: %s", esp_err_to_name(ret));
        return ret;
    }

    // Set promiscuous filter to capture all frame types
    wifi_promiscuous_filter_t filter = {
        .filter_mask = WIFI_PROMIS_FILTER_MASK_ALL
    };
    ret = esp_wifi_set_promiscuous_filter(&filter);
    if (ret != ESP_OK) {
        ESP_LOGE(TAG, "Set filter failed: %s", esp_err_to_name(ret));
        return ret;
    }

    // Register promiscuous callback
    ret = esp_wifi_set_promiscuous_rx_cb(wifi_promiscuous_rx_cb);
    if (ret != ESP_OK) {
        ESP_LOGE(TAG, "Set callback failed: %s", esp_err_to_name(ret));
        return ret;
    }

    ESP_LOGI(TAG, "WiFi monitor initialized successfully");
    return ESP_OK;
}

/**
 * @brief Start WiFi monitoring
 */
esp_err_t wifi_monitor_start(void) {
    ESP_LOGI(TAG, "Starting WiFi monitor mode");

    esp_err_t ret = esp_wifi_set_promiscuous(true);
    if (ret != ESP_OK) {
        ESP_LOGE(TAG, "Enable promiscuous failed: %s", esp_err_to_name(ret));
        return ret;
    }

    monitor_running = true;
    wifi_monitor_reset_stats();

    ESP_LOGI(TAG, "WiFi monitor started - capturing all 802.11 frames");
    return ESP_OK;
}

/**
 * @brief Stop WiFi monitoring
 */
esp_err_t wifi_monitor_stop(void) {
    ESP_LOGI(TAG, "Stopping WiFi monitor mode");

    esp_err_t ret = esp_wifi_set_promiscuous(false);
    if (ret != ESP_OK) {
        ESP_LOGE(TAG, "Disable promiscuous failed: %s", esp_err_to_name(ret));
        return ret;
    }

    monitor_running = false;

    ESP_LOGI(TAG, "WiFi monitor stopped");
    return ESP_OK;
}

/**
 * @brief Set WiFi channel for monitoring
 */
esp_err_t wifi_monitor_set_channel(uint8_t channel) {
    ESP_LOGI(TAG, "Setting channel to %d", channel);

    esp_err_t ret = esp_wifi_set_channel(channel, WIFI_SECOND_CHAN_NONE);
    if (ret != ESP_OK) {
        ESP_LOGE(TAG, "Set channel failed: %s", esp_err_to_name(ret));
        return ret;
    }

    return ESP_OK;
}

/**
 * @brief Get WiFi collapse detection statistics
 */
esp_err_t wifi_monitor_get_stats(wifi_collapse_stats_t *out_stats) {
    if (!out_stats) {
        return ESP_ERR_INVALID_ARG;
    }

    memcpy(out_stats, &stats, sizeof(wifi_collapse_stats_t));
    return ESP_OK;
}

/**
 * @brief Reset WiFi collapse detection statistics
 */
void wifi_monitor_reset_stats(void) {
    memset(&stats, 0, sizeof(wifi_collapse_stats_t));
    ESP_LOGI(TAG, "Statistics reset");
}

/**
 * @brief Check if current conditions indicate WiFi collapse
 */
bool wifi_monitor_is_collapsed(void) {
    // Need minimum sample size
    if (stats.total_frames < 100) {
        return false;
    }

    bool high_retry = stats.retry_rate > threshold_retry_rate_percent;
    bool high_nav = stats.avg_nav > threshold_avg_nav_collapse_us;
    bool high_collision = (float)stats.collision_events / stats.total_frames > (threshold_collision_percent / 100.0f);
    bool duration_issues = (float)stats.duration_mismatch_frames / stats.total_frames > (threshold_mismatch_percent / 100.0f);
    bool rate_fallback = stats.avg_phy_rate_mbps < threshold_phy_rate_fallback_mbps;

    // Collapse detected if 3 out of 5 indicators are true
    int indicators = (high_retry ? 1 : 0) +
                     (high_nav ? 1 : 0) +
                     (high_collision ? 1 : 0) +
                     (duration_issues ? 1 : 0) +
                     (rate_fallback ? 1 : 0);

    return indicators >= 3;
}

/**
 * @brief Get string representation of frame type
 */
const char* wifi_frame_type_str(uint8_t type, uint8_t subtype) {
    if (type == FRAME_TYPE_MANAGEMENT) {
        switch (subtype) {
            case MGMT_ASSOC_REQ: return "ASSOC_REQ";
            case MGMT_ASSOC_RESP: return "ASSOC_RESP";
            case MGMT_REASSOC_REQ: return "REASSOC_REQ";
            case MGMT_REASSOC_RESP: return "REASSOC_RESP";
            case MGMT_PROBE_REQ: return "PROBE_REQ";
            case MGMT_PROBE_RESP: return "PROBE_RESP";
            case MGMT_BEACON: return "BEACON";
            case MGMT_ATIM: return "ATIM";
            case MGMT_DISASSOC: return "DISASSOC";
            case MGMT_AUTH: return "AUTH";
            case MGMT_DEAUTH: return "DEAUTH";
            case MGMT_ACTION: return "ACTION";
            default: return "MGMT_UNKNOWN";
        }
    } else if (type == FRAME_TYPE_CONTROL) {
        switch (subtype) {
            case CTRL_BLOCK_ACK_REQ: return "BLOCK_ACK_REQ";
            case CTRL_BLOCK_ACK: return "BLOCK_ACK";
            case CTRL_PS_POLL: return "PS_POLL";
            case CTRL_RTS: return "RTS";
            case CTRL_CTS: return "CTS";
            case CTRL_ACK: return "ACK";
            case CTRL_CF_END: return "CF_END";
            case CTRL_CF_END_ACK: return "CF_END_ACK";
            default: return "CTRL_UNKNOWN";
        }
    } else if (type == FRAME_TYPE_DATA) {
        switch (subtype) {
            case DATA_DATA: return "DATA";
            case DATA_NULL: return "NULL";
            case DATA_QOS_DATA: return "QOS_DATA";
            case DATA_QOS_NULL: return "QOS_NULL";
            default: return "DATA_UNKNOWN";
        }
    }

    return "UNKNOWN";
}