#include "WiFiManager.h"
#include <Arduino.h>
#include <ArduinoJson.h>
#include <ArduinoLog.h>
#include <DNSServer.h>
#include <LittleFS.h>

#ifndef DNS_PORT
#define DNS_PORT 53
#endif

#ifdef ESP32
#include <WiFi.h>
#include <WiFiMulti.h>
#include <esp_wifi.h>
WiFiMulti wifiMulti;
WiFiEventCb disconnectedEventHandler;
typedef wifi_ap_record_t bss_info;
int scanned_ap_num;
#elif defined(ESP8266)
#include <ESP8266WiFiMulti.h>
ESP8266WiFiMulti wifiMulti;
WiFiEventHandler disconnectedEventHandler;
typedef nullptr_t WiFiEventInfo_t;
#endif

std::vector<JsonObject> *connection_queue;
JsonDocument doc;
const uint8_t *bssid;

// keep track of current connection
uint currentConfig;
uint currentTry;
bool inRecovery = false;

unsigned long tryStartMillis;
unsigned long lastCheckMillis;

int sortByRSSI(const void *a, const void *b) {
    return (((bss_info *)b)->rssi - ((bss_info *)a)->rssi);
}

bool matches_ssid(const std::vector<const bss_info *> &networks, JsonVariant config) {
    for (auto &&network : networks) {
        if (memcmp((char *)network->ssid, config["ssid"].as<String>().c_str(), strlen((const char *)network->ssid)) == 0) {
            return true;
        }
    }
    return false;
}

bool matches_bssid(const std::vector<const bss_info *> &networks, JsonVariant config) {
    char bssid_char[sizeof(bss_info().bssid) * 3];
    for (auto &&network : networks) {
        bssid = network->bssid;
        snprintf_P(bssid_char, sizeof(bssid_char), "%02X:%02X:%02X:%02X:%02X:%02X", bssid[0], bssid[1], bssid[2], bssid[3], bssid[4], bssid[5]);

        for (auto &&config_bssid : config["bssid"].as<JsonArray>()) {
            if (config_bssid == bssid_char) {
                return true;
            }
        }
    }
    return false;
}

void startRecoveryAP() {
    inRecovery = true;
    WiFi.disconnect();
    WiFi.softAP("ESP reached recovery state", NULL);
}

void disconnectHandler(WiFiEvent_t event, WiFiEventInfo_t info) {
    if (inRecovery) {
        return;
    }

    if (connection_queue->size() == 0) {
        Log.noticeln("No configurations found to try. Starting recovery sequence");
        startRecoveryAP();
        return;
    }

    Log.noticeln("In disc-handler");
    JsonVariant config = connection_queue->at(currentConfig);
    if (currentTry < config["retries"]) {
        currentTry++;
        Log.noticeln("trying again to connect");
    } else {
        if (currentConfig + 1 == connection_queue->size()) {
            Log.noticeln("Station disconnected. No more configurations to try. Starting recovery sequence");
            startRecoveryAP();
            return;
        } else {
            currentTry = 0;
            currentConfig++;
            config = connection_queue->at(currentConfig);
        }
    }

    Log.noticeln("Station disconnected. Connecting to %s", config["name"].as<String>().c_str());

    String psk;
    const uint8_t *bssid = nullptr;

    if (!config["psk"].isNull()) {
        psk = config["psk"].as<String>();
    }

    if (!config["bssid"].isNull() && config["bssid"].size() > 0) {
        Log.noticeln("Found bssid %s", config["bssid"].as<String>().c_str());
        bssid = (const uint8_t *)config["bssid"].as<String>().c_str();
    }

    Log.noticeln("Connecting with: %s %s %s", config["ssid"].as<String>().c_str(), psk, bssid);
    tryStartMillis = millis();
    if (!config["hostname"].isNull()) {
        Log.noticeln("Set hostname %s", config["hostname"].as<String>().c_str());
#ifdef ESP32
        WiFi.mode(WIFI_MODE_NULL);
#endif
        WiFi.setHostname(config["hostname"].as<String>().c_str());
    }

    WiFi.begin(config["ssid"].as<String>(), psk, 0, bssid);
    Log.noticeln("Set Wifi to ssid: '%s' psk: '%s' bssid: %s", WiFi.SSID().c_str(), WiFi.psk().c_str(), WiFi.BSSIDstr().c_str());
}

