// main.c — Add support for wifi_cfg and ip address assignment

#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_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"
#include "wifi_cfg.h"

static const char *TAG = "main";

static void event_handler(void* arg, esp_event_base_t event_base,
                          int32_t event_id, void* event_data)
{
    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));
    }
}

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;

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;
    }

    if (iperf_args.abort->count != 0) {
        iperf_stop();
        return 0;
    }

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

    cfg.flag |= (iperf_args.udp->count == 0) ? IPERF_FLAG_TCP : IPERF_FLAG_UDP;

    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 == 0) {
            ESP_LOGE(TAG, "Invalid IP address: %s", iperf_args.ip->sval[0]);
            return 1;
        }
    }

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

    cfg.bw_lim = 0;
    if (iperf_args.bw->count != 0) {
        cfg.bw_lim = iperf_args.bw->ival[0];
    }

    if (iperf_args.interval->count != 0) {
        cfg.interval = iperf_args.interval->ival[0];
    }

    iperf_start(&cfg);

    return 0;
}

static void register_iperf(void)
{
    iperf_args.ip       = arg_str0("c", "client_ip", "<ip>", "Destination IP for client mode");
    iperf_args.server   = arg_lit0("s", "server", "Run as server");
    iperf_args.udp      = arg_lit0("u", "udp", "Use UDP (default TCP)");
    iperf_args.client   = arg_lit0("c", "client", "Run as client");
    iperf_args.port     = arg_int0("p", "port", "<n>", "Port (default 5001)");
    iperf_args.interval = arg_int0("i", "interval", "<s>", "Report interval");
    iperf_args.time     = arg_int0("t", "time", "<s>", "Test duration");
    iperf_args.bw       = arg_int0("b", "bandwidth", "<kbps>", "UDP bandwidth");
    iperf_args.abort    = arg_lit0(NULL, "abort", "Abort iperf");
    iperf_args.end      = arg_end(2);

    const esp_console_cmd_t cmd = {
        .command = "iperf",
        .help    = "iperf traffic test (server/client, TCP/UDP)",
        .hint    = NULL,
        .func    = &cmd_iperf,
        .argtable= &iperf_args
    };
    ESP_ERROR_CHECK(esp_console_cmd_register(&cmd));
}

void app_main(void)
{
    ESP_ERROR_CHECK(nvs_flash_init());
    ESP_ERROR_CHECK(esp_netif_init());
    ESP_ERROR_CHECK(esp_event_loop_create_default());

    ESP_ERROR_CHECK(esp_event_handler_instance_register(
        IP_EVENT, IP_EVENT_STA_GOT_IP, &event_handler, NULL, NULL));

    wifi_cfg_init();
    wifi_cfg_apply_from_nvs();

    register_iperf();

    ESP_LOGI(TAG, "System init complete. Waiting for serial config or console commands.");
}