#include <string.h>

#include "freertos/FreeRTOS.h"
#include "freertos/task.h"
#include "freertos/event_groups.h"

#include "esp_event_loop.h"
#include "esp_heap_caps.h"
#include "esp_log.h"
#include "esp_wifi.h"
#include "nvs_flash.h"

#include "driver/gpio.h"
#include "driver/uart.h"
#include "driver/spi_common.h"
#include "driver/spi_master.h"

#include "lwip/err.h"
#include "lwip/sys.h"

#include "esp32-wifi-manager.h"

#define MAX_ESSID (32)
#define MAX_PASSWORD (64)

static const char *TAG = "WM";
nvs_handle nvs;

typedef struct wifi_manager_ap_info {
  char ssid[MAX_ESSID];
  char password[MAX_PASSWORD];
} wifi_manager_ap_info_t;

typedef struct ap_store_base {
  int8_t last;
  uint8_t count;
  uint32_t magic;

  wifi_manager_ap_info_t aps[CONFIG_WIFI_MANAGER_MAX_AP];
} ap_store_t;

ap_store_t ap_store;
static EventGroupHandle_t wm_event_group;

uint8_t wifi_manager_ap_count() {
  return ap_store.count;
};

static esp_err_t wifi_event_handler(void* arg, esp_event_base_t event_base, int32_t event_id, void* event_data) {
  ESP_LOGI(TAG, "wifi event: %d", event_id);

    if(event_id == WIFI_EVENT_STA_START) {
      ESP_LOGI(TAG, "SYSTEM_EVENT_STA_START");
      //ESP_ERROR_CHECK(esp_wifi_connect());
    } else if(event_id == WIFI_EVENT_STA_CONNECTED) {
      ESP_LOGI(TAG, "WIFI_EVENT_STA_CONNECTED");
      xEventGroupSetBits(wm_event_group, WIFI_CONNECTED);
    } else if(event_id == IP_EVENT_STA_GOT_IP) {
      ip_event_got_ip_t* event = (ip_event_got_ip_t*) event_data;

      ESP_LOGI(TAG, "SYSTEM_EVENT_STA_GOT_IP");
      ESP_LOGI(TAG, "Got IP: %s\n",
          ip4addr_ntoa(&event->ip_info.ip));

      xEventGroupSetBits(wm_event_group, WM_GOT_IP);
    } else if(event_id == SYSTEM_EVENT_STA_DISCONNECTED) {
      wifi_event_sta_disconnected_t *event = ( wifi_event_sta_disconnected_t*)event_data;

      ESP_LOGI(TAG, "SYSTEM_EVENT_STA_DISCONNECTED: reason: %d", event->reason);

      xEventGroupClearBits(wm_event_group, WIFI_CONNECTED);

    } else if(event_id == SYSTEM_EVENT_SCAN_DONE) {
      ESP_LOGI(TAG, "SYSTEM_EVENT_SCAN_DONE");

      uint8_t i = 0;
      uint8_t j = 0;
      uint16_t ap_count;
      wifi_ap_record_t *records;

      xEventGroupClearBits(wm_event_group, WIFI_SCANNING);

      ESP_ERROR_CHECK(esp_wifi_scan_stop());
      ESP_ERROR_CHECK(esp_wifi_scan_get_ap_num(&ap_count));

      if(ap_count > 0) {
        ap_store.last = -1;

        records = malloc(ap_count * sizeof(wifi_ap_record_t));

        ESP_LOGI(TAG, "Found %d APs", ap_count);

        ESP_ERROR_CHECK(esp_wifi_scan_get_ap_records(&ap_count, (wifi_ap_record_t *)records));

        for(i = 0; i < ap_count; i++) {
          ESP_LOGI(TAG, "AP: %s", records[i].ssid);
        }

        for(i = 0; i < ap_count; i++) {
          ESP_LOGI(TAG, "AP: %s", records[i].ssid);

          for(j = 0; j < ap_store.count; j++)
            if(strcmp(&records[i].ssid, &ap_store.aps[j]) == 0) {
              ap_store.last = j;
              goto found;
            }
        }

found:
        xEventGroupClearBits(wm_event_group, WIFI_SCANNING);
        xEventGroupSetBits(wm_event_group, WIFI_IDLE);

        if(ap_store.last != -1) {
          ESP_LOGI(TAG, "Got matching ESSID: %s", ap_store.aps[ap_store.last].ssid);

          wifi_manager_connect();
        } else {
          ESP_LOGI(TAG, "No matching APs found");
        }
      } else {
        ESP_LOGI(TAG, "No APs found");
      }
  }

  return ESP_OK;
}

uint8_t wifi_manager_add_ap(char *ssid, char *password) {
  strcpy(ap_store.aps[ap_store.count].ssid, (char *)ssid);
  strcpy(ap_store.aps[ap_store.count].password, (char *)password);

  ++ap_store.count;

  if(ap_store.last == -1)
    ap_store.last = 0;

  wifi_manager_save_config();

  return ap_store.count;
};

void wifi_manager_scan() {
  wifi_scan_config_t scan_config = {
    .ssid = &ap_store.aps[ap_store.last].ssid,
    .show_hidden = true,
    .scan_type = WIFI_SCAN_TYPE_ACTIVE,
  };

  ESP_ERROR_CHECK(esp_wifi_scan_start(&scan_config, false));

  ESP_LOGI(TAG, "setting WIFI_SCANNING bit");
  xEventGroupSetBits(wm_event_group, WIFI_SCANNING);
  xEventGroupClearBits(wm_event_group, WIFI_IDLE);
};