#ifdef ESP8266
void disconnectHandler_ESP8266(const WiFiEventStationModeDisconnected &event = {}) {
    disconnectHandler(WiFiEvent_t::WIFI_EVENT_ANY, nullptr);
}
#endif

void getScanResults(int count) {
    Log.noticeln("Scan done. Found %i radios", count);

    char bssid_char[sizeof(bss_info().bssid) * 3];

    std::vector<const bss_info *> networks;
    networks.resize(count);
    size_t max_ssid_length = 0;
    for (int i = 0; i < count; i++) {
#ifdef ESP32
        networks.at(i) = (bss_info *)WiFiScanClass::getScanInfoByIndex(i);
#elif defined(ESP8266)
        networks.at(i) = WiFi.getScanInfoByIndex(i);
#endif
        max_ssid_length = max(strlen((const char *)networks.at(i)->ssid), max_ssid_length);
    }

    std::sort(networks.begin(), networks.end(), sortByRSSI);

    char buffer[128];

    for (auto &network : networks) {
        bssid = network->bssid;
        snprintf_P(bssid_char, sizeof(bssid_char), "%02X:%02X:%02X:%02X:%02X:%02X", bssid[0], bssid[1], bssid[2], bssid[3], bssid[4], bssid[5]);
#ifdef ESP32
        const bool ssid_is_hidden = strlen((const char *)network->ssid) == 0;
        const uint8_t channel = network->primary;
#elif defined(ESP8266)
        const bool ssid_is_hidden = network->is_hidden;
        const uint8_t channel = network->channel;
#endif
        snprintf_P(buffer, sizeof(buffer), "SSID: %-*s, RSSI: %i, BSSID: %s, hidden: %i, channel: %i", max_ssid_length, network->ssid, network->rssi, bssid_char, ssid_is_hidden, channel);
        Log.noticeln("%s", buffer);
    }

    connection_queue = new std::vector<JsonObject>;
    connection_queue->reserve(doc["WiFiSta"].size()); // worst case: all configurations are valid

    // TODO: add every found bssid for a given ssid to the queue (in case one bssid with same ssid has a different psk)
    for (JsonVariant config : doc["WiFiSta"].as<JsonArray>()) {
        Log.noticeln("bssid is Null: %d, BSSID is empty: %d", config["bssid"].isNull(), config["bssid"].as<String>().length());
        Log.noticeln("if (%t or (%t and (%t or %t) or %t))", config["alwaysTry"].as<bool>(), matches_ssid(networks, config), config["bssid"].isNull(), config["bssid"].as<String>().length() == 0, matches_bssid(networks, config));
        if (config["alwaysTry"].as<bool>() or (matches_ssid(networks, config) and ((config["bssid"].isNull() or config["bssid"].as<String>().length() == 0) or matches_bssid(networks, config)))) {
            connection_queue->push_back(config);
            Log.noticeln("Adding %s to queue", config["name"].as<String>().c_str());
        }
    }

#ifdef ESP32
    disconnectHandler(WiFiEvent_t::ARDUINO_EVENT_WIFI_STA_START, arduino_event_info_t{}); // start connecting to configurations
#endif
#ifdef ESP8266
    disconnectHandler(WiFiEvent_t::WIFI_EVENT_ANY, nullptr); // start connecting to configurations
#endif
}

void getScanResultsESP32(WiFiEvent_t event, WiFiEventInfo_t info) {
    getScanResults(WiFi.scanComplete()); // assume scan always completed since we are in scan complete CB
}

