/*
 * rt_promisc.c
 * Shared promiscuous mode dispatcher.
 *
 * Multiplexes up to RT_PROMISC_MAX_HANDLERS consumers behind a single
 * IRAM callback registered with esp_wifi_set_promiscuous_rx_cb().
 */
#include "sdkconfig.h"

#ifdef CONFIG_MODULE_REDTEAM

#include <string.h>
#include <stdatomic.h>

#include "freertos/FreeRTOS.h"
#include "freertos/semphr.h"
#include "esp_wifi.h"
#include "esp_log.h"

#include "rt_promisc.h"

#define TAG "RT_PROMISC"

/* ============================================================
 * Internal state
 * ============================================================ */

static rt_promisc_handler_t s_handlers[RT_PROMISC_MAX_HANDLERS];
static int                  s_handler_count = 0;
static SemaphoreHandle_t    s_mutex         = NULL;
static atomic_bool          s_enabled       = ATOMIC_VAR_INIT(false);
static bool                 s_inited        = false;

/* ============================================================
 * IRAM dispatcher — called from WiFi driver context
 * ============================================================ */

static void IRAM_ATTR promisc_dispatcher(void *buf, wifi_promiscuous_pkt_type_t type)
{
    /*
     * We iterate without taking the mutex because:
     *  - This runs in IRAM from the WiFi RX ISR context
     *  - s_handler_count is only modified while promiscuous is disabled
     *  - Handlers are only added/removed via register/unregister which
     *    require the caller to disable promiscuous first (or accept races
     *    on the count — benign: worst case a handler is skipped once).
     */
    int n = s_handler_count;
    for (int i = 0; i < n; i++) {
        if (s_handlers[i].cb) {
            s_handlers[i].cb(buf, type);
        }
    }
}

/* ============================================================
 * Public API
 * ============================================================ */

void rt_promisc_init(void)
{
    if (s_inited) return;

    s_mutex = xSemaphoreCreateMutex();
    configASSERT(s_mutex);

    memset(s_handlers, 0, sizeof(s_handlers));
    s_handler_count = 0;
    s_inited = true;

    ESP_LOGI(TAG, "Promiscuous dispatcher initialised (max %d handlers)",
             RT_PROMISC_MAX_HANDLERS);
}

esp_err_t rt_promisc_register(const rt_promisc_handler_t *h)
{
    if (!h || !h->cb) return ESP_ERR_INVALID_ARG;
    if (!s_inited) rt_promisc_init();

    xSemaphoreTake(s_mutex, portMAX_DELAY);

    /* Check for duplicates */
    for (int i = 0; i < s_handler_count; i++) {
        if (s_handlers[i].cb == h->cb) {
            xSemaphoreGive(s_mutex);
            ESP_LOGW(TAG, "Handler '%s' already registered", h->tag ? h->tag : "?");
            return ESP_OK;
        }
    }

    if (s_handler_count >= RT_PROMISC_MAX_HANDLERS) {
        xSemaphoreGive(s_mutex);
        ESP_LOGE(TAG, "Handler table full (%d/%d)", s_handler_count,
                 RT_PROMISC_MAX_HANDLERS);
        return ESP_ERR_NO_MEM;
    }

    s_handlers[s_handler_count] = *h;
    s_handler_count++;

    ESP_LOGI(TAG, "Registered handler '%s' (filter=0x%04"PRIx32") [%d/%d]",
             h->tag ? h->tag : "?", h->filter_mask,
             s_handler_count, RT_PROMISC_MAX_HANDLERS);

    xSemaphoreGive(s_mutex);
    return ESP_OK;
}

esp_err_t rt_promisc_unregister(const rt_promisc_handler_t *h)
{
    if (!h || !h->cb) return ESP_ERR_INVALID_ARG;
    if (!s_inited) return ESP_ERR_INVALID_STATE;

    xSemaphoreTake(s_mutex, portMAX_DELAY);

    for (int i = 0; i < s_handler_count; i++) {
        if (s_handlers[i].cb == h->cb) {
            /* Shift remaining entries down */
            for (int j = i; j < s_handler_count - 1; j++) {
                s_handlers[j] = s_handlers[j + 1];
            }
            s_handler_count--;
            memset(&s_handlers[s_handler_count], 0, sizeof(rt_promisc_handler_t));

            ESP_LOGI(TAG, "Unregistered handler '%s' [%d/%d]",
                     h->tag ? h->tag : "?",
                     s_handler_count, RT_PROMISC_MAX_HANDLERS);

            xSemaphoreGive(s_mutex);
            return ESP_OK;
        }
    }

    xSemaphoreGive(s_mutex);
    ESP_LOGW(TAG, "Handler '%s' not found", h->tag ? h->tag : "?");
    return ESP_ERR_NOT_FOUND;
}

esp_err_t rt_promisc_enable(void)
{
    if (!s_inited) rt_promisc_init();

    xSemaphoreTake(s_mutex, portMAX_DELAY);

    /* Combine filters from all registered handlers */
    uint32_t combined = 0;
    for (int i = 0; i < s_handler_count; i++) {
        combined |= s_handlers[i].filter_mask;
    }

    /* If no handlers registered but caller still wants promisc (e.g. for TX),
     * default to management frames */
    if (combined == 0) {
        combined = WIFI_PROMIS_FILTER_MASK_MGMT;
    }

    xSemaphoreGive(s_mutex);

    /* Set the single dispatcher callback */
    esp_err_t ret = esp_wifi_set_promiscuous_rx_cb(promisc_dispatcher);
    if (ret != ESP_OK) {
        ESP_LOGE(TAG, "set_promiscuous_rx_cb failed: %s", esp_err_to_name(ret));
        return ret;
    }

    wifi_promiscuous_filter_t filter = { .filter_mask = combined };
    esp_wifi_set_promiscuous_filter(&filter);

    ret = esp_wifi_set_promiscuous(true);
    if (ret == ESP_OK) {
        atomic_store(&s_enabled, true);
        ESP_LOGI(TAG, "Promiscuous enabled (filter=0x%04"PRIx32")", combined);
    } else {
        ESP_LOGE(TAG, "set_promiscuous(true) failed: %s", esp_err_to_name(ret));
    }

    return ret;
}

esp_err_t rt_promisc_disable(void)
{
    esp_err_t ret = esp_wifi_set_promiscuous(false);
    if (ret == ESP_OK) {
        atomic_store(&s_enabled, false);
        ESP_LOGI(TAG, "Promiscuous disabled");
    }
    return ret;
}

esp_err_t rt_promisc_set_channel(uint8_t channel)
{
    if (channel < 1 || channel > 14) return ESP_ERR_INVALID_ARG;
    return esp_wifi_set_channel(channel, WIFI_SECOND_CHAN_NONE);
}

bool rt_promisc_is_enabled(void)
{
    return atomic_load(&s_enabled);
}

#endif /* CONFIG_MODULE_REDTEAM */