/*
 * rt_config.c
 * NVS-backed storage for known WiFi networks and C2 fallback addresses.
 */
#include "sdkconfig.h"
#include "rt_config.h"

#ifdef CONFIG_MODULE_REDTEAM

#include <string.h>
#include <stdio.h>

#include "nvs_flash.h"
#include "nvs.h"
#include "esp_log.h"
#include "esp_wifi.h"

static const char *TAG = "RT_CFG";
static const char *NVS_NS = "rt_cfg";

/* ============================================================
 * Init
 * ============================================================ */
void rt_config_init(void)
{
    nvs_handle_t h;
    esp_err_t err = nvs_open(NVS_NS, NVS_READWRITE, &h);
    if (err == ESP_OK) {
        nvs_close(h);
        ESP_LOGI(TAG, "NVS namespace '%s' ready", NVS_NS);
    } else {
        ESP_LOGE(TAG, "NVS open failed: %s", esp_err_to_name(err));
    }
}

/* ============================================================
 * Known WiFi networks
 * ============================================================ */

static void net_key_ssid(int idx, char *out, size_t len)
{
    snprintf(out, len, "n_%d", idx);
}

static void net_key_pass(int idx, char *out, size_t len)
{
    snprintf(out, len, "p_%d", idx);
}

int rt_config_net_count(void)
{
    nvs_handle_t h;
    if (nvs_open(NVS_NS, NVS_READONLY, &h) != ESP_OK)
        return 0;

    int32_t count = 0;
    nvs_get_i32(h, "rt_count", &count);
    nvs_close(h);
    return (int)count;
}

int rt_config_net_list(rt_network_t *out, int max_count)
{
    nvs_handle_t h;
    if (nvs_open(NVS_NS, NVS_READONLY, &h) != ESP_OK)
        return 0;

    int32_t count = 0;
    nvs_get_i32(h, "rt_count", &count);
    if (count > max_count) count = max_count;
    if (count > CONFIG_RT_MAX_KNOWN_NETWORKS) count = CONFIG_RT_MAX_KNOWN_NETWORKS;

    char key[16];
    for (int i = 0; i < count; i++) {
        memset(&out[i], 0, sizeof(rt_network_t));

        net_key_ssid(i, key, sizeof(key));
        size_t len = RT_SSID_MAX_LEN;
        nvs_get_str(h, key, out[i].ssid, &len);

        net_key_pass(i, key, sizeof(key));
        len = RT_PASS_MAX_LEN;
        nvs_get_str(h, key, out[i].pass, &len);
    }

    nvs_close(h);
    return (int)count;
}

bool rt_config_net_add(const char *ssid, const char *pass)
{
    if (!ssid || !ssid[0]) return false;

    nvs_handle_t h;
    if (nvs_open(NVS_NS, NVS_READWRITE, &h) != ESP_OK)
        return false;

    int32_t count = 0;
    nvs_get_i32(h, "rt_count", &count);

    /* Check if SSID already exists → update */
    char key[16];
    for (int i = 0; i < count; i++) {
        net_key_ssid(i, key, sizeof(key));
        char existing[RT_SSID_MAX_LEN] = {0};
        size_t len = RT_SSID_MAX_LEN;
        if (nvs_get_str(h, key, existing, &len) == ESP_OK) {
            if (strcmp(existing, ssid) == 0) {
                /* Update password */
                net_key_pass(i, key, sizeof(key));
                nvs_set_str(h, key, pass ? pass : "");
                nvs_commit(h);
                nvs_close(h);
                ESP_LOGI(TAG, "Updated network '%s'", ssid);
                return true;
            }
        }
    }

    /* New entry */
    if (count >= CONFIG_RT_MAX_KNOWN_NETWORKS) {
        nvs_close(h);
        ESP_LOGW(TAG, "Known networks full (%d)", (int)count);
        return false;
    }

    net_key_ssid(count, key, sizeof(key));
    nvs_set_str(h, key, ssid);

    net_key_pass(count, key, sizeof(key));
    nvs_set_str(h, key, pass ? pass : "");

    count++;
    nvs_set_i32(h, "rt_count", count);
    nvs_commit(h);
    nvs_close(h);

    ESP_LOGI(TAG, "Added network '%s' (total: %d)", ssid, (int)count);
    return true;
}