wl_status_t WiFiManager::tick(unsigned long update_interval = 100) {
    unsigned long currentMillis = millis();
    if (!(currentMillis - lastCheckMillis >= update_interval)) {
        return this->lastState;
    }

    Log.noticeln("In tick");

    lastCheckMillis = currentMillis;

    this->lastState = WiFi.status();

    if (this->lastState == WL_CONNECTED) {
        return this->lastState;
    }

    if (connection_queue == nullptr || currentConfig >= connection_queue->size()) {
        Log.warningln("Tried to tick but config still loading");
        return WL_IDLE_STATUS;
    }

    Log.noticeln("In tick with %d configs", connection_queue->size());

    if (currentMillis - tryStartMillis >= connection_queue->at(currentConfig)["timeout"]) {
        Log.noticeln("Disconnect from WiFi (reached timeout)");
        WiFi.disconnect();
    }
    return WL_DISCONNECTED;
}

WiFiManager::WiFiManager() {}
WiFiManager::~WiFiManager() {}
void WiFiManager::begin() {

    // TODO: load connections
}
bool WiFiManager::loadFromJson(File file) {

    // Deserialize the JSON document
    DeserializationError error = deserializeJson(doc, file);
    if (error)
        Serial.println(F("Failed to read file, using default configuration"));

    String setupAP;

    Log.noticeln("Found %i configurations, overflowed: %T ", doc["WiFiSta"].size(), doc.overflowed());

    setupAP = doc["SetupAP"].as<String>();

    for (JsonVariant i : doc["WiFiSta"].as<JsonArray>()) {
        Log.noticeln("Name: %s", i["name"].as<String>().c_str());
        Log.noticeln("SSID: %s", i["ssid"].as<String>().c_str());
        Log.noticeln("PSK: %s", i["psk"].as<String>().c_str());
    }

#ifdef ESP32
    WiFi.onEvent(disconnectHandler, WiFiEvent_t::ARDUINO_EVENT_WIFI_STA_DISCONNECTED);
    WiFi.onEvent(getScanResultsESP32, WiFiEvent_t::ARDUINO_EVENT_WIFI_SCAN_DONE);
    WiFi.scanNetworks(true, true);
#elif defined(ESP8266)
    disconnectedEventHandler = WiFi.onStationModeDisconnected(disconnectHandler_ESP8266);
    WiFi.scanNetworksAsync(getScanResults, true);
#endif

    loadConnections();

    file.close();
    return true;
}

bool WiFiManager::loadConnections() {
    WiFi.persistent(false);
    WiFi.mode(WIFI_STA);

    IPAddress local_ip, gateway, subnet;
    local_ip.fromString("10.0.0.1");
    gateway.fromString("10.0.0.1");
    subnet.fromString(F("255.255.255.0"));

    uint8_t newMACAddress[WL_MAC_ADDR_LENGTH] = {0x32, 0xAE, 0xA4, 0x07, 0x0D, 0x66};

    const bool ap = false;

    // TODO: can't set AP channel, can't scan AP channel -> set ap, then start async scan seems to work
    if (ap) {

        if (setApMac(newMACAddress))
            Log.infoln("set MAC for SOFTAP_IF to %s", WiFi.softAPmacAddress().c_str());
        else
            Log.warningln("Unable to set MAC for SOFTAP_IF. Using %s instead", WiFi.softAPmacAddress().c_str());
        WiFi.softAPConfig(local_ip, gateway, subnet);
        WiFi.softAP("ESP Setup", "tempPW123", 5, 0, 1);
        Serial.println(WiFi.softAPIP());
        Serial.println(WiFi.macAddress());
    }

    return true;
}

bool WiFiManager::setApMac(uint8_t mac[WL_MAC_ADDR_LENGTH]) {
#ifdef ESP32
    esp_wifi_set_mac(WIFI_IF_AP, &mac[0]);
#elif defined(ESP8266)
    wifi_set_macaddr(SOFTAP_IF, &mac[0]);
#endif

    char charMacAddress[(sizeof(*mac) * 3)];
    snprintf_P(charMacAddress, sizeof(charMacAddress), "%02X:%02X:%02X:%02X:%02X:%02X", mac[0], mac[1], mac[2], mac[3], mac[4], mac[5]);
    String actualMac = WiFi.softAPmacAddress();
    return strncmp(charMacAddress, actualMac.c_str(), sizeof(charMacAddress));
}