#include "Wifi.hpp"
#include "sdkconfig.h"
#include "esp_log.h"
#include "esp_netif.h"
#include "nvs_flash.h"
#include "nvs.h"
#include <cstring>
#include <algorithm>

static const char* TAG = "Wifi";

Wifi& Wifi::instance() {
    static Wifi instance;
    return instance;
}

void Wifi::eventHandler(void* arg, esp_event_base_t event_base,
                        int32_t event_id, void* event_data) {
    Wifi& wifi = Wifi::instance();

    if (event_base == WIFI_EVENT) {
        switch (event_id) {
            case WIFI_EVENT_STA_START:
                ESP_LOGI(TAG, "WiFi station started");
                break;
            case WIFI_EVENT_STA_DISCONNECTED:
                ESP_LOGI(TAG, "WiFi disconnected");
                wifi.connected_ = false;
                wifi.currentSSID_ = "";
                if (wifi.statusCallback_) {
                    wifi.statusCallback_(false);
                }
                xEventGroupSetBits(wifi.wifiEventGroup_, WIFI_FAIL_BIT);
                break;
            case WIFI_EVENT_SCAN_DONE:
                ESP_LOGI(TAG, "WiFi scan done");
                xEventGroupSetBits(wifi.wifiEventGroup_, WIFI_SCAN_DONE_BIT);
                break;
            default:
                break;
        }
    } else if (event_base == IP_EVENT) {
        if (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));
            wifi.connected_ = true;
            if (wifi.statusCallback_) {
                wifi.statusCallback_(true);
            }
            xEventGroupSetBits(wifi.wifiEventGroup_, WIFI_CONNECTED_BIT);
        }
    }
}

void Wifi::init() {
    if (initialized_) {
        return;
    }

    ESP_LOGI(TAG, "Initializing WiFi");

    wifiEventGroup_ = xEventGroupCreate();

    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();
    ESP_ERROR_CHECK(esp_wifi_init(&cfg));

    ESP_ERROR_CHECK(esp_event_handler_instance_register(WIFI_EVENT,
                                                        ESP_EVENT_ANY_ID,
                                                        &eventHandler,
                                                        nullptr,
                                                        nullptr));
    ESP_ERROR_CHECK(esp_event_handler_instance_register(IP_EVENT,
                                                        IP_EVENT_STA_GOT_IP,
                                                        &eventHandler,
                                                        nullptr,
                                                        nullptr));

    ESP_ERROR_CHECK(esp_wifi_set_mode(WIFI_MODE_STA));
    ESP_ERROR_CHECK(esp_wifi_start());

    initialized_ = true;

    auto savedNetworks = getSavedNetworks();
    if (!savedNetworks.empty()) {
        std::string password;
        if (getSavedPassword(savedNetworks[0].c_str(), password)) {
            ESP_LOGI(TAG, "Auto-connecting to saved network: %s", savedNetworks[0].c_str());
            connect(savedNetworks[0].c_str(), password.c_str());
        }
    }

    ESP_LOGI(TAG, "WiFi initialized");
}

void Wifi::scan(std::function<void(std::vector<wifi_ap_record_t>&)> callback) {
    scanCallback_ = callback;

    xEventGroupClearBits(wifiEventGroup_, WIFI_SCAN_DONE_BIT);

    wifi_scan_config_t scan_config = {};
    scan_config.ssid = nullptr;
    scan_config.bssid = nullptr;
    scan_config.channel = 0;
    scan_config.show_hidden = false;
    scan_config.scan_type = WIFI_SCAN_TYPE_ACTIVE;
    scan_config.scan_time.active.min = 100;
    scan_config.scan_time.active.max = 300;

    ESP_ERROR_CHECK(esp_wifi_scan_start(&scan_config, false));

    xTaskCreate([](void* arg) {
        Wifi& wifi = Wifi::instance();

        xEventGroupWaitBits(wifi.wifiEventGroup_, WIFI_SCAN_DONE_BIT,
                           pdTRUE, pdFALSE, pdMS_TO_TICKS(10000));

        uint16_t ap_count = 0;
        esp_wifi_scan_get_ap_num(&ap_count);

        std::vector<wifi_ap_record_t> ap_records(ap_count);
        if (ap_count > 0) {
            esp_wifi_scan_get_ap_records(&ap_count, ap_records.data());
        }

        ESP_LOGI(TAG, "Found %d access points", ap_count);

        if (wifi.scanCallback_) {
            wifi.scanCallback_(ap_records);
        }

        vTaskDelete(nullptr);
    }, "wifi_scan", 4096, nullptr, 5, nullptr);
}

void Wifi::connect(const char* ssid, const char* password) {
    ESP_LOGI(TAG, "Connecting to %s", ssid);

    xEventGroupClearBits(wifiEventGroup_, WIFI_CONNECTED_BIT | WIFI_FAIL_BIT);

    wifi_config_t wifi_config = {};
    strncpy((char*)wifi_config.sta.ssid, ssid, sizeof(wifi_config.sta.ssid) - 1);
    strncpy((char*)wifi_config.sta.password, password, sizeof(wifi_config.sta.password) - 1);
    wifi_config.sta.threshold.authmode = WIFI_AUTH_WPA2_PSK;
    wifi_config.sta.pmf_cfg.capable = true;
    wifi_config.sta.pmf_cfg.required = false;

    currentSSID_ = ssid;

    ESP_ERROR_CHECK(esp_wifi_set_config(WIFI_IF_STA, &wifi_config));
    ESP_ERROR_CHECK(esp_wifi_connect());
}