bool rt_config_net_remove(const char *ssid)
{
    if (!ssid || !ssid[0]) return false;

    nvs_handle_t h;
    if (nvs_open(NVS_NS, NVS_READWRITE, &h) != ESP_OK)
        return false;

    int32_t count = 0;
    nvs_get_i32(h, "rt_count", &count);

    int found = -1;
    char key[16];
    for (int i = 0; i < count; i++) {
        net_key_ssid(i, key, sizeof(key));
        char existing[RT_SSID_MAX_LEN] = {0};
        size_t len = RT_SSID_MAX_LEN;
        if (nvs_get_str(h, key, existing, &len) == ESP_OK) {
            if (strcmp(existing, ssid) == 0) {
                found = i;
                break;
            }
        }
    }

    if (found < 0) {
        nvs_close(h);
        return false;
    }

    /* Shift entries down to fill the gap */
    for (int i = found; i < count - 1; i++) {
        char src_key[16], dst_key[16];
        char buf[RT_PASS_MAX_LEN];
        size_t len;

        /* Copy SSID[i+1] → SSID[i] */
        net_key_ssid(i + 1, src_key, sizeof(src_key));
        net_key_ssid(i, dst_key, sizeof(dst_key));
        len = RT_SSID_MAX_LEN;
        memset(buf, 0, sizeof(buf));
        nvs_get_str(h, src_key, buf, &len);
        nvs_set_str(h, dst_key, buf);

        /* Copy PASS[i+1] → PASS[i] */
        net_key_pass(i + 1, src_key, sizeof(src_key));
        net_key_pass(i, dst_key, sizeof(dst_key));
        len = RT_PASS_MAX_LEN;
        memset(buf, 0, sizeof(buf));
        nvs_get_str(h, src_key, buf, &len);
        nvs_set_str(h, dst_key, buf);
    }

    /* Erase last entries — reuse key[16] from above */
    net_key_ssid(count - 1, key, sizeof(key));
    nvs_erase_key(h, key);
    net_key_pass(count - 1, key, sizeof(key));
    nvs_erase_key(h, key);

    count--;
    nvs_set_i32(h, "rt_count", count);
    nvs_commit(h);
    nvs_close(h);

    ESP_LOGI(TAG, "Removed network '%s' (total: %d)", ssid, (int)count);
    return true;
}

/* ============================================================
 * C2 fallback addresses
 * ============================================================ */

int rt_config_c2_count(void)
{
    nvs_handle_t h;
    if (nvs_open(NVS_NS, NVS_READONLY, &h) != ESP_OK)
        return 0;

    int32_t count = 0;
    nvs_get_i32(h, "c2_count", &count);
    nvs_close(h);
    return (int)count;
}

int rt_config_c2_list(rt_c2_addr_t *out, int max_count)
{
    nvs_handle_t h;
    if (nvs_open(NVS_NS, NVS_READONLY, &h) != ESP_OK)
        return 0;

    int32_t count = 0;
    nvs_get_i32(h, "c2_count", &count);
    if (count > max_count) count = max_count;
    if (count > CONFIG_RT_MAX_C2_FALLBACKS) count = CONFIG_RT_MAX_C2_FALLBACKS;

    for (int i = 0; i < count; i++) {
        memset(&out[i], 0, sizeof(rt_c2_addr_t));
        char key[16];
        snprintf(key, sizeof(key), "c2_%d", i);
        size_t len = RT_ADDR_MAX_LEN;
        nvs_get_str(h, key, out[i].addr, &len);
    }

    nvs_close(h);
    return (int)count;
}

