#include "wifi_con.h"
#include "creds.h"
#include "esp_netif_types.h"
#include "esp_wifi_types_generic.h"

static void wifiEventHandler(void *arg, esp_event_base_t eventBase,
                             int32_t eventId, void *eventData) {
  static int retryNum = 0;
  if (eventBase == WIFI_EVENT && eventId == WIFI_EVENT_STA_START) {
    ESP_LOGI(WIFI_TAG, "Connecting to WiFi...");
    esp_wifi_connect();
  } else if (eventBase == WIFI_EVENT &&
             eventId == WIFI_EVENT_STA_DISCONNECTED) {
    if (retryNum < MAX_FAILURES) {
      ESP_LOGI(WIFI_TAG, "retrying to connect.. %d", retryNum);
      esp_wifi_connect();
      retryNum++;
    } else {
      xEventGroupSetBits(wifiEventGroup, WIFI_FAILURE);
    }
  }
}

static void ipEventHandler(void *arg, esp_event_base_t eventBase,
                           int32_t eventId, void *eventData) {
  if (eventBase == IP_EVENT && eventId == IP_EVENT_STA_GOT_IP) {
    ip_event_got_ip_t *event = (ip_event_got_ip_t *)eventData;
    ESP_LOGI(WIFI_TAG, "IP Address Assigned: " IPSTR,
             IP2STR(&event->ip_info.ip));
    xEventGroupSetBits(wifiEventGroup, WIFI_SUCCESS);
  }
}

esp_err_t initWifi() {

  wifiEventGroup = xEventGroupCreate();

  ESP_ERROR_CHECK(esp_netif_init());

  ESP_ERROR_CHECK(esp_event_loop_create_default());
  esp_netif_create_default_wifi_sta();

  wifi_init_config_t config = WIFI_INIT_CONFIG_DEFAULT();
  ESP_ERROR_CHECK(esp_wifi_init(&config));

  esp_event_handler_instance_t instanceWIFI;
  esp_event_handler_instance_t instanceGotIP;

  ESP_ERROR_CHECK(esp_event_handler_instance_register(
      WIFI_EVENT, ESP_EVENT_ANY_ID, &wifiEventHandler, NULL, &instanceWIFI));

  ESP_ERROR_CHECK(esp_event_handler_instance_register(
      IP_EVENT, IP_EVENT_STA_GOT_IP, &ipEventHandler, NULL, &instanceGotIP));

  wifi_config_t wifiConfig = {.sta = {
                                  .ssid = SSID,
                                  .password = WIFI_PASS,
                                  .threshold.authmode = WIFI_AUTH_WPA2_PSK,
                                  .sae_pwe_h2e = WPA3_SAE_PWE_BOTH,
                              }};

  ESP_ERROR_CHECK(esp_wifi_set_mode(WIFI_MODE_STA));
  ESP_ERROR_CHECK(esp_wifi_set_config(WIFI_IF_STA, &wifiConfig));
  ESP_ERROR_CHECK(esp_wifi_start());

  ESP_LOGI(WIFI_TAG, "WiFi init station complete");

  EventBits_t bits =
      xEventGroupWaitBits(wifiEventGroup, WIFI_SUCCESS | WIFI_FAILURE, pdFALSE,
                          pdFALSE, portMAX_DELAY);

  if (bits & WIFI_SUCCESS) {
    return WIFI_SUCCESS;
  } else if (bits & WIFI_FAILURE) {
    return WIFI_FAILURE;
  } else {
    return WIFI_FAILURE;
  }
}