#include "wifi.h"

#include <netinet/in.h>

#include "esp_mac.h"
#include "esp_wifi.h"
#include "esp_netif_types.h"
#include "esp_err.h"

#define CONFIG_SYS_LOG_LEVEL SYS_LOG_LEVEL_INF
#define SYS_LOG_DOMAIN "WIFI"
#include "sys_log.h"

static struct
{
    uint8_t init;
    uint8_t curr_mode; // 0: 已关闭, 1: ap, 2: sta, 3: ap+sta
    wifi_netif_mode_t cfg_mode;
    wifi_init_t init_param;
    esp_netif_t *netif;
    uint8_t mac[6];
    union
    {
        wifi_ap_connect_status_t ap_status;
        wifi_sta_connect_status_t sta_status;
    };
} s_cm;

static void _wifi_start_ap(void);
static void _wifi_start_sta(void);
static void _wiri_stop_ap(void);
static void _wiri_stop_sta(void);
static void _wifi_event_handler_ap(void *arg, esp_event_base_t event_base, int32_t event_id, void *event_data);  // 由 _wifi_start_ap() 启动时设置的事件回调函数，处理 AP 模式下客户端的连接状态
static void _wifi_event_handler_sta(void *arg, esp_event_base_t event_base, int32_t event_id, void *event_data); // 由 _wifi_start_sta() 启动时设置的事件回调函数，处理 STA 模式下的连接状态

/**
 * @brief 初始化设置。可重复设置
 *
 * @param init_struct
 */
void wifi_netif_init(wifi_init_t *init_struct)
{
    if (s_cm.init == 0)
    {
        s_cm.init = 1;
        ESP_ERROR_CHECK(esp_netif_init());
    }
    s_cm.init_param = *init_struct;
}

/**
 * @brief 设置 WIFI 模式。在下次执行 wifi_start() 时根据此模式启动 WIFI
 *
 * @param mode @ref wifi_netif_mode_t
 */
void wifi_set_mode(wifi_netif_mode_t mode)
{
    s_cm.cfg_mode = mode;
}

void wifi_start(void)
{
    switch (s_cm.cfg_mode)
    {
    case WIFI_NETIF_MODE_AP:
        _wifi_start_ap();
        break;

    case WIFI_NETIF_MODE_STA:
    default:
        _wifi_start_sta();
        break;
    }
}

/**
 * @brief 关闭 WIFI
 *
 */
void wifi_stop(void)
{
    if (s_cm.curr_mode == 0)
    {
        SYS_LOG_WRN("wifi has been stoped");
        return;
    }

    if (s_cm.curr_mode == 1)
    {
        _wiri_stop_ap();
    }
    else if (s_cm.curr_mode == 2)
    {
        _wiri_stop_sta();
    }

    s_cm.curr_mode = 0;
}

/**
 * @brief （仅在 AP 模式下）获取已连接的工作站信息
 *
 * @param netif_sta_list
 * @return int
 */
int wifi_get_sta_list(wifi_sta_mac_ip_list_t *netif_sta_list)
{
    wifi_sta_list_t sta_list;
    memset(&sta_list, 0, sizeof(sta_list));
    memset(netif_sta_list, 0, sizeof(wifi_sta_mac_ip_list_t));
    esp_err_t err = esp_wifi_ap_get_sta_list(&sta_list);
    if (err != ESP_OK)
    {
        return err;
    }

    return esp_wifi_ap_get_sta_list_with_ip(&sta_list, netif_sta_list);
}

/**
 * @brief （仅在 AP 模式下）获取已连接的 AP 信息
 *
 * @param netif_ap[out]
 * @return int
 */
int wifi_get_ap_info(esp_netif_ip_info_t *netif_ap)
{
    if (s_cm.curr_mode == 2)
    {
        return esp_netif_get_ip_info(s_cm.netif, netif_ap);
    }
    return -1;
}

/**
 * @brief 初始化 WIFI 协议栈，启动 WIFI 的 AP 模式，启动 DHCP 服务
 *
 */
