#include "dpp_wifi.h"

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_ERROR_CHECK(esp_supp_dpp_start_listen());
        ESP_LOGI(WIFI_TAG, "Started listening for DPP Authentication");
    } else if (event_base == WIFI_EVENT && event_id == WIFI_EVENT_STA_DISCONNECTED) {
        if (s_retry_num < WIFI_MAX_RETRY_NUM) {
            esp_wifi_connect();
            s_retry_num++;
            ESP_LOGI(WIFI_TAG, "retry to connect to the AP");
        } else {
            xEventGroupSetBits(s_dpp_event_group, DPP_CONNECT_FAIL_BIT);
        }
        ESP_LOGI(WIFI_TAG, "connect to the AP fail");
    } else if (event_base == WIFI_EVENT && event_id == WIFI_EVENT_STA_CONNECTED) {
	    ESP_LOGI(WIFI_TAG, "Successfully connected to the AP ssid : %s ", s_dpp_wifi_configuration.sta.ssid);
    } 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(WIFI_TAG, "got ip:" IPSTR, IP2STR(&event->ip_info.ip));
        s_retry_num = 0;
        xEventGroupSetBits(s_dpp_event_group, DPP_CONNECTED_BIT);
    }
}

void dpp_enrollee_event_cb(esp_supp_dpp_event_t event, void *data)
{
    switch (event) {
    case ESP_SUPP_DPP_URI_READY:
        if (data != NULL) {
            esp_qrcode_config_t cfg = ESP_QRCODE_CONFIG_DEFAULT();

            ESP_LOGI(WIFI_TAG, "Scan below QR Code to configure the enrollee:");
            esp_qrcode_generate(&cfg, (const char *)data);
        }
        break;
    case ESP_SUPP_DPP_CFG_RECVD:
        memcpy(&s_dpp_wifi_configuration, data, sizeof(s_dpp_wifi_configuration));
        s_retry_num = 0;
        esp_wifi_set_config(ESP_IF_WIFI_STA, &s_dpp_wifi_configuration);
        esp_wifi_connect();
        break;
    case ESP_SUPP_DPP_FAIL:
        if (s_retry_num < 5) {
            ESP_LOGI(WIFI_TAG, "DPP Auth failed (Reason: %s), retry...", esp_err_to_name((int)data));
            ESP_ERROR_CHECK(esp_supp_dpp_start_listen());
            s_retry_num++;
        } else {
            xEventGroupSetBits(s_dpp_event_group, DPP_AUTH_FAIL_BIT);
        }
        break;
    default:
        break;
    }
}

esp_err_t dpp_enrollee_bootstrap(void)
{
    esp_err_t ret;
    size_t pkey_len = strlen(EXAMPLE_DPP_BOOTSTRAPPING_KEY);
    char *key = NULL;

    if (pkey_len) {
        /* Currently only NIST P-256 curve is supported, add prefix/postfix accordingly */
        char prefix[] = "30310201010420";
        char postfix[] = "a00a06082a8648ce3d030107";

        if (pkey_len != CURVE_SEC256R1_PKEY_HEX_DIGITS) {
            ESP_LOGI(WIFI_TAG, "Invalid key length! Private key needs to be 32 bytes (or 64 hex digits) long");
            return ESP_FAIL;
        }

        key = malloc(sizeof(prefix) + pkey_len + sizeof(postfix));
        if (!key) {
            ESP_LOGI(WIFI_TAG, "Failed to allocate for bootstrapping key");
            return ESP_ERR_NO_MEM;
        }
        sprintf(key, "%s%s%s", prefix, EXAMPLE_DPP_BOOTSTRAPPING_KEY, postfix);
    }

    /* Currently only supported method is QR Code */
    ret = esp_supp_dpp_bootstrap_gen(EXAMPLE_DPP_LISTEN_CHANNEL_LIST, DPP_BOOTSTRAP_QR_CODE,
                                     key, EXAMPLE_DPP_DEVICE_INFO);

    if (key)
        free(key);

    return ret;
}

void save_wifi_configuration(const uint8_t *ssid, const uint8_t *psk) {
    nvs_handle_t handle;
    ESP_ERROR_CHECK(nvs_open("wifi_creds", NVS_READWRITE, &handle));

    // Save SSID
    ESP_ERROR_CHECK(nvs_set_str(handle, "ssid", (const char *)ssid));

    // Save PSK
    ESP_ERROR_CHECK(nvs_set_str(handle, "psk", (const char *)psk));

    // Commit the changes
    ESP_ERROR_CHECK(nvs_commit(handle));

    // Close the handle
    nvs_close(handle);
}

