#include <string.h>
#include "freertos/FreeRTOS.h"
#include "freertos/task.h"
#include "freertos/event_groups.h"
#include "esp_log.h"
#include "esp_check.h"

#include "utils.h"
#include "net.h"

static const char *TAG = "net";

#if !CONFIG_IDF_TARGET_LINUX

#include "esp_wifi.h"
#include "esp_mac.h"
#include "esp_event.h"
#include "esp_netif_net_stack.h"
#include "esp_netif.h"
#include "nvs_flash.h"
#include "lwip/inet.h"
#include "lwip/netdb.h"
#include "lwip/sockets.h"
#if IP_NAPT
#include "lwip/lwip_napt.h"
#endif
#include "lwip/err.h"
#include "lwip/sys.h"
#include "mdns.h"

/* The event group allows multiple bits for each event, but we only care about two events:
 * - we are connected to the AP with an IP
 * - we failed to connect after the maximum amount of retries */
#define WIFI_CONNECTED_BIT BIT0
#define WIFI_FAIL_BIT      BIT1

static const char *TAG_AP = "net SoftAP";
static const char *TAG_STA = "net Sta";

/* FreeRTOS event group to signal when we are connected/disconnected */
static EventGroupHandle_t s_wifi_event_group;

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

    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) {
        if (s_retry_num < 5) {
            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, IP2STR(&event->ip_info.ip));
        s_retry_num = 0;
        xEventGroupSetBits(s_wifi_event_group, WIFI_CONNECTED_BIT);
    }
    ESP_LOGI(TAG_STA, "%s, %li", event_base, event_id);
}

static void setup_mdns()
{
    ESP_ERROR_CHECK(mdns_init());

    char* id;
    asprintf(
        &id,
        "shelly%s-"MACFMT"",
        CONFIG_SHELLSP_MODEL, MAC2STR(get_wifi_mac())
    );

    ESP_LOGI(TAG, "mDNS hostname: %s", id);
    mdns_hostname_set(id);
    mdns_instance_name_set(id); 

    mdns_txt_item_t serviceTxtData[6] = {
        {"discoverable", "true"},
        {"fw_id", "20230913-112003/v1.14.0-gcb84623"},
        {"fw_version", "1.0"},
        {"app", "switch1"},
        {"arch", "esp8266"},
        {"id", id},
    };
    mdns_service_add(NULL, "_http", "_tcp", 80, serviceTxtData, 6);
}

static bool wifi_init_softap(void)
{
    esp_netif_t* netif = esp_netif_create_default_wifi_ap();

    wifi_config_t wifi_ap_config = {
        .ap = {
            .channel = 1,
            .authmode = WIFI_AUTH_OPEN,
            .max_connection = 5,
            .pmf_cfg = {
                .required = false,
            },
        },
    };

    char* ssid = (char*)wifi_ap_config.ap.ssid;
    snprintf(
        ssid, 32,
        "shelly%s-"MACFMT"",
        CONFIG_SHELLSP_MODEL, MAC2STR(get_wifi_mac())
    );
    wifi_ap_config.ap.ssid_len = strlen(ssid);

    ESP_ERROR_CHECK(esp_wifi_set_config(WIFI_IF_AP, &wifi_ap_config));

    // Stop DHCP so we can configure it
    ESP_ERROR_CHECK(esp_netif_dhcps_stop(netif));

    // Setting a static IP for setup
    // Device is expected to be available at 192.168.33.1
    esp_netif_ip_info_t ip = {0};
    ip.ip.addr = ipaddr_addr(CONFIG_SHELLSP_AP_IP_ADDR);
    ip.netmask.addr = ipaddr_addr(CONFIG_SHELLSP_AP_IP_NETMASK);
    ip.gw.addr = ipaddr_addr(CONFIG_SHELLSP_AP_IP_ADDR);
    ESP_ERROR_CHECK(esp_netif_set_ip_info(netif, &ip));

    // Turn the DHCP server back on
    ESP_ERROR_CHECK(esp_netif_dhcps_start(netif));

    ESP_LOGI(TAG_AP, "wifi_init_softap finished. SSID:%s", wifi_ap_config.ap.ssid);

    return true;
}

static bool wifi_init_sta(void)
{
    wifi_config_t wifi_cfg;
    ESP_RETURN_ON_FALSE(esp_wifi_get_config(WIFI_IF_STA, &wifi_cfg) == ESP_OK, false, TAG_STA, "failed to get wifi sta config");

    ESP_LOGI(TAG_STA, "BLA BLA SSID IS: %s", (char*)wifi_cfg.sta.ssid);

    if (!strlen((const char*)wifi_cfg.sta.ssid))
        return false;

    esp_netif_create_default_wifi_sta();

    return true;
}

esp_err_t net_connect(void)
{
    /* Initialize event group */
    s_wifi_event_group = xEventGroupCreate();

    /* Register Event handler */
    ESP_ERROR_CHECK(
        esp_event_handler_instance_register(
            WIFI_EVENT,
            ESP_EVENT_ANY_ID,
            &wifi_event_handler,
            NULL,
            NULL
        )
    );
    ESP_ERROR_CHECK(
        esp_event_handler_instance_register(
            IP_EVENT,
            IP_EVENT_STA_GOT_IP,
            &wifi_event_handler,
            NULL,
            NULL
        )
    );

    /*Initialize WiFi */
    wifi_init_config_t cfg = WIFI_INIT_CONFIG_DEFAULT();
    ESP_ERROR_CHECK(esp_wifi_init(&cfg));

    ESP_ERROR_CHECK(esp_wifi_set_storage(WIFI_STORAGE_FLASH));


    // Set APSTA ahead of time so we can check data
    ESP_ERROR_CHECK(esp_wifi_set_mode(WIFI_MODE_APSTA));

    if (!wifi_init_softap())
    {
        ESP_LOGE(TAG, "Failed to initialize SoftAP");
    }

    bool s = wifi_init_sta();
    if (s)
         ESP_ERROR_CHECK(esp_wifi_set_mode(WIFI_MODE_STA));

    ESP_ERROR_CHECK(esp_wifi_start());

    setup_mdns();

    if (!s)
    {
        ESP_LOGI(TAG, "Invalid STA config, not waiting for STA");
        return ESP_OK;
    }

    /*
     * Wait 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(
        s_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 station, setting STA mode");
        ESP_ERROR_CHECK(esp_wifi_set_mode(WIFI_MODE_STA));
    } else if (bits & WIFI_FAIL_BIT) {
        ESP_LOGI(TAG, "Failed to connect to station");
    } else {
        ESP_LOGE(TAG, "UNEXPECTED EVENT");
        return ESP_FAIL;
    }

    return ESP_OK;
}

#else /* CONFIG_IDF_TARGET_LINUX */

esp_err_t net_connect(void) {
    ESP_LOGI(TAG, "STUB net_connect()");
    return ESP_OK;
}

#endif /* !CONFIG_IDF_TARGET_LINUX */