#include "config.h"
#include "frame_parser.h"
#include <stdio.h>
#include <string.h>
#include <arpa/inet.h>

// Radiotap header parsing (simplified)
static int skip_radiotap_header(const uint8_t *packet, size_t len, size_t *offset) {
    if (len < 8) return -1;
    
    // Check for radiotap header (starts with version 0)
    if (packet[0] == 0 && packet[1] == 0) {
        // Radiotap header present
        uint16_t rt_len = *(uint16_t *)&packet[2];
        if (rt_len > len || rt_len < 8) return -1;
        
        // Extract RSSI and rate from radiotap (simplified - full parsing is complex)
        // For now, just skip the header
        *offset = rt_len;
        return 0;
    }
    
    // No radiotap header
    *offset = 0;
    return 0;
}

int parse_80211_frame(const uint8_t *packet, size_t len, wifi_frame_info_t *frame_info) {
    size_t offset = 0;
    
    // Skip radiotap header if present
    if (skip_radiotap_header(packet, len, &offset) < 0) {
        return -1;
    }
    
    if (len < offset + 24) {
        return -1;  // Too short for 802.11 header
    }
    
    const uint8_t *frame = packet + offset;
    
    // Parse Frame Control (bytes 0-1)
    frame_info->frame_control = frame[0] | (frame[1] << 8);
    frame_info->type = (frame[0] >> 2) & 0x03;
    frame_info->subtype = (frame[0] >> 4) & 0x0F;
    frame_info->to_ds = (frame[1] >> 0) & 0x01;
    frame_info->from_ds = (frame[1] >> 1) & 0x01;
    frame_info->retry = (frame[1] >> 3) & 0x01;
    
    // Parse Duration/ID (bytes 2-3)
    frame_info->duration_id = frame[2] | (frame[3] << 8);
    
    // Parse Addresses based on To DS / From DS bits
    // addr1 = RA (Receiver Address)
    // addr2 = TA (Transmitter Address)
    // addr3 = BSSID/SA/DA
    memcpy(frame_info->addr1, &frame[4], 6);   // RA
    memcpy(frame_info->addr2, &frame[10], 6);  // TA
    memcpy(frame_info->addr3, &frame[16], 6);  // BSSID/SA/DA
    
    // Check for Address 4 (only 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) {
        if (len < offset + 30) return -1;
        memcpy(frame_info->addr4, &frame[22], 6);
    }
    
    // Parse Sequence Control (bytes 22-23 or 28-29)
    uint16_t seq_offset = frame_info->has_addr4 ? 28 : 22;
    if (len < offset + seq_offset + 2) return -1;
    
    frame_info->seq_ctrl = frame[seq_offset] | (frame[seq_offset + 1] << 8);
    frame_info->fragment_num = frame_info->seq_ctrl & 0x0F;
    frame_info->sequence_num = (frame_info->seq_ctrl >> 4) & 0x0FFF;
    
    // Initialize PHY info (would need radiotap parsing for accurate values)
    frame_info->rssi = -100;  // Default
    frame_info->mcs = 0;
    frame_info->spatial_streams = 1;
    frame_info->bandwidth = 0;
    frame_info->sgi = false;
    frame_info->phy_rate_kbps = 0;
    
    frame_info->frame_len = len - offset;
    
    return 0;
}

const char *get_frame_type_name(uint8_t type, uint8_t subtype) {
    switch (type) {
        case FRAME_TYPE_MANAGEMENT:
            switch (subtype) {
                case 0: return "ASSOC_REQ";
                case 1: return "ASSOC_RESP";
                case 2: return "REASSOC_REQ";
                case 3: return "REASSOC_RESP";
                case 4: return "PROBE_REQ";
                case 5: return "PROBE_RESP";
                case 8: return "BEACON";
                case 10: return "DISASSOC";
                case 11: return "AUTH";
                case 12: return "DEAUTH";
                default: return "MGMT_UNKNOWN";
            }
        case FRAME_TYPE_CONTROL:
            switch (subtype) {
                case 11: return "RTS";
                case 12: return "CTS";
                case 13: return "ACK";
                default: return "CTRL_UNKNOWN";
            }
        case FRAME_TYPE_DATA:
            switch (subtype) {
                case 0: return "DATA";
                case 1: return "DATA_CF_ACK";
                case 2: return "DATA_CF_POLL";
                case 3: return "DATA_CF_ACK_POLL";
                case 4: return "NULL";
                case 8: return "QOS_DATA";
                case 12: return "QOS_NULL";
                default: return "DATA_UNKNOWN";
            }
        default:
            return "UNKNOWN";
    }
}

void print_frame_info(const wifi_frame_info_t *frame_info, uint64_t timestamp_us) {
    const char *type_name = get_frame_type_name(frame_info->type, frame_info->subtype);
    
    printf("[%llu.%03llu] %s: ", 
           timestamp_us / 1000000, (timestamp_us / 1000) % 1000);
    
    printf("TA=%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]);
    
    printf("RA=%02x:%02x:%02x:%02x:%02x:%02x, ",
           frame_info->addr1[0], frame_info->addr1[1], frame_info->addr1[2],
           frame_info->addr1[3], frame_info->addr1[4], frame_info->addr1[5]);
    
    printf("Size=%u bytes, ", frame_info->frame_len);
    printf("Dur=%u us, ", frame_info->duration_id);
    printf("RSSI=%d dBm", frame_info->rssi);
    
    if (frame_info->retry) {
        printf(", Retry=YES");
    } else {
        printf(", Retry=no");
    }
    
    if (frame_info->mcs > 0) {
        printf(", MCS=%u", frame_info->mcs);
    }
    
    if (frame_info->spatial_streams > 1) {
        printf(", SS=%u", frame_info->spatial_streams);
    }
    
    printf("\n");
}