bool rt_config_c2_add(const char *addr)
{
    if (!addr || !addr[0]) return false;

    nvs_handle_t h;
    if (nvs_open(NVS_NS, NVS_READWRITE, &h) != ESP_OK)
        return false;

    int32_t count = 0;
    nvs_get_i32(h, "c2_count", &count);

    /* Check duplicate */
    for (int i = 0; i < count; i++) {
        char key[16];
        snprintf(key, sizeof(key), "c2_%d", i);
        char existing[RT_ADDR_MAX_LEN] = {0};
        size_t len = RT_ADDR_MAX_LEN;
        if (nvs_get_str(h, key, existing, &len) == ESP_OK) {
            if (strcmp(existing, addr) == 0) {
                nvs_close(h);
                return true;  /* Already exists */
            }
        }
    }

    if (count >= CONFIG_RT_MAX_C2_FALLBACKS) {
        nvs_close(h);
        ESP_LOGW(TAG, "C2 fallbacks full (%d)", (int)count);
        return false;
    }

    char key[16];
    snprintf(key, sizeof(key), "c2_%d", (int)count);
    nvs_set_str(h, key, addr);

    count++;
    nvs_set_i32(h, "c2_count", count);
    nvs_commit(h);
    nvs_close(h);

    ESP_LOGI(TAG, "Added C2 fallback '%s' (total: %d)", addr, (int)count);
    return true;
}

bool rt_config_c2_remove(const char *addr)
{
    if (!addr || !addr[0]) return false;

    nvs_handle_t h;
    if (nvs_open(NVS_NS, NVS_READWRITE, &h) != ESP_OK)
        return false;

    int32_t count = 0;
    nvs_get_i32(h, "c2_count", &count);

    int found = -1;
    for (int i = 0; i < count; i++) {
        char key[16];
        snprintf(key, sizeof(key), "c2_%d", i);
        char existing[RT_ADDR_MAX_LEN] = {0};
        size_t len = RT_ADDR_MAX_LEN;
        if (nvs_get_str(h, key, existing, &len) == ESP_OK) {
            if (strcmp(existing, addr) == 0) {
                found = i;
                break;
            }
        }
    }

    if (found < 0) {
        nvs_close(h);
        return false;
    }

    /* Shift down */
    for (int i = found; i < count - 1; i++) {
        char src_key[16], dst_key[16], buf[RT_ADDR_MAX_LEN];
        size_t len = RT_ADDR_MAX_LEN;
        snprintf(src_key, sizeof(src_key), "c2_%d", i + 1);
        snprintf(dst_key, sizeof(dst_key), "c2_%d", i);
        memset(buf, 0, sizeof(buf));
        nvs_get_str(h, src_key, buf, &len);
        nvs_set_str(h, dst_key, buf);
    }

    char key[16];
    snprintf(key, sizeof(key), "c2_%d", (int)(count - 1));
    nvs_erase_key(h, key);

    count--;
    nvs_set_i32(h, "c2_count", count);
    nvs_commit(h);
    nvs_close(h);

    ESP_LOGI(TAG, "Removed C2 fallback '%s' (total: %d)", addr, (int)count);
    return true;
}

/* ============================================================
 * Original MAC storage
 * ============================================================ */

void rt_config_save_orig_mac(void)
{
    uint8_t mac[6];
    if (esp_wifi_get_mac(WIFI_IF_STA, mac) != ESP_OK) {
        ESP_LOGW(TAG, "Failed to read STA MAC");
        return;
    }

    nvs_handle_t h;
    if (nvs_open(NVS_NS, NVS_READWRITE, &h) != ESP_OK)
        return;

    nvs_set_blob(h, "orig_mac", mac, 6);
    nvs_commit(h);
    nvs_close(h);

    ESP_LOGI(TAG, "Saved original MAC: %02X:%02X:%02X:%02X:%02X:%02X",
             mac[0], mac[1], mac[2], mac[3], mac[4], mac[5]);
}

bool rt_config_get_orig_mac(uint8_t mac[6])
{
    nvs_handle_t h;
    if (nvs_open(NVS_NS, NVS_READONLY, &h) != ESP_OK)
        return false;

    size_t len = 6;
    esp_err_t err = nvs_get_blob(h, "orig_mac", mac, &len);
    nvs_close(h);
    return (err == ESP_OK && len == 6);
}

#endif /* CONFIG_MODULE_REDTEAM */