#include <string.h>

#include <freertos/FreeRTOS.h>
#include <esp_wifi.h>
#include <esp_log.h>
#include <esp_event.h>
#include <lwip/err.h>
#include <lwip/sys.h>

#include "wifi.h"

static EventGroupHandle_t wifi_event_group;
static int retries = 0;

esp_err_t init_wifi() {
  wifi_event_group = xEventGroupCreate();
  esp_netif_init();
  esp_err_t err = esp_event_loop_create_default();
  if(err != ESP_OK) {
    return err;
  }
  esp_netif_create_default_wifi_sta();
  wifi_init_config_t wifi_initiation = WIFI_INIT_CONFIG_DEFAULT();
  return esp_wifi_init(&wifi_initiation);
}

static void wifi_event_handler(void* arg, esp_event_base_t event_base, int32_t event_id, void* event_data) {
  const char* TAG = "wifi_event_handler";

  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 (retries++ < MAX_WIFI_RECONNECT_RETRIES) {
      esp_wifi_connect();
      ESP_LOGW(TAG, "retry to connect to the AP");
    } else {
      xEventGroupSetBits(wifi_event_group, WIFI_FAIL_BIT);
    }
    ESP_LOGE(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));
    retries = 0;
    xEventGroupSetBits(wifi_event_group, WIFI_CONNECTED_BIT);
  }
}

esp_err_t connect_wifi(const char* ssid, const char* password) {
  const char* TAG = "connect_wifi";
  esp_err_t err = esp_event_handler_register(WIFI_EVENT, ESP_EVENT_ANY_ID, wifi_event_handler, NULL);
  if(err != ESP_OK) {
    return err;
  }
  err = esp_event_handler_register(IP_EVENT, IP_EVENT_STA_GOT_IP, wifi_event_handler, NULL);
  if(err != ESP_OK) {
    return err;
  }

  wifi_config_t wifi_config = {
    .sta = {
      .ssid = "",
      .password = "",
      .scan_method = WIFI_FAST_SCAN,
      .threshold={.authmode = WIFI_AUTH_WPA2_PSK},
      .sae_pwe_h2e = WPA3_SAE_PWE_BOTH,
      .failure_retry_cnt = MAX_WIFI_RECONNECT_RETRIES,
      .sae_h2e_identifier = "",
    }
  };

  strcpy((char*)wifi_config.sta.ssid, ssid);
  strcpy((char*)wifi_config.sta.password, password);

  err = esp_wifi_set_mode(WIFI_MODE_STA);
  if(err != ESP_OK) {
    return err;
  }
  err = esp_wifi_set_config(WIFI_IF_STA, &wifi_config);
  if(err != ESP_OK) {
    return err;
  }
  err = esp_wifi_start();
  if(err != ESP_OK) {
    return err;
  }

  ESP_LOGI(TAG, "wifi_init_sta finished.");

  EventBits_t bits = xEventGroupWaitBits(
    wifi_event_group,
    WIFI_CONNECTED_BIT | WIFI_FAIL_BIT,
    pdFALSE,
    pdFALSE,
    portMAX_DELAY
  );

  if (bits & WIFI_CONNECTED_BIT) {
    ESP_LOGI(TAG, "connected to ap SSID:%s", ssid);
  } else if (bits & WIFI_FAIL_BIT) {
    ESP_LOGE(TAG, "Failed to connect to SSID:%s Password:%s", wifi_config.sta.ssid, wifi_config.sta.password);
  } else {
    ESP_LOGE(TAG, "UNEXPECTED EVENT");
  }

  return ESP_OK;
}