/*
 * mcs_telemetry.c
 *
 * Copyright (c) 2026 Umber Networks & Robert McMahon
 * All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions are met:
 *
 * 1. Redistributions of source code must retain the above copyright notice,
 * this list of conditions and the following disclaimer.
 *
 * 2. Redistributions in binary form must reproduce the above copyright
 * notice, this list of conditions and the following disclaimer in the
 * documentation and/or other materials provided with the distribution.
 *
 * 3. Neither the name of the copyright holder nor the names of its
 * contributors may be used to endorse or promote products derived from
 * this software without prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
 * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
 * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
 * ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
 * LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
 * CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
 * SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
 * INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
 * CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
 * ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
 * POSSIBILITY OF SUCH DAMAGE.
 */

#include "mcs_telemetry.h"
#include "wifi_monitor.h"
#include "esp_log.h"
#include "esp_timer.h"
#include <string.h>
#include <time.h>

static const char *TAG = "MCS_Telemetry";

static bool s_telemetry_running = false;
static mcs_telemetry_cb_t s_user_callback = NULL;
static mcs_telemetry_stats_t s_stats = {0};

// 802.11ax MCS rates for 20MHz, 800ns GI, 1 Spatial Stream
// MCS 0-11: Standard HE rates
static const uint32_t PHY_RATES_20MHZ_1SS_800NS[] = {
    8650,    // MCS 0: BPSK 1/2
    17200,   // MCS 1: QPSK 1/2
    25800,   // MCS 2: QPSK 3/4
    34400,   // MCS 3: 16-QAM 1/2
    51600,   // MCS 4: 16-QAM 3/4
    68800,   // MCS 5: 64-QAM 2/3
    77400,   // MCS 6: 64-QAM 3/4
    86000,   // MCS 7: 64-QAM 5/6
    103200,  // MCS 8: 256-QAM 3/4
    114700,  // MCS 9: 256-QAM 5/6
    129000,  // MCS 10: 1024-QAM 3/4
    143400   // MCS 11: 1024-QAM 5/6
};

// 802.11ax MCS rates for 20MHz, 400ns GI, 1 Spatial Stream
static const uint32_t PHY_RATES_20MHZ_1SS_400NS[] = {
    9600,    // MCS 0
    19200,   // MCS 1
    28800,   // MCS 2
    38400,   // MCS 3
    57600,   // MCS 4
    76800,   // MCS 5
    86400,   // MCS 6
    96000,   // MCS 7
    115200,  // MCS 8
    128100,  // MCS 9
    144100,  // MCS 10
    160200   // MCS 11
};

// 802.11ax MCS rates for 40MHz, 800ns GI, 1 Spatial Stream
static const uint32_t PHY_RATES_40MHZ_1SS_800NS[] = {
    17200,   // MCS 0
    34400,   // MCS 1
    51600,   // MCS 2
    68800,   // MCS 3
    103200,  // MCS 4
    137600,  // MCS 5
    154900,  // MCS 6
    172100,  // MCS 7
    206500,  // MCS 8
    229400,  // MCS 9
    258100,  // MCS 10
    286800   // MCS 11
};

// 802.11ax MCS rates for 40MHz, 400ns GI, 1 Spatial Stream
static const uint32_t PHY_RATES_40MHZ_1SS_400NS[] = {
    19200,   // MCS 0
    38400,   // MCS 1
    57600,   // MCS 2
    76800,   // MCS 3
    115200,  // MCS 4
    153600,  // MCS 5
    172800,  // MCS 6
    192000,  // MCS 7
    230400,  // MCS 8
    256200,  // MCS 9
    288200,  // MCS 10
    320300   // MCS 11
};

/**
 * @brief Get device index by MAC address, or create new entry
 */
static int mcs_get_device_index(const uint8_t *mac) {
    int empty_slot = -1;

    for (int i = 0; i < MCS_TELEMETRY_MAX_DEVICES; i++) {
        if (memcmp(s_stats.devices[i].mac, mac, 6) == 0) {
            return i;  // Found existing device
        }
        if (empty_slot < 0 && s_stats.devices[i].sample_count == 0) {
            empty_slot = i;  // Found empty slot
        }
    }

    // Use empty slot or create new entry
    if (empty_slot >= 0) {
        memcpy(s_stats.devices[empty_slot].mac, mac, 6);
        memset(&s_stats.devices[empty_slot], 0, sizeof(mcs_device_telemetry_t));
        memcpy(s_stats.devices[empty_slot].mac, mac, 6);
        s_stats.total_devices++;
        return empty_slot;
    }

    // No space - return oldest device (simple round-robin replacement)
    return 0;
}

/**
 * @brief Update device telemetry with new sample
 */
