#include "nx/firmware/Wifi.h"

#include "freertos/FreeRTOS.h"
#include "freertos/task.h"
#include "freertos/event_groups.h"
#include "esp_system.h"
#include "esp_wifi.h"
#include "esp_event.h"
#include "esp_log.h"
#include "nvs_flash.h"

#include "lwip/err.h"
#include "lwip/sys.h"
#include "lwip/dns.h"
#include "lwip/inet.h"
#include "lwip/sockets.h"

#define TAG "WIFI"

#define WIFI_MAXIMUM_RETRY  -1
#define WIFI_RECONNECT_DELAY_SEC 15
#define WIFI_AP_CHANNEL 11
#define WIFI_MAX_CONNECTION 3

/* 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

#define WIFI_AP_PASSWORD ""

#define DEVICE_NAME_MAX_LEN 32

static const int POWER_SAVE_AP_BEACON_INTERVAL = 10;

static EventGroupHandle_t wifiEventGroup;

static int retryCount = 0;

static WifiSettings* settings = NULL;

static WifiCallback errorCallback = NULL;
static WifiCallback connectedCallback = NULL;

static char deviceName[DEVICE_NAME_MAX_LEN] = {0};

static bool initWifiStation(void);

static bool initWifiAp(void);

// --------- Public API --------- //

bool nxInitWifi(WifiSettings* wifiSettings)
{
  settings = wifiSettings;
  if (settings->mode == NX_WIFI_MODE_AP) {
    return initWifiAp();
  }
  return initWifiStation();
}

void nxSetWifiConnectedCallback(WifiCallback callback)
{
  connectedCallback = callback;
}

void nxSetWifiErrorCallback(WifiCallback callback)
{
  errorCallback = callback;
}


// ------------------------------ //


static void generateDeviceName()
{
  ESP_LOGI(TAG, "Generating device name with prefix %s", settings->devicePrefix);

  uint8_t mac[6];
  esp_read_mac(mac, 0); // 0 means WiFi base MAC address
  sprintf(deviceName, "%s%x%x%x%x%x%x", settings->devicePrefix,
      mac[0], mac[1], mac[2], mac[3], mac[4], mac[5]);
  ESP_LOGI(TAG, "Device name generated: %s", deviceName);
}

static void handleStaEvents(esp_event_base_t event_base, int32_t event_id, void* event_data)
{
  ESP_LOGI(TAG, "WiFi STA event received; event ID: %i", event_id);

  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) {
    ESP_LOGE(TAG, "Cannot connect to the AP");
    if (errorCallback) {
      errorCallback();
    }

    if (WIFI_MAXIMUM_RETRY < 0 || retryCount < WIFI_MAXIMUM_RETRY) {
      ESP_LOGI(TAG, "Sleeping %i seconds before AP reconnection attempt", WIFI_RECONNECT_DELAY_SEC);
      vTaskDelay(WIFI_RECONNECT_DELAY_SEC * 1000 / portTICK_PERIOD_MS);

      ESP_LOGI(TAG, "AP reconnection...");
      esp_wifi_connect();
      retryCount++;
    }
    else {
      xEventGroupSetBits(wifiEventGroup, WIFI_FAIL_BIT);
    }
  }
  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, "IP address acquired: " IPSTR, IP2STR(&event->ip_info.ip));
    retryCount = 0;
    if (connectedCallback) {
      connectedCallback();
    }
    xEventGroupSetBits(wifiEventGroup, WIFI_CONNECTED_BIT);
  }
}

static void handleApEvents(esp_event_base_t event_base, int32_t event_id, void* event_data)
{
  ESP_LOGI(TAG, "WiFi AP event received; event ID: %i", event_id);

  if (event_id == WIFI_EVENT_AP_STACONNECTED) {
    wifi_event_ap_staconnected_t* event = (wifi_event_ap_staconnected_t*) event_data;
    ESP_LOGI(TAG, "Station "MACSTR" joined, AID=%d",
        MAC2STR(event->mac), event->aid);
  } else if (event_id == WIFI_EVENT_AP_STADISCONNECTED) {
    wifi_event_ap_stadisconnected_t* event = (wifi_event_ap_stadisconnected_t*) event_data;
    ESP_LOGI(TAG, "Station "MACSTR" left, AID=%d",
        MAC2STR(event->mac), event->aid);
  }
}

static void eventHandler(void* arg, esp_event_base_t event_base,
    int32_t event_id, void* event_data)
{
  ESP_LOGI(TAG, "Handling WiFi event in mode %i", settings->mode);

  if (settings->mode == NX_WIFI_MODE_STA) {
    handleStaEvents(event_base, event_id, event_data);
  }
  else {
    handleApEvents(event_base, event_id, event_data);
  }
}


static bool initWifiAp(void)
{
  generateDeviceName(NX_WIFI_MODE_AP);
  ESP_LOGI(TAG, "Initializing WiFi AP");

  ESP_ERROR_CHECK(esp_netif_init());
  ESP_ERROR_CHECK(esp_event_loop_create_default());
  esp_netif_create_default_wifi_ap();

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

  ESP_ERROR_CHECK(esp_event_handler_instance_register(WIFI_EVENT,
      ESP_EVENT_ANY_ID,
      &eventHandler,
      NULL,
      NULL));

  wifi_config_t wifiConfig = {
      .ap = {
          .ssid_len = strlen(deviceName),
          .channel = WIFI_AP_CHANNEL,
          .password = WIFI_AP_PASSWORD,
          .max_connection = WIFI_MAX_CONNECTION,
          .authmode = WIFI_AUTH_WPA_WPA2_PSK
      },
  };
  strcpy((char*)(wifiConfig.ap.ssid), deviceName);

  if (strlen((const char*)wifiConfig.ap.password) == 0) {
    wifiConfig.ap.authmode = WIFI_AUTH_OPEN;
  }

  ESP_ERROR_CHECK(esp_wifi_set_mode(WIFI_MODE_AP));
  ESP_ERROR_CHECK(esp_wifi_set_config(WIFI_IF_AP, &wifiConfig));
  ESP_ERROR_CHECK(esp_wifi_start());

  ESP_LOGI(TAG, "WiFi AP initialization finished. SSID: %s password: %s channel: %d",
      wifiConfig.ap.ssid, wifiConfig.ap.password, wifiConfig.ap.channel);
  return true;
}

static void configStaticApAddress(esp_netif_t* sta)
{
  ESP_LOGI(TAG, "Setting static IP address");
  esp_netif_dhcpc_stop(sta);
  esp_netif_ip_info_t ip_info;

  IP4_ADDR(&ip_info.ip,
      settings->ip4addr[0],
      settings->ip4addr[1],
      settings->ip4addr[2],
      settings->ip4addr[3]
  );
  IP4_ADDR(&ip_info.gw,
      settings->ip4gw[0],
      settings->ip4gw[1],
      settings->ip4gw[2],
      settings->ip4gw[3]
  );
  IP4_ADDR(&ip_info.netmask,
      settings->ip4mask[0],
      settings->ip4mask[1],
      settings->ip4mask[2],
      settings->ip4mask[3]
  );
  ESP_LOGI(TAG, "Calling esp_netif_set_ip_info");
  esp_netif_set_ip_info(sta, &ip_info);
}

static void registerWifiStaHandlers(void)
{
  ESP_LOGI(TAG, "Registering WiFi station handlers");

  esp_event_handler_instance_t instance_any_id;
  esp_event_handler_instance_t instance_got_ip;
  ESP_ERROR_CHECK(esp_event_handler_instance_register(WIFI_EVENT,
      ESP_EVENT_ANY_ID,
      &eventHandler,
      NULL,
      &instance_any_id));
  ESP_ERROR_CHECK(esp_event_handler_instance_register(IP_EVENT,
      IP_EVENT_STA_GOT_IP,
      &eventHandler,
      NULL,
      &instance_got_ip));
}

static void configureWifiSta(void)
{
  ESP_LOGI(TAG, "Configuring WiFi station");

  wifi_config_t wifi_config = {
      .sta = {
          .threshold.authmode = WIFI_AUTH_WPA2_PSK,
          .listen_interval = POWER_SAVE_AP_BEACON_INTERVAL,
          .pmf_cfg = {
              .capable = true,
              .required = false
          },
      },
  };

  strcpy((char*)wifi_config.sta.ssid, settings->wname);
  strcpy((char*)wifi_config.sta.password, settings->wpass);

  ESP_ERROR_CHECK(esp_wifi_set_mode(WIFI_MODE_STA) );
  ESP_ERROR_CHECK(esp_wifi_set_config(ESP_IF_WIFI_STA, &wifi_config) );
  ESP_LOGI(TAG, "WiFi station configured");
}

static bool waitForStaConnection(void)
{
  ESP_LOGI(TAG, "Waiting for the WiFi station to connect");

  bool result = false;

  /* 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(wifiEventGroup,
      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 the AP; SSID: %s", settings->wname);

    //resolveBrokerAddr();

    if (*settings->usePowerSave) {
      esp_wifi_set_ps(WIFI_PS_MAX_MODEM);
    }

    result = true;
  } else if (bits & WIFI_FAIL_BIT) {
    ESP_LOGI(TAG, "Failed to connect to SSID:%s, password:%s",
        settings->wname, settings->wpass);
    if (errorCallback) {
      errorCallback();
    }
  } else {
    ESP_LOGE(TAG, "UNEXPECTED EVENT");
  }
  return result;
}

static bool initWifiStation(void)
{
  generateDeviceName(NX_WIFI_MODE_STA);
  ESP_LOGI(TAG, "Initializing WiFi station");

  wifiEventGroup = xEventGroupCreate();

  ESP_ERROR_CHECK(esp_netif_init());
  ESP_ERROR_CHECK(esp_event_loop_create_default());

  esp_netif_t* sta = esp_netif_create_default_wifi_sta();

  if (*settings->useStaticAddr) {
    configStaticApAddress(sta);
  }

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

  registerWifiStaHandlers();

  configureWifiSta();

  ESP_ERROR_CHECK(esp_wifi_start() );

  ESP_LOGI(TAG, "Updating the host name");
  ESP_ERROR_CHECK(esp_netif_set_hostname(sta, deviceName));

  ESP_LOGI(TAG, "Updating DNS server IPv4");

   esp_netif_dns_info_t dns;
   dns.ip.u_addr.ip4.addr = PP_HTONL(LWIP_MAKEU32(
       settings->dns[0],
       settings->dns[1],
       settings->dns[2],
       settings->dns[3]));
   dns.ip.type = IPADDR_TYPE_V4;
   ESP_ERROR_CHECK(esp_netif_set_dns_info(sta, ESP_NETIF_DNS_MAIN, &dns));


  return waitForStaConnection();
}

const char* nxGetWifiDeviceName(void)
{
  if (settings && settings->devicePrefix && strlen(deviceName) == 0) {
    generateDeviceName();
    return deviceName;
  }
  else if (strlen(deviceName)) {
    return deviceName;
  }
  return "unknown";
}