#include <stdio.h>
#include <string.h>
#include <inttypes.h>
#include "freertos/FreeRTOS.h"
#include "freertos/task.h"
#include "freertos/event_groups.h"
#include "esp_system.h"
#include "esp_wifi.h"
#include "esp_event.h"
#include "esp_log.h"
#include "nvs_flash.h"
#include "esp_console.h"
#include "argtable3/argtable3.h"
#include "esp_netif.h"
#include "lwip/inet.h"
#include "iperf.h"

#define WIFI_SSID      CONFIG_WIFI_SSID
#define WIFI_PASS      CONFIG_WIFI_PASSWORD
#define WIFI_MAXIMUM_RETRY  CONFIG_WIFI_MAXIMUM_RETRY

static const char *TAG = "main";

static EventGroupHandle_t s_wifi_event_group;
#define WIFI_CONNECTED_BIT BIT0
#define WIFI_FAIL_BIT      BIT1

static int s_retry_num = 0;

static void 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_START) {
        esp_wifi_connect();
    } else if (event_base == WIFI_EVENT && event_id == WIFI_EVENT_STA_CONNECTED) {
        wifi_event_sta_connected_t* event = (wifi_event_sta_connected_t*) event_data;
        
        // Get RSSI
        wifi_ap_record_t ap_info;
        esp_wifi_sta_get_ap_info(&ap_info);
        
        ESP_LOGI(TAG, "Connected to AP SSID:%s, BSSID:" MACSTR " Channel:%d RSSI:%d dBm",
                 event->ssid, MAC2STR(event->bssid), event->channel, ap_info.rssi);
    } else if (event_base == WIFI_EVENT && event_id == WIFI_EVENT_STA_DISCONNECTED) {
        if (s_retry_num < WIFI_MAXIMUM_RETRY) {
            esp_wifi_connect();
            s_retry_num++;
            ESP_LOGI(TAG, "retry to connect to the AP");
        } else {
            xEventGroupSetBits(s_wifi_event_group, WIFI_FAIL_BIT);
        }
        ESP_LOGI(TAG,"connect to the AP fail");
    } else if (event_base == IP_EVENT && event_id == IP_EVENT_STA_GOT_IP) {
        ip_event_got_ip_t* event = (ip_event_got_ip_t*) event_data;
        ESP_LOGI(TAG, "got ip:" IPSTR " gw:" IPSTR " netmask:" IPSTR, 
                 IP2STR(&event->ip_info.ip),
                 IP2STR(&event->ip_info.gw),
                 IP2STR(&event->ip_info.netmask));
        s_retry_num = 0;
        xEventGroupSetBits(s_wifi_event_group, WIFI_CONNECTED_BIT);
    }
}

void wifi_init_sta(void)
{
    s_wifi_event_group = xEventGroupCreate();

    ESP_ERROR_CHECK(esp_netif_init());
    ESP_ERROR_CHECK(esp_event_loop_create_default());
    
    esp_netif_t *sta_netif = esp_netif_create_default_wifi_sta();

#ifdef CONFIG_USE_STATIC_IP
    // Configure static IP
    ESP_LOGI(TAG, "Configuring static IP: %s", CONFIG_STATIC_IP_ADDR);
    
    esp_netif_dhcpc_stop(sta_netif);
    
    esp_netif_ip_info_t ip_info;
    memset(&ip_info, 0, sizeof(esp_netif_ip_info_t));
    
    ip_info.ip.addr = ipaddr_addr(CONFIG_STATIC_IP_ADDR);
    ip_info.gw.addr = ipaddr_addr(CONFIG_STATIC_GATEWAY_ADDR);
    ip_info.netmask.addr = ipaddr_addr(CONFIG_STATIC_NETMASK_ADDR);
    
    ESP_ERROR_CHECK(esp_netif_set_ip_info(sta_netif, &ip_info));
#endif

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

    esp_event_handler_instance_t instance_any_id;
    esp_event_handler_instance_t instance_got_ip;
    ESP_ERROR_CHECK(esp_event_handler_instance_register(WIFI_EVENT,
                                                        ESP_EVENT_ANY_ID,
                                                        &event_handler,
                                                        NULL,
                                                        &instance_any_id));
    ESP_ERROR_CHECK(esp_event_handler_instance_register(IP_EVENT,
                                                        IP_EVENT_STA_GOT_IP,
                                                        &event_handler,
                                                        NULL,
                                                        &instance_got_ip));

    wifi_config_t wifi_config = {
        .sta = {
            .ssid = WIFI_SSID,
            .password = WIFI_PASS,
            .threshold.authmode = WIFI_AUTH_WPA2_PSK,
        },
    };
    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_LOGI(TAG, "wifi_init_sta finished.");

    EventBits_t bits = xEventGroupWaitBits(s_wifi_event_group,
            WIFI_CONNECTED_BIT | WIFI_FAIL_BIT,
            pdFALSE,
            pdFALSE,
            portMAX_DELAY);

    if (bits & WIFI_CONNECTED_BIT) {
        ESP_LOGI(TAG, "WiFi connection successful");
    } else if (bits & WIFI_FAIL_BIT) {
        ESP_LOGI(TAG, "Failed to connect to SSID:%s", WIFI_SSID);
    } else {
        ESP_LOGE(TAG, "UNEXPECTED EVENT");
    }
}