static void mcs_update_device_telemetry(int dev_idx, const mcs_sample_t *sample) {
    if (dev_idx < 0 || dev_idx >= MCS_TELEMETRY_MAX_DEVICES) {
        return;
    }

    mcs_device_telemetry_t *dev = &s_stats.devices[dev_idx];

    // Update sample buffer (sliding window)
    dev->samples[dev->sample_idx] = *sample;
    dev->sample_idx = (dev->sample_idx + 1) % 16;

    // Update counters
    dev->sample_count++;
    dev->total_frames++;
    dev->total_bytes += sample->frame_len;
    dev->last_update_ms = sample->timestamp_ms;

    if (sample->is_retry) {
        dev->retry_frames++;
    }

    // Update MCS distribution
    if (sample->mcs <= MCS_TELEMETRY_MAX_MCS) {
        dev->mcs_count[sample->mcs]++;

        // Update dominant MCS (most frequent in recent samples)
        uint32_t max_count = 0;
        for (int i = 0; i <= MCS_TELEMETRY_MAX_MCS; i++) {
            if (dev->mcs_count[i] > max_count) {
                max_count = dev->mcs_count[i];
                dev->current_mcs = i;
            }
        }
    }

    // Update SS distribution
    if (sample->ss >= 1 && sample->ss <= MCS_TELEMETRY_MAX_SS) {
        dev->ss_count[sample->ss]++;

        // Update dominant SS
        uint32_t max_count = 0;
        for (int i = 1; i <= MCS_TELEMETRY_MAX_SS; i++) {
            if (dev->ss_count[i] > max_count) {
                max_count = dev->ss_count[i];
                dev->current_ss = i;
            }
        }
    }

    // Update RSSI statistics
    if (dev->sample_count == 1) {
        dev->avg_rssi = sample->rssi;
        dev->min_rssi = sample->rssi;
        dev->max_rssi = sample->rssi;
    } else {
        // Running average
        dev->avg_rssi = ((int16_t)dev->avg_rssi * (dev->sample_count - 1) + sample->rssi) / dev->sample_count;
        if (sample->rssi < dev->min_rssi) dev->min_rssi = sample->rssi;
        if (sample->rssi > dev->max_rssi) dev->max_rssi = sample->rssi;
    }

    // Update PHY rate statistics
    uint32_t total_rate = dev->avg_phy_rate_kbps * (dev->sample_count - 1) + sample->phy_rate_kbps;
    dev->avg_phy_rate_kbps = total_rate / dev->sample_count;
    if (sample->phy_rate_kbps > dev->max_phy_rate_kbps) {
        dev->max_phy_rate_kbps = sample->phy_rate_kbps;
    }
}

esp_err_t mcs_telemetry_init(mcs_telemetry_cb_t callback) {
    ESP_LOGI(TAG, "Initializing MCS telemetry");

    s_user_callback = callback;
    s_telemetry_running = false;
    memset(&s_stats, 0, sizeof(mcs_telemetry_stats_t));

    return ESP_OK;
}

esp_err_t mcs_telemetry_start(void) {
    ESP_LOGI(TAG, "Starting MCS telemetry capture");

    s_telemetry_running = true;
    s_stats.window_start_ms = esp_timer_get_time() / 1000;

    return ESP_OK;
}

esp_err_t mcs_telemetry_stop(void) {
    ESP_LOGI(TAG, "Stopping MCS telemetry capture");

    s_telemetry_running = false;
    s_stats.window_end_ms = esp_timer_get_time() / 1000;

    return ESP_OK;
}

esp_err_t mcs_telemetry_process_frame(const wifi_frame_info_t *frame_info, const wifi_pkt_rx_ctrl_t *rx_ctrl) {
    (void)rx_ctrl;  /* Optional: not used, frame_info contains RX metadata */
    if (!s_telemetry_running || !frame_info) {
        return ESP_ERR_INVALID_ARG;
    }

    // Get device index by MAC address (Addr2 = transmitter)
    int dev_idx = mcs_get_device_index(frame_info->addr2);
    if (dev_idx < 0) {
        return ESP_ERR_NO_MEM;
    }

    // Create sample from frame info
    mcs_sample_t sample = {0};
    sample.timestamp_ms = esp_timer_get_time() / 1000;
    sample.mcs = frame_info->mcs;
    sample.ss = 1;  // TODO: Extract from HT/VHT/HE headers
    sample.rssi = frame_info->rssi;
    sample.channel = frame_info->channel;
    sample.bandwidth = (frame_info->bandwidth == 0) ? MCS_BW_20MHZ :
                       (frame_info->bandwidth == 1) ? MCS_BW_40MHZ : MCS_BW_20MHZ;
    sample.frame_len = frame_info->frame_len;
    sample.is_retry = frame_info->retry;
    sample.sig_mode = frame_info->sig_mode;

    // Calculate PHY rate if we have MCS info
    if (sample.mcs <= MCS_TELEMETRY_MAX_MCS && sample.ss >= 1 && sample.ss <= MCS_TELEMETRY_MAX_SS) {
        sample.phy_rate_kbps = mcs_calculate_phy_rate_ax(sample.mcs, sample.ss, sample.bandwidth, frame_info->sgi);
    } else {
        // Fallback to frame_info's calculated rate or estimate from rate index
        sample.phy_rate_kbps = frame_info->phy_rate_kbps;
    }

    // Update device telemetry
    mcs_update_device_telemetry(dev_idx, &sample);
    s_stats.total_frames_captured++;

    // TODO: Parse HT/VHT/HE headers to extract actual SS count
    // This requires parsing the PLCP/HT Control/VHT Control/HE Control fields
    // which follow the MAC header in HT/VHT/HE frames

    return ESP_OK;
}