bool load_wifi_configuration() {
    nvs_handle_t handle;
    esp_err_t err = nvs_open("wifi_creds", NVS_READONLY, &handle);
    if (err != ESP_OK) {
        // Handle error opening NVS
        return false; // Indicates failure to load credentials
    }

    // Initialize ssid and psk fields in the struct
    memset(&s_dpp_wifi_configuration, 0, sizeof(s_dpp_wifi_configuration));

    // Read SSID
    size_t ssid_size = sizeof(s_dpp_wifi_configuration.sta.ssid);
    if (nvs_get_str(handle, "ssid", (char *)s_dpp_wifi_configuration.sta.ssid, &ssid_size) != ESP_OK || strlen((char *)s_dpp_wifi_configuration.sta.ssid) == 0) {
        nvs_close(handle);
        return false; // Indicates no valid SSID stored
    }

    // Read PSK
    size_t psk_size = sizeof(s_dpp_wifi_configuration.sta.password);
    if (nvs_get_str(handle, "psk", (char *)s_dpp_wifi_configuration.sta.password, &psk_size) != ESP_OK || strlen((char *)s_dpp_wifi_configuration.sta.password) == 0) {
        nvs_close(handle);
        return false; // Indicates no valid PSK stored
    }

    // Close the handle
    nvs_close(handle);
    return true; // Indicates valid credentials loaded
}

void dpp_enrollee_init(void)
{
    s_dpp_event_group = xEventGroupCreate();

    ESP_ERROR_CHECK(esp_netif_init());

    ESP_ERROR_CHECK(esp_event_loop_create_default());
    esp_netif_create_default_wifi_sta();

    ESP_ERROR_CHECK(esp_event_handler_register(WIFI_EVENT, ESP_EVENT_ANY_ID, &event_handler, NULL));
    ESP_ERROR_CHECK(esp_event_handler_register(IP_EVENT, IP_EVENT_STA_GOT_IP, &event_handler, NULL));

    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_supp_dpp_init(dpp_enrollee_event_cb));
    ESP_ERROR_CHECK(dpp_enrollee_bootstrap());
    ESP_ERROR_CHECK(esp_wifi_start());

    /* 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(s_dpp_event_group,
                                           DPP_CONNECTED_BIT | DPP_CONNECT_FAIL_BIT | DPP_AUTH_FAIL_BIT,
                                           pdFALSE,
                                           pdFALSE,
                                           portMAX_DELAY);

    /* xEventGroupWaitBits() returns the bits before the call returned, hence we can test which event actually
     * happened. */
    if (bits & DPP_CONNECTED_BIT) {
        ESP_LOGI(WIFI_TAG, "connected to ap SSID:%s",
                 s_dpp_wifi_configuration.sta.ssid);
        save_wifi_configuration(s_dpp_wifi_configuration.sta.ssid, s_dpp_wifi_configuration.sta.password);
    } else if (bits & DPP_CONNECT_FAIL_BIT) {
        ESP_LOGI(WIFI_TAG, "Failed to connect to SSID:%s",
                 s_dpp_wifi_configuration.sta.ssid);
    } else if (bits & DPP_AUTH_FAIL_BIT) {
        ESP_LOGI(WIFI_TAG, "DPP Authentication failed after %d retries", s_retry_num);
    } else {
        ESP_LOGE(WIFI_TAG, "UNEXPECTED EVENT");
    }

    esp_supp_dpp_deinit();
    ESP_ERROR_CHECK(esp_event_handler_unregister(IP_EVENT, IP_EVENT_STA_GOT_IP, &event_handler));
    //ESP_ERROR_CHECK(esp_event_handler_unregister(WIFI_EVENT, ESP_EVENT_ANY_ID, &event_handler));
    vEventGroupDelete(s_dpp_event_group);
}

void launch_dpp(void)
{
    if (load_wifi_configuration()) {
        ESP_ERROR_CHECK(esp_netif_init());
        ESP_ERROR_CHECK(esp_event_loop_create_default());
        esp_netif_create_default_wifi_sta();
        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_start());
        // Use the configuration to connect to Wi-Fi
        ESP_ERROR_CHECK(esp_event_handler_register(WIFI_EVENT, ESP_EVENT_ANY_ID, &event_handler, NULL));
        esp_wifi_set_config(ESP_IF_WIFI_STA, &s_dpp_wifi_configuration);
        if(esp_wifi_connect() != ESP_OK) // Connect to the network
        {
            ESP_LOGI(WIFI_TAG, "No valid Wi-Fi credentials stored.");
            ESP_LOGI(WIFI_TAG, "Initiating dpp...");
            dpp_enrollee_init();
        }
    } else {
        // No valid Wi-Fi credentials found
        ESP_LOGI(WIFI_TAG, "No valid Wi-Fi credentials stored.");
        ESP_LOGI(WIFI_TAG, "Initiating dpp...");
        dpp_enrollee_init();
    }
}