/* Console command structure for iperf */
static struct {
    struct arg_str *ip;
    struct arg_lit *server;
    struct arg_lit *udp;
    struct arg_lit *client;
    struct arg_int *port;
    struct arg_int *interval;
    struct arg_int *time;
    struct arg_int *bw;
    struct arg_lit *abort;
    struct arg_end *end;
} iperf_args;

/* Console command structure for wifi config */
static struct {
    struct arg_str *ssid;
    struct arg_str *password;
    struct arg_end *end;
} wifi_args;

static int cmd_wifi(int argc, char **argv)
{
    int nerrors = arg_parse(argc, argv, (void **)&wifi_args);
    if (nerrors != 0) {
        arg_print_errors(stderr, wifi_args.end, argv[0]);
        return 1;
    }

    if (wifi_args.ssid->count == 0) {
        ESP_LOGE(TAG, "Please provide SSID with -s option");
        return 1;
    }

    wifi_config_t wifi_config;
    memset(&wifi_config, 0, sizeof(wifi_config_t));

    // Set SSID
    strncpy((char *)wifi_config.sta.ssid, wifi_args.ssid->sval[0], sizeof(wifi_config.sta.ssid) - 1);

    // Set password if provided
    if (wifi_args.password->count > 0) {
        strncpy((char *)wifi_config.sta.password, wifi_args.password->sval[0], sizeof(wifi_config.sta.password) - 1);
        wifi_config.sta.threshold.authmode = WIFI_AUTH_WPA2_PSK;
    }

    ESP_LOGI(TAG, "Updating WiFi configuration...");
    ESP_LOGI(TAG, "SSID: %s", wifi_config.sta.ssid);

    // Stop WiFi
    esp_wifi_disconnect();
    esp_wifi_stop();

    // Update config
    ESP_ERROR_CHECK(esp_wifi_set_config(WIFI_IF_STA, &wifi_config));

    // Restart WiFi
    ESP_ERROR_CHECK(esp_wifi_start());
    esp_wifi_connect();

    ESP_LOGI(TAG, "WiFi configuration updated. Attempting to connect...");

    return 0;
}

static void register_wifi(void)
{
    wifi_args.ssid = arg_str1("s", "ssid", "<ssid>", "WiFi SSID");
    wifi_args.password = arg_str0("p", "password", "<password>", "WiFi password");
    wifi_args.end = arg_end(1);

    const esp_console_cmd_t wifi_cmd = {
        .command = "wifi",
        .help = "Configure WiFi credentials (wifi -s <ssid> -p <password>)",
        .hint = NULL,
        .func = &cmd_wifi,
        .argtable = &wifi_args
    };

    ESP_ERROR_CHECK(esp_console_cmd_register(&wifi_cmd));
}