void Wifi::disconnect() {
    ESP_LOGI(TAG, "Disconnecting");
    esp_wifi_disconnect();
    connected_ = false;
    currentSSID_ = "";
}

bool Wifi::isConnected() {
    return connected_;
}

std::string Wifi::getCurrentSSID() {
    return currentSSID_;
}

void Wifi::setStatusCallback(std::function<void(bool connected)> callback) {
    statusCallback_ = callback;
}

void Wifi::saveNetwork(const char* ssid, const char* password) {
    nvs_handle_t handle;
    esp_err_t err = nvs_open(NVS_NAMESPACE, NVS_READWRITE, &handle);
    if (err != ESP_OK) {
        ESP_LOGE(TAG, "Error opening NVS: %s", esp_err_to_name(err));
        return;
    }

    char key[32];
    snprintf(key, sizeof(key), "pw_%s", ssid);

    size_t key_len = strlen(key);
    if (key_len > 15) {
        key[15] = '\0';
    }

    err = nvs_set_str(handle, key, password);
    if (err != ESP_OK) {
        ESP_LOGE(TAG, "Error saving password: %s", esp_err_to_name(err));
    }

    auto networks = getSavedNetworks();
    bool found = false;
    for (const auto& net : networks) {
        if (net == ssid) {
            found = true;
            break;
        }
    }

    if (!found) {
        networks.insert(networks.begin(), ssid);
        if (networks.size() > MAX_SAVED_NETWORKS) {
            networks.pop_back();
        }

        std::string networkList;
        for (size_t i = 0; i < networks.size(); i++) {
            if (i > 0) networkList += ";";
            networkList += networks[i];
        }
        nvs_set_str(handle, "networks", networkList.c_str());
    }

    nvs_commit(handle);
    nvs_close(handle);

    ESP_LOGI(TAG, "Saved network: %s", ssid);
}

void Wifi::removeNetwork(const char* ssid) {
    nvs_handle_t handle;
    esp_err_t err = nvs_open(NVS_NAMESPACE, NVS_READWRITE, &handle);
    if (err != ESP_OK) {
        ESP_LOGE(TAG, "Error opening NVS: %s", esp_err_to_name(err));
        return;
    }

    char key[32];
    snprintf(key, sizeof(key), "pw_%s", ssid);
    size_t key_len = strlen(key);
    if (key_len > 15) {
        key[15] = '\0';
    }
    nvs_erase_key(handle, key);

    auto networks = getSavedNetworks();
    std::string ssidStr(ssid);
    networks.erase(std::remove(networks.begin(), networks.end(), ssidStr), networks.end());

    std::string networkList;
    for (size_t i = 0; i < networks.size(); i++) {
        if (i > 0) networkList += ";";
        networkList += networks[i];
    }
    nvs_set_str(handle, "networks", networkList.c_str());

    nvs_commit(handle);
    nvs_close(handle);

    ESP_LOGI(TAG, "Removed network: %s", ssid);
}

std::vector<std::string> Wifi::getSavedNetworks() {
    std::vector<std::string> networks;

    nvs_handle_t handle;
    esp_err_t err = nvs_open(NVS_NAMESPACE, NVS_READONLY, &handle);
    if (err != ESP_OK) {
        return networks;
    }

    size_t required_size = 0;
    err = nvs_get_str(handle, "networks", nullptr, &required_size);
    if (err == ESP_OK && required_size > 0) {
        std::string networkList(required_size, '\0');
        nvs_get_str(handle, "networks", &networkList[0], &required_size);

        size_t start = 0;
        size_t end = networkList.find(';');
        while (end != std::string::npos) {
            networks.push_back(networkList.substr(start, end - start));
            start = end + 1;
            end = networkList.find(';', start);
        }
        if (start < networkList.size() - 1) {
            networks.push_back(networkList.substr(start));
        }
    }

    nvs_close(handle);
    return networks;
}

bool Wifi::getSavedPassword(const char* ssid, std::string& password) {
    nvs_handle_t handle;
    esp_err_t err = nvs_open(NVS_NAMESPACE, NVS_READONLY, &handle);
    if (err != ESP_OK) {
        return false;
    }

    char key[32];
    snprintf(key, sizeof(key), "pw_%s", ssid);
    size_t key_len = strlen(key);
    if (key_len > 15) {
        key[15] = '\0';
    }

    size_t required_size = 0;
    err = nvs_get_str(handle, key, nullptr, &required_size);
    if (err == ESP_OK && required_size > 0) {
        password.resize(required_size);
        nvs_get_str(handle, key, &password[0], &required_size);
        nvs_close(handle);
        return true;
    }

    nvs_close(handle);
    return false;
}