#include <nvs_flash.h>
#include <string.h>
#include "freertos/FreeRTOS.h"
#include "freertos/task.h"
#include "freertos/event_groups.h"
#include "wifi_conn.h"
#include "esp_wifi.h"
#include "esp_wifi_types.h"
#include "esp_event.h"
#include "esp_log.h"
#include "sntp_cli.h"

#include "application.h"
#include "dhcp_wd.h"

#define WIFI_MAX_RETRY 5

/* FreeRTOS event group to signal when we are connected*/
EventGroupHandle_t g_wifi_event_group;

static const char *TAG = "wifi station";

static int s_retry_num = 0;
static dhcp_wd_handle_t dhcp_wd = NULL;

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_DISCONNECTED) {
        ESP_LOGI(TAG, "STA disconnected");
        xEventGroupClearBits(g_wifi_event_group, WIFI_CONNECTED_BIT);

        if (dhcp_wd) {
            dhcp_watchdog_notify(dhcp_wd, DHCP_WD_NOTIFY_DISCONNECTED);
        }

//        if (s_retry_num < WIFI_MAX_RETRY) {
//            esp_wifi_disconnect();
            esp_wifi_connect(); // eventually, reconnect works, but it takes a while
            s_retry_num++;
            ESP_LOGI(TAG, "Retry to connect (try %d)", s_retry_num+1);
//        } else {
//            xEventGroupSetBits(g_wifi_event_group, WIFI_FAIL_BIT);
//        }
    } else if (event_base == IP_EVENT && event_id == WIFI_EVENT_STA_CONNECTED) {
        if (dhcp_wd) {
            dhcp_watchdog_notify(dhcp_wd, DHCP_WD_NOTIFY_CONNECTED);
        }
    } 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_retry_num = 0;
        xEventGroupSetBits(g_wifi_event_group, WIFI_CONNECTED_BIT);

        if (g_Settings.ntp_enable) {
            sntp_cli_start();
        }

        if (dhcp_wd) {
            dhcp_watchdog_notify(dhcp_wd, DHCP_WD_NOTIFY_GOT_IP);
        }
    }
}

bool wifi_wait_connect()
{
    s_retry_num = 0;

    /* Waiting until either the connection is established (WIFI_CONNECTED_BIT) or connection failed for the maximum
     * number of re-tries (WIFI_FAIL_BIT). The bits are set by event_handler() (see above) */
    EventBits_t bits = xEventGroupWaitBits(g_wifi_event_group,
                                           WIFI_CONNECTED_BIT | WIFI_FAIL_BIT,
                                           pdFALSE,
                                           pdFALSE,
                                           portMAX_DELAY);

    /* xEventGroupWaitBits() returns the bits before the call returned, hence we can test which event actually
     * happened. */
    if (bits & WIFI_CONNECTED_BIT) {
        ESP_LOGI(TAG, "Connected to AP!");
        return true;
    } else if (bits & WIFI_FAIL_BIT) {
        ESP_LOGI(TAG, "Failed to connect.");
        return false;
    } else {
        ESP_LOGE(TAG, "UNEXPECTED EVENT");
        return false;
    }
}

void initialise_wifi(void)
{
    esp_err_t rv = ESP_OK;

    #define ERROR_CHECK(x) do { if (ESP_OK != (rv = (x))) goto fail; } while (0)

    assert(g_Settings.sta_enabled || g_Settings.ap_enabled);

    g_wifi_event_group = xEventGroupCreate();

    ESP_LOGD(TAG, "create wifi netif");
    esp_netif_t *netif_sta = NULL;
    if (g_Settings.sta_enabled) {
        netif_sta = esp_netif_create_default_wifi_sta();
    }
    if (g_Settings.ap_enabled) {
        esp_netif_create_default_wifi_ap();
    }

    ESP_LOGD(TAG, "initing wifi");
    wifi_init_config_t cfg = WIFI_INIT_CONFIG_DEFAULT();
    ERROR_CHECK(esp_wifi_init(&cfg));

    ESP_LOGD(TAG, "set storage, set sta");
    ERROR_CHECK(esp_wifi_set_storage(WIFI_STORAGE_FLASH));

    wifi_mode_t mode;
    if (g_Settings.sta_enabled && g_Settings.ap_enabled) {
        mode = WIFI_MODE_APSTA;
    } else if (g_Settings.sta_enabled) {
        mode = WIFI_MODE_STA;
    } else {
        // This function is never called if both STA and AP are disabled!
        mode = WIFI_MODE_AP;
    }

    ESP_LOGD(TAG, "set wifi mode");
    ERROR_CHECK(esp_wifi_set_mode(mode));

    if (g_Settings.ap_enabled) {
        ESP_LOGD(TAG, "set AP IP config");

        tcpip_adapter_ip_info_t ip_info;
        ERROR_CHECK(tcpip_adapter_dhcps_stop(TCPIP_ADAPTER_IF_AP));

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

        ERROR_CHECK(tcpip_adapter_set_ip_info(TCPIP_ADAPTER_IF_AP, &ip_info));
        ERROR_CHECK(tcpip_adapter_dhcps_start(TCPIP_ADAPTER_IF_AP));
    }

    const bool use_dhcp = (g_Settings.dhcp_enable || g_Settings.static_ip == 0)
            && g_Settings.sta_enabled;

    if (g_Settings.sta_enabled) {
        if (!use_dhcp) {
            ESP_LOGD(TAG, "set STA static IP");

            tcpip_adapter_ip_info_t ip_info;
            ERROR_CHECK(tcpip_adapter_dhcpc_stop(TCPIP_ADAPTER_IF_STA));

            memset(&ip_info, 0, sizeof(ip_info));
            ip_info.ip.addr = g_Settings.static_ip;
            ip_info.gw.addr = g_Settings.static_ip_gw;
            ip_info.netmask.addr = g_Settings.static_ip_mask;
            ERROR_CHECK(tcpip_adapter_set_ip_info(TCPIP_ADAPTER_IF_STA, &ip_info));

            tcpip_adapter_dns_info_t dnsinfo = {
                    // special esp version of ip address because why not I guess.
                    // ipv4 has union tag 0, so it needn't be set
                    .ip.u_addr.ip4.addr = g_Settings.static_dns,
            };

            ERROR_CHECK(tcpip_adapter_set_dns_info(TCPIP_ADAPTER_IF_STA, ESP_NETIF_DNS_MAIN, &dnsinfo));
        }

        ESP_LOGD(TAG, "register WiFi events");
        ERROR_CHECK(esp_event_handler_register(WIFI_EVENT, ESP_EVENT_ANY_ID, &event_handler, NULL));
        ERROR_CHECK(esp_event_handler_register(IP_EVENT, IP_EVENT_STA_GOT_IP, &event_handler, NULL));
    }

    ESP_LOGD(TAG, "wifi start");
    ERROR_CHECK(esp_wifi_start() );

    ESP_LOGI(TAG, "WiFi config done");

    if (g_Settings.dhcp_wd_enable && use_dhcp) {
        ESP_LOGD(TAG, "Start gateway ping watchdog");
        ERROR_CHECK(dhcp_watchdog_start(netif_sta, true, &dhcp_wd));
    }

    return;

fail:
    ESP_LOGE(TAG, "Error setting up WiFi: %d - %s", rv, esp_err_to_name(rv));
}