esp_err_t mcs_telemetry_get_stats(mcs_telemetry_stats_t *stats) {
    if (!stats) {
        return ESP_ERR_INVALID_ARG;
    }

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

    return ESP_OK;
}

void mcs_telemetry_reset(void) {
    ESP_LOGI(TAG, "Resetting telemetry statistics");

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

esp_err_t mcs_telemetry_to_json(char *json_buffer, size_t buffer_len, const char *device_id) {
    if (!json_buffer || buffer_len == 0) {
        return ESP_ERR_INVALID_ARG;
    }

#define MCS_JSON_TAIL_RESERVE 4  /* room for "]\0" and safety */
    if (buffer_len <= MCS_JSON_TAIL_RESERVE) {
        return ESP_ERR_NO_MEM;
    }

    uint32_t now_ms = esp_timer_get_time() / 1000;
    int written = snprintf(json_buffer, buffer_len,
        "{\"device_id\":\"%s\",\"timestamp\":%lu,\"total_frames\":%lu,\"devices\":[",
        device_id ? device_id : "unknown", now_ms, s_stats.total_frames_captured);

    if (written < 0 || (size_t)written >= buffer_len) {
        return ESP_ERR_NO_MEM;
    }

    size_t offset = (size_t)written;
    bool first = true;

    for (int i = 0; i < MCS_TELEMETRY_MAX_DEVICES; i++) {
        mcs_device_telemetry_t *dev = &s_stats.devices[i];
        if (dev->sample_count == 0) continue;

        size_t space = buffer_len - offset - MCS_JSON_TAIL_RESERVE;
        if (space < 2) break;

        if (!first) {
            written = snprintf(json_buffer + offset, space, ",");
            if (written < 0) break;
            offset += (size_t)((written < (int)space) ? written : (space - 1));
            space = buffer_len - offset - MCS_JSON_TAIL_RESERVE;
            if (space < 2) break;
        }
        first = false;

        written = snprintf(json_buffer + offset, space,
            "{\"mac\":\"%02x:%02x:%02x:%02x:%02x:%02x\","
            "\"mcs\":%u,\"ss\":%u,\"rssi\":%d,"
            "\"channel\":%u,\"bandwidth\":%u,"
            "\"frames\":%lu,\"retries\":%lu,"
            "\"phy_rate_kbps\":%lu}",
            dev->mac[0], dev->mac[1], dev->mac[2], dev->mac[3], dev->mac[4], dev->mac[5],
            dev->current_mcs, dev->current_ss, dev->avg_rssi,
            dev->samples[0].channel, dev->samples[0].bandwidth,
            dev->total_frames, dev->retry_frames,
            dev->avg_phy_rate_kbps);

        if (written < 0) break;
        offset += (size_t)((written < (int)space) ? written : (space - 1));
    }

    {
        size_t tail_space = buffer_len - offset;
        if (tail_space < 3) {
            return ESP_ERR_NO_MEM;
        }
        written = snprintf(json_buffer + offset, tail_space, "]}");
        if (written < 0 || (size_t)written >= tail_space) {
            return ESP_ERR_NO_MEM;
        }
    }

    return ESP_OK;
#undef MCS_JSON_TAIL_RESERVE
}

uint32_t mcs_calculate_phy_rate_ax(uint8_t mcs, uint8_t ss, mcs_bandwidth_t bandwidth, bool sgi) {
    if (mcs > MCS_TELEMETRY_MAX_MCS || ss < 1 || ss > MCS_TELEMETRY_MAX_SS) {
        return 0;
    }

    const uint32_t *rate_table = NULL;

    if (bandwidth == MCS_BW_20MHZ) {
        rate_table = sgi ? PHY_RATES_20MHZ_1SS_400NS : PHY_RATES_20MHZ_1SS_800NS;
    } else if (bandwidth == MCS_BW_40MHZ) {
        rate_table = sgi ? PHY_RATES_40MHZ_1SS_400NS : PHY_RATES_40MHZ_1SS_800NS;
    } else {
        return 0;
    }

    // PHY rate = base rate (1SS) * spatial streams
    uint32_t base_rate = rate_table[mcs];
    return base_rate * ss;
}