static void _wifi_start_ap(void)
{
    if (s_cm.curr_mode == 1)
    {
        SYS_LOG_WRN("wifi has been started");
        return;
    }

    if (s_cm.curr_mode)
    {
        wifi_stop();
    }

    s_cm.curr_mode = 1;

    // Initialize networking stack
    ESP_ERROR_CHECK(esp_event_loop_create_default());

    // start the DHCP server
    if (s_cm.netif == NULL)
    {
        s_cm.netif = esp_netif_create_default_wifi_ap();
        assert(s_cm.netif);
        // stop DHCP server
        ESP_ERROR_CHECK(esp_netif_dhcps_stop(s_cm.netif));
        // assign a static IP to the network interface
        esp_netif_ip_info_t info;
        memset(&info, 0, sizeof(info));
        IP4_ADDR(&info.ip, s_cm.init_param.ap.ip_v4[0], s_cm.init_param.ap.ip_v4[1], s_cm.init_param.ap.ip_v4[2], s_cm.init_param.ap.ip_v4[3]);
        IP4_ADDR(&info.gw, s_cm.init_param.ap.gw_v4[0], s_cm.init_param.ap.gw_v4[1], s_cm.init_param.ap.gw_v4[2], s_cm.init_param.ap.gw_v4[3]); // ESP acts as router, so gw addr will be its own addr
        IP4_ADDR(&info.netmask, s_cm.init_param.ap.mask_v4[0], s_cm.init_param.ap.mask_v4[1], s_cm.init_param.ap.mask_v4[2], s_cm.init_param.ap.mask_v4[3]);
        ESP_ERROR_CHECK(esp_netif_set_ip_info(s_cm.netif, &info));
        ESP_ERROR_CHECK(esp_netif_dhcps_start(s_cm.netif));
    }

    /* Initialise ESP32 in SoftAP curr_mode */
    const char *ssid = s_cm.init_param.ap.ssid;
    const char *password = s_cm.init_param.ap.password;
    wifi_config_t wifi_config;
    memset(&wifi_config, 0, sizeof(wifi_config_t));
    memcpy(wifi_config.ap.ssid, ssid, strlen(ssid));
    wifi_config.ap.ssid_len = strlen(ssid);
    wifi_config.ap.max_connection = s_cm.init_param.ap.max_connection;
    if (strlen(password) == 0)
    {
        wifi_config.ap.authmode = WIFI_AUTH_OPEN;
    }
    else
    {
        strcpy((char *)wifi_config.ap.password, password);
        wifi_config.ap.authmode = WIFI_AUTH_WPA_WPA2_PSK;
    }

    wifi_init_config_t cfg = WIFI_INIT_CONFIG_DEFAULT();
    ESP_ERROR_CHECK(esp_wifi_init(&cfg));
    ESP_ERROR_CHECK(esp_event_handler_register(WIFI_EVENT, ESP_EVENT_ANY_ID, _wifi_event_handler_ap, NULL));
    ESP_ERROR_CHECK(esp_wifi_set_mode(WIFI_MODE_AP));
    ESP_ERROR_CHECK(esp_wifi_set_config(ESP_IF_WIFI_AP, &wifi_config));
    ESP_ERROR_CHECK(esp_wifi_start());

    wifi_ap_connect_cb connect_cb = s_cm.init_param.ap.connect_cb;
    if (connect_cb)
    {
        s_cm.ap_status = WIFI_AP_CONNECT_GOT_STA;
        connect_cb(s_cm.ap_status, 0);
    }

    SYS_LOG_INF("done");
}

/**
 * @brief 初始化 WIFI 协议栈，启动 WIFI 的 STA 模式
 *
 */
static void _wifi_start_sta(void)
{
    if (s_cm.curr_mode == 2)
    {
        SYS_LOG_WRN("wifi has been started");
        return;
    }

    if (s_cm.curr_mode)
    {
        wifi_stop();
    }

    s_cm.curr_mode = 2;

    /* 参考自 examples/wifi/fast_scan/main/fast_scan.c */
    ESP_ERROR_CHECK(esp_event_loop_create_default());

    // Initialize default station as network interface instance (esp-netif)
    if (s_cm.netif == NULL)
    {
        s_cm.netif = esp_netif_create_default_wifi_sta();
        assert(s_cm.netif);
        ESP_ERROR_CHECK(esp_event_handler_instance_register(WIFI_EVENT, ESP_EVENT_ANY_ID, &_wifi_event_handler_sta, s_cm.netif, NULL));
        ESP_ERROR_CHECK(esp_event_handler_instance_register(IP_EVENT, IP_EVENT_STA_GOT_IP, &_wifi_event_handler_sta, s_cm.netif, NULL));
    }

    // Initialize and start WiFi
    const char *ssid = s_cm.init_param.sta.ssid;
    const char *password = s_cm.init_param.sta.password;
    wifi_config_t wifi_config;
    memset(&wifi_config, 0, sizeof(wifi_config_t));
    strncpy((char *)wifi_config.sta.ssid, ssid, sizeof(wifi_config.sta.ssid));
    strncpy((char *)wifi_config.sta.password, password, sizeof(wifi_config.sta.password));
    wifi_config.sta.scan_method = WIFI_FAST_SCAN;
    wifi_config.sta.sort_method = WIFI_CONNECT_AP_BY_SIGNAL;
    wifi_config.sta.threshold.rssi = -127;
    wifi_config.sta.threshold.authmode = WIFI_AUTH_OPEN;

    wifi_init_config_t cfg = WIFI_INIT_CONFIG_DEFAULT();
    ESP_ERROR_CHECK(esp_wifi_init(&cfg));
    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());

    wifi_sta_connect_cb connect_cb = s_cm.init_param.sta.connect_cb;
    if (connect_cb)
    {
        uint8_t ip_v4[4];
        memset(ip_v4, 0, sizeof(ip_v4));
        s_cm.sta_status = WIFI_STA_CONNECT_CONNECTING;
        connect_cb(s_cm.sta_status, ip_v4);
    }

    SYS_LOG_INF("done");
}