static int cmd_iperf(int argc, char **argv)
{
    int nerrors = arg_parse(argc, argv, (void **)&iperf_args);
    if (nerrors != 0) {
        arg_print_errors(stderr, iperf_args.end, argv[0]);
        return 1;
    }

    /* Check if needs to abort */
    if (iperf_args.abort->count != 0) {
        iperf_stop();
        return 0;
    }

    iperf_cfg_t cfg;
    memset(&cfg, 0, sizeof(cfg));

    /* Set protocol type */
    cfg.flag |= (iperf_args.udp->count == 0) ? IPERF_FLAG_TCP : IPERF_FLAG_UDP;

    /* Set server/client mode */
    if (iperf_args.server->count != 0) {
        cfg.flag |= IPERF_FLAG_SERVER;
        cfg.sport = IPERF_DEFAULT_PORT;
        if (iperf_args.port->count != 0) {
            cfg.sport = iperf_args.port->ival[0];
        }
    } else {
        cfg.flag |= IPERF_FLAG_CLIENT;
        cfg.dport = IPERF_DEFAULT_PORT;
        if (iperf_args.port->count != 0) {
            cfg.dport = iperf_args.port->ival[0];
        }
        
        if (iperf_args.ip->count == 0) {
            ESP_LOGE(TAG, "Please input destination IP address");
            return 1;
        }
        cfg.dip = ipaddr_addr(iperf_args.ip->sval[0]);
        if (cfg.dip == IPADDR_NONE) {
            ESP_LOGE(TAG, "Invalid IP address: %s", iperf_args.ip->sval[0]);
            return 1;
        }
    }

    /* Set test time */
    cfg.time = IPERF_DEFAULT_TIME;
    if (iperf_args.time->count != 0) {
        cfg.time = iperf_args.time->ival[0];
    }

    /* Set bandwidth limit for UDP */
    cfg.bw_lim = 0;
    if (iperf_args.bw->count != 0) {
        cfg.bw_lim = iperf_args.bw->ival[0];
    }

    ESP_LOGI(TAG, "mode=%s-%s sport=%u, dip=%u.%u.%u.%u:%u, interval=%" PRIu32 ", time=%" PRIu32,
             (cfg.flag & IPERF_FLAG_TCP) ? "tcp" : "udp",
             (cfg.flag & IPERF_FLAG_SERVER) ? "server" : "client",
             cfg.sport,
             cfg.dip & 0xFF, (cfg.dip >> 8) & 0xFF, (cfg.dip >> 16) & 0xFF, (cfg.dip >> 24) & 0xFF, cfg.dport,
             cfg.interval, cfg.time);

    iperf_start(&cfg);
    return 0;
}

static void register_iperf(void)
{
    iperf_args.ip = arg_str0("c", "client", "<ip>", "run in client mode, connecting to <host>");
    iperf_args.server = arg_lit0("s", "server", "run in server mode");
    iperf_args.udp = arg_lit0("u", "udp", "use UDP rather than TCP");
    iperf_args.port = arg_int0("p", "port", "<port>", "server port to listen on/connect to");
    iperf_args.interval = arg_int0("i", "interval", "<interval>", "seconds between periodic bandwidth reports");
    iperf_args.time = arg_int0("t", "time", "<time>", "time in seconds to transmit for (default 30 secs)");
    iperf_args.bw = arg_int0("b", "bandwidth", "<bandwidth>", "bandwidth to send at in Mbits/sec");
    iperf_args.abort = arg_lit0("a", "abort", "abort running iperf");
    iperf_args.end = arg_end(1);

    const esp_console_cmd_t iperf_cmd = {
        .command = "iperf",
        .help = "iperf command",
        .hint = NULL,
        .func = &cmd_iperf,
        .argtable = &iperf_args
    };

    ESP_ERROR_CHECK(esp_console_cmd_register(&iperf_cmd));
}

static void initialize_console(void)
{
    esp_console_repl_t *repl = NULL;
    esp_console_repl_config_t repl_config = ESP_CONSOLE_REPL_CONFIG_DEFAULT();
    repl_config.prompt = "iperf>";

    esp_console_dev_uart_config_t uart_config = ESP_CONSOLE_DEV_UART_CONFIG_DEFAULT();
    ESP_ERROR_CHECK(esp_console_new_repl_uart(&uart_config, &repl_config, &repl));

    register_iperf();
    register_wifi();

    ESP_ERROR_CHECK(esp_console_start_repl(repl));
}

void app_main(void)
{
    ESP_LOGI(TAG, "ESP32 iperf starting...");

    // Initialize NVS
    esp_err_t ret = nvs_flash_init();
    if (ret == ESP_ERR_NVS_NO_FREE_PAGES || ret == ESP_ERR_NVS_NEW_VERSION_FOUND) {
        ESP_ERROR_CHECK(nvs_flash_erase());
        ret = nvs_flash_init();
    }
    ESP_ERROR_CHECK(ret);

    // Initialize WiFi
    wifi_init_sta();

    // Initialize console
    initialize_console();

    ESP_LOGI(TAG, "Ready. Type 'iperf --help' for usage");
}