void wifi_manager_connect() {
  wifi_config_t wifi_config = {
    .sta = {
      .channel = 11,
    }
  };

  memcpy((char *)&wifi_config.sta.ssid, ap_store.aps[ap_store.last].ssid, strlen(ap_store.aps[ap_store.last].ssid) + 1);
  memcpy((char *)&wifi_config.sta.password, ap_store.aps[ap_store.last].password, strlen(ap_store.aps[ap_store.last].password) + 1);

  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_init_sta finished.");
  ESP_LOGI(TAG, "connect to ap SSID:%s",
      wifi_config.sta.ssid);

  ESP_ERROR_CHECK(esp_wifi_connect());
};

void wifi_manager_main_loop() {
  ap_store.last = 0;
  vTaskDelay(3000 / portTICK_PERIOD_MS);

  while(true) {
    EventBits_t ev_bits = xEventGroupGetBits(wm_event_group);

    if(ev_bits & AP_AVAILABLE &&
      ((ev_bits & (WIFI_SCANNING|WIFI_CONNECTED|WIFI_CONNECTING)) == 0)
    ) {
      ESP_LOGI(TAG, "AP Available, starting WiFi connect");

      wifi_manager_connect();
      //wifi_manager_scan();

      vTaskDelay(10000 / portTICK_PERIOD_MS);
    }

    vTaskDelay(3000 / portTICK_PERIOD_MS);
  }
}

void wifi_manager_save_config() {
  ESP_LOGI(TAG, "writing %d byte config", sizeof(ap_store));
  ESP_LOGI(TAG, "ap_store.magic = 0x%4X", ap_store.magic);
  ESP_LOGI(TAG, "ap_store.count = %d", ap_store.count);
  ESP_LOGI(TAG, "ap_store.last = %d", ap_store.last);

  uint8_t i = 0;

  for(; i < ap_store.count; i++) {
    ESP_LOGI(TAG, "AP: %s", ap_store.aps[i].ssid);
  }

  if(ap_store.count > 0) {
    ESP_LOGI(TAG, "setting AP_AVAILABLE");

    xEventGroupSetBits(wm_event_group, AP_AVAILABLE);
  }

  ESP_ERROR_CHECK(nvs_set_blob(nvs, "config", &ap_store, (size_t)sizeof(ap_store)));
  nvs_commit(nvs);
};

void wifi_manager_reset_store() {
  ESP_LOGW(TAG, "resetting store");

  memset(&ap_store, 0, sizeof(ap_store));

  xEventGroupClearBits(wm_event_group, AP_AVAILABLE);

  ap_store.magic = WIFI_MANAGER_MAGIC;
  ap_store.last = -1;
  ap_store.count = 0;

  wifi_manager_save_config();
};

void wifi_manager_load_config() {
  size_t store_size = 0;
  esp_err_t err;

  err = nvs_get_blob(nvs, "config", NULL, &store_size);
  if (err != ESP_OK && err != ESP_ERR_NVS_NOT_FOUND) ESP_ERROR_CHECK(err);

  err = nvs_get_blob(nvs, "config", &ap_store, &store_size);
  if (err != ESP_OK && err != ESP_ERR_NVS_NOT_FOUND) ESP_ERROR_CHECK(err);

  if(ap_store.magic != WIFI_MANAGER_MAGIC) {
    ESP_LOGI(TAG, "load_config got bad magic: 0x%8X", ap_store.magic);

    wifi_manager_reset_store();

    return wifi_manager_save_config();
  }

  if(ap_store.count > 0) {
    xEventGroupSetBits(wm_event_group, AP_AVAILABLE);
  }

  uint8_t i = 0;
  ESP_LOGI(TAG, "load_config found %d APs", ap_store.count);
  for(; i < ap_store.count; i++)
    ESP_LOGI(TAG, "AP: %s", ap_store.aps[i].ssid);
};

EventGroupHandle_t wifi_manager_start() {
  ESP_ERROR_CHECK(nvs_open("wm-config", NVS_READWRITE, &nvs));

  wm_event_group = xEventGroupCreate();

  wifi_manager_load_config();

  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();

  //cfg.event_handler = wifi_event_handler;

  ESP_ERROR_CHECK(esp_wifi_init(&cfg));
  ESP_ERROR_CHECK(esp_wifi_start());

  esp_event_handler_instance_t instance_got_ip;
  esp_event_handler_instance_t instance_any_id;
  ESP_ERROR_CHECK(esp_event_handler_instance_register(WIFI_EVENT,
        ESP_EVENT_ANY_ID,
        &wifi_event_handler,
        NULL,
        &instance_any_id));
  ESP_ERROR_CHECK(esp_event_handler_instance_register(IP_EVENT,
        ESP_EVENT_ANY_ID,
        &wifi_event_handler,
        NULL,
        &instance_got_ip));

  xTaskCreate(&wifi_manager_main_loop, "wm_main_loop",  4096, NULL, 6, NULL);

  return wm_event_group;
};

int wifi_manager_bootstrap_config() {
  int count = 0;

  ESP_LOGI(TAG, "Bootstrapping config from IDF config");

  if(strlen(CONFIG_WIFI_MANAGER_TEST_AP) > 0) {
    wifi_manager_add_ap(CONFIG_WIFI_MANAGER_TEST_AP, CONFIG_WIFI_MANAGER_TEST_PWD);
    count++;
  }

  if(strlen(CONFIG_WIFI_MANAGER_TEST_AP1) > 0) {
    wifi_manager_add_ap(CONFIG_WIFI_MANAGER_TEST_AP1, CONFIG_WIFI_MANAGER_TEST_PWD1);
    count++;
  }

  if(strlen(CONFIG_WIFI_MANAGER_TEST_AP2) > 0) {
    wifi_manager_add_ap(CONFIG_WIFI_MANAGER_TEST_AP2, CONFIG_WIFI_MANAGER_TEST_PWD2);
    count++;
  }

  return count;
}