#include "wifi_manager.h"

#include <string.h>

#include "esp_log.h"
#include "esp_event.h"
#include "esp_netif.h"
#include "esp_wifi.h"

#include "lwip/inet.h"
#include "lwip/ip4_addr.h"

#include "freertos/event_groups.h"

#define WIFI_CONNECTED_BIT BIT0

static const char *TAG = "WIFI_MANAGER";

static EventGroupHandle_t s_wifi_event_group = NULL;
static esp_netif_t *s_sta_netif = NULL;
static bool s_connected = false;

static void wifi_event_handler(void *arg,
                               esp_event_base_t event_base,
                               int32_t event_id,
                               void *event_data)
{
    (void)arg;

    if (event_base == WIFI_EVENT && event_id == WIFI_EVENT_STA_START)
    {
        ESP_LOGI(TAG, "Wi-Fi STA start");
        esp_wifi_connect();
    }
    else if (event_base == WIFI_EVENT && event_id == WIFI_EVENT_STA_DISCONNECTED)
    {
        ESP_LOGW(TAG, "Wi-Fi disconnected, retrying");
        s_connected = false;
        if (s_wifi_event_group != NULL)
            xEventGroupClearBits(s_wifi_event_group, WIFI_CONNECTED_BIT);
        esp_wifi_connect();
    }
    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, IP2STR(&event->ip_info.ip));
        s_connected = true;

        if (s_wifi_event_group != NULL)
            xEventGroupSetBits(s_wifi_event_group, WIFI_CONNECTED_BIT);
    }
}

static esp_err_t wifi_manager_apply_static_ip(const wifi_settings_t *cfg)
{
    esp_err_t err;
    esp_netif_ip_info_t ip_info;
    ip4_addr_t ip, gw, netmask;

    if (cfg == NULL || s_sta_netif == NULL)
        return ESP_ERR_INVALID_ARG;

    if (!ip4addr_aton(cfg->static_ip, &ip) ||
        !ip4addr_aton(cfg->gateway, &gw) ||
        !ip4addr_aton(cfg->netmask, &netmask))
    {
        ESP_LOGE(TAG, "Invalid static IP configuration");
        return ESP_ERR_INVALID_ARG;
    }

    err = esp_netif_dhcpc_stop(s_sta_netif);
    if (err != ESP_OK && err != ESP_ERR_ESP_NETIF_DHCP_ALREADY_STOPPED)
        return err;

    memset(&ip_info, 0, sizeof(ip_info));
    ip_info.ip.addr = ip.addr;
    ip_info.gw.addr = gw.addr;
    ip_info.netmask.addr = netmask.addr;

    err = esp_netif_set_ip_info(s_sta_netif, &ip_info);
    if (err != ESP_OK)
        return err;

    ESP_LOGI(TAG,
             "Static IP configured: ip=%s gw=%s netmask=%s",
             cfg->static_ip,
             cfg->gateway,
             cfg->netmask);

    return ESP_OK;
}

esp_err_t wifi_manager_init(void)
{
    static bool initialized = false;
    wifi_init_config_t wifi_init_cfg = WIFI_INIT_CONFIG_DEFAULT();

    if (initialized)
        return ESP_OK;

    esp_err_t err;

    err = esp_netif_init();
    if (err != ESP_OK && err != ESP_ERR_INVALID_STATE)
        return err;

    err = esp_event_loop_create_default();
    if (err != ESP_OK && err != ESP_ERR_INVALID_STATE)
        return err;

    s_sta_netif = esp_netif_create_default_wifi_sta();
    if (s_sta_netif == NULL)
        return ESP_FAIL;

    s_wifi_event_group = xEventGroupCreate();
    if (s_wifi_event_group == NULL)
        return ESP_ERR_NO_MEM;

    ESP_ERROR_CHECK(esp_wifi_init(&wifi_init_cfg));

    ESP_ERROR_CHECK(esp_event_handler_register(WIFI_EVENT,
                                               ESP_EVENT_ANY_ID,
                                               &wifi_event_handler,
                                               NULL));

    ESP_ERROR_CHECK(esp_event_handler_register(IP_EVENT,
                                               IP_EVENT_STA_GOT_IP,
                                               &wifi_event_handler,
                                               NULL));

    initialized = true;
    return ESP_OK;
}

esp_err_t wifi_manager_start_sta(const wifi_settings_t *cfg)
{
    wifi_config_t wifi_cfg;

    if (cfg == NULL)
        return ESP_ERR_INVALID_ARG;

    memset(&wifi_cfg, 0, sizeof(wifi_cfg));
    strncpy((char *)wifi_cfg.sta.ssid, cfg->ssid, sizeof(wifi_cfg.sta.ssid) - 1U);
    strncpy((char *)wifi_cfg.sta.password, cfg->password, sizeof(wifi_cfg.sta.password) - 1U);

    wifi_cfg.sta.threshold.authmode = WIFI_AUTH_WPA2_PSK;
    wifi_cfg.sta.pmf_cfg.capable = true;
    wifi_cfg.sta.pmf_cfg.required = false;

    ESP_ERROR_CHECK(esp_wifi_set_mode(WIFI_MODE_STA));

    if (!cfg->dhcp)
    {
        esp_err_t err = wifi_manager_apply_static_ip(cfg);
        if (err != ESP_OK)
            return err;
    }
    else
    {
        esp_netif_dhcpc_start(s_sta_netif);
    }

    ESP_ERROR_CHECK(esp_wifi_set_config(WIFI_IF_STA, &wifi_cfg));
    ESP_ERROR_CHECK(esp_wifi_start());

    return ESP_OK;
}

bool wifi_manager_wait_connected(TickType_t timeout_ticks)
{
    EventBits_t bits;

    if (s_wifi_event_group == NULL)
        return false;

    bits = xEventGroupWaitBits(s_wifi_event_group,
                               WIFI_CONNECTED_BIT,
                               pdFALSE,
                               pdFALSE,
                               timeout_ticks);

    return ((bits & WIFI_CONNECTED_BIT) != 0U);
}

bool wifi_manager_is_connected(void)
{
    return s_connected;
}