static void _wiri_stop_ap(void)
{
    /* Deinitialise ESP32 netif */
    ESP_ERROR_CHECK(esp_wifi_stop());
    if (s_cm.netif)
    {
        esp_netif_dhcp_status_t status;
        ESP_ERROR_CHECK(esp_netif_dhcps_get_status(s_cm.netif, &status));
        if (ESP_NETIF_DHCP_STARTED == status)
        {
            ESP_ERROR_CHECK(esp_netif_dhcps_stop(s_cm.netif));
        }
        esp_netif_destroy_default_wifi(s_cm.netif);
        s_cm.netif = NULL;
    }
    ESP_ERROR_CHECK(esp_wifi_deinit());

    // Deinitialize networking stack
    ESP_ERROR_CHECK(esp_event_loop_delete_default());

    wifi_ap_connect_cb connect_cb = s_cm.init_param.ap.connect_cb;
    if (connect_cb)
    {
        s_cm.ap_status = WIFI_AP_CONNECT_STOPED;
        connect_cb(s_cm.ap_status, 0);
    }

    SYS_LOG_INF("done");

    s_cm.curr_mode = 0;
}

static void _wiri_stop_sta(void)
{
    /* Deinitialise ESP32 netif */
    ESP_ERROR_CHECK(esp_wifi_stop());
    if (s_cm.netif)
    {
        esp_netif_destroy_default_wifi(s_cm.netif);
        s_cm.netif = NULL;
    }
    ESP_ERROR_CHECK(esp_wifi_deinit());

    // Deinitialize networking stack
    ESP_ERROR_CHECK(esp_event_loop_delete_default());

    wifi_sta_connect_cb connect_cb = s_cm.init_param.sta.connect_cb;
    if (connect_cb)
    {
        uint8_t ip_v4[4];
        memset(ip_v4, 0, sizeof(ip_v4));
        s_cm.sta_status = WIFI_STA_CONNECT_STOPED;
        connect_cb(s_cm.sta_status, ip_v4);
    }

    SYS_LOG_INF("done");

    s_cm.curr_mode = 0;
}

static void _wifi_event_handler_ap(void *arg, esp_event_base_t event_base, int32_t event_id, void *event_data)
{
    switch (event_id)
    {
    case WIFI_EVENT_AP_STACONNECTED:
    {
        wifi_event_ap_staconnected_t *event = (wifi_event_ap_staconnected_t *)event_data;
        SYS_LOG_INF("station " MACSTR " join, AID=%d",
                    MAC2STR(event->mac), event->aid);
        memcpy(s_cm.mac, event->mac, 6);
        break;
    }

    case WIFI_EVENT_AP_STADISCONNECTED:
    {
        wifi_event_ap_stadisconnected_t *event = (wifi_event_ap_stadisconnected_t *)event_data;
        SYS_LOG_INF("station " MACSTR " leave, AID=%d",
                    MAC2STR(event->mac), event->aid);
        break;
    }

    default:
        return;
    }

    wifi_ap_connect_cb connect_cb = s_cm.init_param.ap.connect_cb;
    if (connect_cb)
    {
        wifi_sta_mac_ip_list_t netif_sta_list;
        wifi_get_sta_list(&netif_sta_list);
        s_cm.ap_status = WIFI_AP_CONNECT_GOT_STA;
        connect_cb(s_cm.ap_status, netif_sta_list.num);
    }
}

void get_wifi_mac(uint8_t *mac)
{
    memcpy(mac, s_cm.mac, 6);
}

static void _wifi_event_handler_sta(void *arg, esp_event_base_t event_base, int32_t event_id, void *event_data)
{
    wifi_sta_connect_status_t status;
    if (event_base == WIFI_EVENT && event_id == WIFI_EVENT_STA_START)
    {
        status = WIFI_STA_CONNECT_CONNECTING;
        esp_wifi_connect();
    }
    else if (event_base == WIFI_EVENT && event_id == WIFI_EVENT_STA_DISCONNECTED)
    {
        status = WIFI_STA_CONNECT_CONNECTING;
        esp_wifi_connect();
    }
    else if (event_base == IP_EVENT && event_id == IP_EVENT_STA_GOT_IP)
    {
        status = WIFI_STA_CONNECT_GOT_IP;
        ip_event_got_ip_t *event = (ip_event_got_ip_t *)event_data;
        SYS_LOG_INF("got ip: " IPSTR, IP2STR(&event->ip_info.ip));
    }
    else
    {
        return;
    }

    wifi_sta_connect_cb connect_cb = s_cm.init_param.sta.connect_cb;
    if (connect_cb)
    {
        uint8_t ip_v4[4];
        memset(ip_v4, 0, sizeof(ip_v4));
        if (s_cm.sta_status != status)
        {
            s_cm.sta_status = status;
            connect_cb(s_cm.sta_status, ip_v4);
        }
    }
}