/*
 * SPDX-FileCopyrightText: 2023-2024 Espressif Systems (Shanghai) CO LTD
 *
 * SPDX-License-Identifier: Unlicense OR CC0-1.0
 */
/*  WiFi softAP & station Example

   This example code is in the Public Domain (or CC0 licensed, at your option.)

   Unless required by applicable law or agreed to in writing, this
   software is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR
   CONDITIONS OF ANY KIND, either express or implied.
*/
#include "esp_event.h"
#include "esp_log.h"
#include "esp_mac.h"
#include "esp_netif.h"
#include "esp_netif_net_stack.h"
#include "esp_netif_types.h"
#include "esp_wifi.h"
#include "lwip/inet.h"
#include "lwip/ip4_addr.h"
#include "lwip/netdb.h"
#include "lwip/sockets.h"
#include "nvs_flash.h"
#include <network.h>
#include <string.h>
#if IP_NAPT
#include "lwip/lwip_napt.h"
#endif
#include "lwip/err.h"
#include "lwip/sys.h"

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

static int s_retry_num = 0;

/* FreeRTOS event group to signal when we are connected/disconnected */
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) {
  if (event_base == WIFI_EVENT && event_id == WIFI_EVENT_AP_STACONNECTED) {
    wifi_event_ap_staconnected_t *event =
        (wifi_event_ap_staconnected_t *)event_data;
    ESP_LOGI(TAG_AP, "Station " MACSTR " joined, AID=%d", MAC2STR(event->mac),
             event->aid);
  } else if (event_base == WIFI_EVENT &&
             event_id == WIFI_EVENT_AP_STADISCONNECTED) {
    wifi_event_ap_stadisconnected_t *event =
        (wifi_event_ap_stadisconnected_t *)event_data;
    ESP_LOGI(TAG_AP, "Station " MACSTR " left, AID=%d, reason:%d",
             MAC2STR(event->mac), event->aid, event->reason);
  } else if (event_base == WIFI_EVENT && event_id == WIFI_EVENT_STA_START) {
    esp_wifi_connect();
    ESP_LOGI(TAG_STA, "Station started");
  } 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_STA, "Got IP:" IPSTR, IP2STR(&event->ip_info.ip));
    s_retry_num = 0;
    xEventGroupSetBits(s_wifi_event_group, WIFI_CONNECTED_BIT);
  }
}

/* Initialize soft AP */
esp_netif_t *wifi_init_softap(void) {
  esp_netif_t *esp_netif_ap = esp_netif_create_default_wifi_ap();

  wifi_config_t wifi_ap_config = {
      .ap =
          {
              .ssid = EXAMPLE_ESP_WIFI_AP_SSID,
              .ssid_len = strlen(EXAMPLE_ESP_WIFI_AP_SSID),
              .channel = EXAMPLE_ESP_WIFI_CHANNEL,
              .password = EXAMPLE_ESP_WIFI_AP_PASSWD,
              .max_connection = EXAMPLE_MAX_STA_CONN,
              .authmode = WIFI_AUTH_WPA2_PSK,
              .pmf_cfg =
                  {
                      .required = false,
                  },
          },
  };

  if (strlen(EXAMPLE_ESP_WIFI_AP_PASSWD) == 0) {
    wifi_ap_config.ap.authmode = WIFI_AUTH_OPEN;
  }

  ESP_ERROR_CHECK(esp_wifi_set_config(WIFI_IF_AP, &wifi_ap_config));

  esp_netif_dhcps_stop(esp_netif_ap);

  ESP_LOGI(TAG_AP, "wifi_init_softap finished. SSID:%s password:%s channel:%d",
           EXAMPLE_ESP_WIFI_AP_SSID, EXAMPLE_ESP_WIFI_AP_PASSWD,
           EXAMPLE_ESP_WIFI_CHANNEL);

  esp_netif_ip_info_t ipinfo;
  IP4_ADDR(&ipinfo.ip, 172, 16, 0, 1);
  IP4_ADDR(&ipinfo.gw, 172, 16, 0, 1);
  IP4_ADDR(&ipinfo.netmask, 255, 255, 255, 0);
  esp_netif_set_ip_info(esp_netif_ap, &ipinfo);
  esp_netif_dhcps_start(esp_netif_ap);

  return esp_netif_ap;
}

/* Initialize wifi station */
esp_netif_t *wifi_init_sta(void) {
  esp_netif_t *esp_netif_sta = esp_netif_create_default_wifi_sta();

  wifi_config_t wifi_sta_config = {
      .sta =
          {
              .ssid = EXAMPLE_ESP_WIFI_STA_SSID,
              .password = EXAMPLE_ESP_WIFI_STA_PASSWD,
              .scan_method = WIFI_ALL_CHANNEL_SCAN,
              .failure_retry_cnt = EXAMPLE_ESP_MAXIMUM_RETRY,
              /* Authmode threshold resets to WPA2 as default if password
               * matches WPA2 standards (password len => 8). If you want to
               * connect the device to deprecated WEP/WPA networks, Please set
               * the threshold value to WIFI_AUTH_WEP/WIFI_AUTH_WPA_PSK and set
               * the password with length and format matching to
               * WIFI_AUTH_WEP/WIFI_AUTH_WPA_PSK standards.
               */
              .threshold.authmode = ESP_WIFI_SCAN_AUTH_MODE_THRESHOLD,
              .sae_pwe_h2e = WPA3_SAE_PWE_BOTH,
          },
  };

  ESP_ERROR_CHECK(esp_wifi_set_config(WIFI_IF_STA, &wifi_sta_config));

  ESP_LOGI(TAG_STA, "wifi_init_sta finished.");

  return esp_netif_sta;
}

void softap_set_dns_addr(esp_netif_t *esp_netif_ap,
                         esp_netif_t *esp_netif_sta) {
  esp_netif_dns_info_t dns;
  esp_netif_get_dns_info(esp_netif_sta, ESP_NETIF_DNS_MAIN, &dns);
  uint8_t dhcps_offer_option = DHCPS_OFFER_DNS;
  ESP_ERROR_CHECK_WITHOUT_ABORT(esp_netif_dhcps_stop(esp_netif_ap));
  ESP_ERROR_CHECK(esp_netif_dhcps_option(
      esp_netif_ap, ESP_NETIF_OP_SET, ESP_NETIF_DOMAIN_NAME_SERVER,
      &dhcps_offer_option, sizeof(dhcps_offer_option)));
  ESP_ERROR_CHECK(
      esp_netif_set_dns_info(esp_netif_ap, ESP_NETIF_DNS_MAIN, &dns));
  ESP_ERROR_CHECK_WITHOUT_ABORT(esp_netif_dhcps_start(esp_netif_ap));
}

void netif_start(void) {
  ESP_ERROR_CHECK(esp_netif_init());
  ESP_ERROR_CHECK(esp_event_loop_create_default());

  // Initialize NVS
  esp_err_t ret = nvs_flash_init();
  if (ret == ESP_ERR_NVS_NO_FREE_PAGES ||
      ret == ESP_ERR_NVS_NEW_VERSION_FOUND) {
    ESP_ERROR_CHECK(nvs_flash_erase());
    ret = nvs_flash_init();
  }
  ESP_ERROR_CHECK(ret);

  /* 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_mode(WIFI_MODE_APSTA));

  /* Initialize AP */
  ESP_LOGI(TAG_AP, "ESP_WIFI_MODE_AP");
  esp_netif_t *esp_netif_ap = wifi_init_softap();

  /* Initialize STA */
  ESP_LOGI(TAG_STA, "ESP_WIFI_MODE_STA");
  esp_netif_t *esp_netif_sta = wifi_init_sta();

  /* Start WiFi */
  ESP_ERROR_CHECK(esp_wifi_start());

  /*
   * 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_STA, "connected to ap SSID:%s password:%s",
             EXAMPLE_ESP_WIFI_STA_SSID, EXAMPLE_ESP_WIFI_STA_PASSWD);
    softap_set_dns_addr(esp_netif_ap, esp_netif_sta);
  } else if (bits & WIFI_FAIL_BIT) {
    ESP_LOGI(TAG_STA, "Failed to connect to SSID:%s, password:%s",
             EXAMPLE_ESP_WIFI_STA_SSID, EXAMPLE_ESP_WIFI_STA_PASSWD);
  } else {
    ESP_LOGE(TAG_STA, "UNEXPECTED EVENT");
    return;
  }

  /* Set sta as the default interface */
  esp_netif_set_default_netif(esp_netif_sta);

  /* Enable napt on the AP netif */
  if (esp_netif_napt_enable(esp_netif_ap) != ESP_OK) {
    ESP_LOGE(TAG_STA, "NAPT not enabled on the netif: %p", esp_netif_ap);
  }
}