#include "wifi_service.h"

#include <WiFi.h>
#include <esp_sntp.h>
#include <esp_wifi.h>
#include <lwip/ip_addr.h>

#include <cstring>

#include "ipv4.h"

namespace roro {


namespace {
IPAddress toIp(uint32_t a) { return IPAddress(a >> 24, a >> 16 & 255, a >> 8 & 255, a & 255); }
IPAddress toIp(const std::string& text) {
    uint32_t a = 0;
    net::parseIpv4(text, a);
    return toIp(a);
}
const IPAddress kNoAddress(static_cast<uint32_t>(0));
// A slot lwIP filled from DHCP has an address and no name; ours are set by name.
bool sntpSlotFromDhcp(int i) {
    const ip_addr_t* a = esp_sntp_getserver(i);
    return !esp_sntp_getservername(i) && a && !ip_addr_isany(a);
}
constexpr int kNtpSlots = 3;  // CONFIG_LWIP_SNTP_MAX_SERVERS
constexpr uint32_t kServersEveryMs = 30000;
}  // namespace

// Joins a Saved Network with its IP setting (S1, Q105): Fixed, or DHCP.
void WifiService::join(const std::string& ssid, const std::string& password) {
    const SavedNetwork* n = saved_.find(ssid);
    fixed_ = n && n->fixed;
    if (fixed_) {
        // Nothing comes from DHCP here: forget the NTP servers a previous network offered.
        for (int i = 0; i < kNtpSlots; i++)
            if (sntpSlotFromDhcp(i)) esp_sntp_setserver(i, nullptr);
        WiFi.config(toIp(n->ip.address), toIp(n->ip.gateway), toIp(net::maskOf(n->ip.prefix)),
                    toIp(settings_.getString(Setting::Dns1)), toIp(settings_.getString(Setting::Dns2)));
    } else {
        WiFi.config(kNoAddress, kNoAddress, kNoAddress);  // DHCP
    }
    WiFi.begin(ssid.c_str(), password.empty() ? nullptr : password.c_str());
}

void WifiService::ipSettingChanged(const std::string& ssid) {
    if (controller_.state() != WifiController::State::Connected || controller_.ssid() != ssid) return;
    WiFi.disconnect();
    apply(controller_.disconnected(millis()));
    controller_.retryNow(millis());
}

// DNS and NTP as decided in Q108 and Q110. Run when connected, when a setting changes, and now
// and then: a DHCP renewal puts DHCP's DNS back and clears the NTP slots it didn't fill.
void WifiService::applyServers(Why why) {
    if (controller_.state() != WifiController::State::Connected) return;
    serversCheckedMs_ = millis();
    bool changed = why != Why::Check;

    bool wasFromSettings = dnsFromSettings_;
    dnsFromSettings_ = fixed_ || settings_.getBool(Setting::DnsAlways);
    if (dnsFromSettings_) {
        IPAddress dns1 = toIp(settings_.getString(Setting::Dns1)), dns2 = toIp(settings_.getString(Setting::Dns2));
        if (WiFi.dnsIP(0) != dns1 || WiFi.dnsIP(1) != dns2) WiFi.setDNS(dns1, dns2);
    } else if (why == Why::SettingsChanged && wasFromSettings) {
        // "Always use my DNS" was switched off: only a new lease brings DHCP's servers back.
        ipSettingChanged(controller_.ssid());
        return;
    }

    // NTP: the servers DHCP offered stay first; ours follow.
    int fromDhcp = 0;
    while (!fixed_ && fromDhcp < kNtpSlots && sntpSlotFromDhcp(fromDhcp)) fromDhcp++;
    std::string wanted[2] = {settings_.getString(Setting::Ntp1), settings_.getString(Setting::Ntp2)};
    bool touched = false;
    for (int i = fromDhcp, mine = 0; i < kNtpSlots; i++, mine++) {
        const char* current = esp_sntp_getservername(i);
        if (mine < 2 && !wanted[mine].empty()) {
            if (current && wanted[mine] == current) continue;
            ntpNames_[mine] = wanted[mine];
            esp_sntp_setservername(i, ntpNames_[mine].c_str());
            touched = true;
        } else if (current || sntpSlotFromDhcp(i)) {
            esp_sntp_setserver(i, nullptr);
            touched = true;
        }
    }
    if (!esp_sntp_enabled()) {
        esp_sntp_setoperatingmode(ESP_SNTP_OPMODE_POLL);
        esp_sntp_init();
    } else if (touched || changed) {
        esp_sntp_restart();
    }
    if (touched || changed) ntpWaiting_ = true;
}

WifiService::Connection WifiService::connection() const {
    Connection c;
    c.connected = controller_.state() == WifiController::State::Connected;
    if (!c.connected) return c;
    c.fixed = fixed_;
    c.dnsFromSettings = dnsFromSettings_;
    c.address = WiFi.localIP().toString().c_str();
    c.mask = WiFi.subnetMask().toString().c_str();
    c.prefix = WiFi.subnetCIDR();
    IPAddress gateway = WiFi.gatewayIP();
    if (gateway != kNoAddress) c.gateway = gateway.toString().c_str();
    for (int i = 0; i < 2; i++)
        if (WiFi.dnsIP(i) != kNoAddress) c.dns[i] = WiFi.dnsIP(i).toString().c_str();
    for (int i = 0; i < kNtpSlots; i++) {
        const char* name = esp_sntp_getservername(i);
        bool answered = esp_sntp_getreachability(i) != 0;
        if (name) c.ntp[c.ntpCount++] = {name, false, answered};
        else if (sntpSlotFromDhcp(i)) c.ntp[c.ntpCount++] = {ipaddr_ntoa(esp_sntp_getserver(i)), true, answered};
    }
    return c;
}

std::string WifiService::ip() const {
    return controller_.state() == WifiController::State::Connected ? WiFi.localIP().toString().c_str() : "";
}

int WifiService::rssi() const {
    return controller_.state() == WifiController::State::Connected ? WiFi.RSSI() : 0;
}

void WifiService::startScan(int channel) {
    if (!radioInitialised_) {
        WiFi.mode(WIFI_STA);
        WiFi.setAutoReconnect(false);  // the controller decides when to reconnect
        esp_wifi_set_country_code("EU", true);
        esp_sntp_servermode_dhcp(true);  // before DHCP runs: take the NTP servers it offers (Q110)
        radioInitialised_ = true;
    }
    if (scanRunning_) return;
    // Hidden networks are listed too (with an empty name); the controller skips them.
    if (WiFi.scanNetworks(true /* async */, true /* hidden */, false, channel ? 120 : 300, channel) == WIFI_SCAN_FAILED)
        return;
    scanRunning_ = true;
}

void WifiService::startListScan(int channel) {
    if (scanRunning_) {
        listScanPending_ = true;  // results of the scan already running will do
        return;
    }
    listScanPending_ = true;
    startScan(channel);
}

void WifiService::endListScans() {
    listScanPending_ = false;
    if (controller_.state() == WifiController::State::Off && radioInitialised_ && !scanRunning_) {
        WiFi.mode(WIFI_OFF);
        radioInitialised_ = false;
    }
}

namespace {
const char* security(wifi_auth_mode_t auth) {
    switch (auth) {
        case WIFI_AUTH_OPEN: return "open";
        case WIFI_AUTH_WEP: return "WEP";
        case WIFI_AUTH_WPA_PSK: return "WPA";
        case WIFI_AUTH_WPA2_PSK:
        case WIFI_AUTH_WPA_WPA2_PSK: return "WPA2";
        case WIFI_AUTH_WPA2_ENTERPRISE: return "WPA2-E";
        case WIFI_AUTH_WPA3_PSK:
        case WIFI_AUTH_WPA2_WPA3_PSK: return "WPA3";
        default: return "?";
    }
}
}  // namespace

void WifiService::finishScan(uint32_t nowMs) {
    int n = WiFi.scanComplete();
    if (n == WIFI_SCAN_RUNNING) return;
    scanRunning_ = false;

    std::vector<ScanResult> results;
    std::vector<ScanEntry> entries;
    for (int i = 0; i < n; i++) {
        std::string ssid = WiFi.SSID(i).c_str();
        wifi_auth_mode_t auth = WiFi.encryptionType(i);
        entries.push_back({ssid, WiFi.BSSIDstr(i).c_str(), WiFi.channel(i), WiFi.RSSI(i), auth == WIFI_AUTH_OPEN,
                           security(auth)});
        if (!ssid.empty()) results.push_back({ssid, WiFi.RSSI(i)});
    }
    WiFi.scanDelete();
    if (listScanPending_) {
        listScan_ = entries;
        listScanPending_ = false;
        listScanSeq_++;
    }
    apply(controller_.scanDone(results, nowMs));
}

void WifiService::startNtp() {
    // Not configTime(): it would overwrite the POSIX TZ the Clock applies.
    applyServers(Why::Joined);
    ntpWaiting_ = true;
}

void WifiService::apply(const WifiController::Step& step) {
    using Action = WifiController::Action;
    switch (step.action) {
        case Action::None: break;
        case Action::StartScan: startScan(); break;
        case Action::Connect: join(step.ssid, step.password); break;
        case Action::Disconnect: WiFi.disconnect(); break;
        case Action::RadioOff:
            WiFi.disconnect(true);
            WiFi.mode(WIFI_OFF);
            radioInitialised_ = scanRunning_ = false;
            break;
        case Action::StartMonitor:
            startScan();  // makes sure the radio is up
            WiFi.disconnect();
            break;
        case Action::StopMonitorAndScan:
            esp_wifi_set_promiscuous(false);
            startScan();
            break;
        case Action::StopMonitorAndOff:
            esp_wifi_set_promiscuous(false);
            WiFi.mode(WIFI_OFF);
            radioInitialised_ = scanRunning_ = false;
            break;
    }
}

void WifiService::tick(uint32_t nowMs) {
    using State = WifiController::State;
    if (scanRunning_) finishScan(nowMs);

    wl_status_t status = WiFi.status();
    if (controller_.state() == State::Connecting) {
        if (status == WL_CONNECTED) {
            apply(controller_.connected(nowMs));
            startNtp();
        } else if (status == WL_CONNECT_FAILED || status == WL_NO_SSID_AVAIL) {
            apply(controller_.disconnected(nowMs));
        }
    } else if (controller_.state() == State::Connected && status != WL_CONNECTED) {
        apply(controller_.disconnected(nowMs));
    }

    apply(controller_.update(nowMs, settings_.getBool(Setting::WifiEnabled)));

    if (controller_.state() == State::Connected && nowMs - serversCheckedMs_ >= kServersEveryMs) applyServers(Why::Check);

    if (ntpWaiting_ && sntp_get_sync_status() == SNTP_SYNC_STATUS_COMPLETED) {
        ntpWaiting_ = false;
        clock_.set(static_cast<int64_t>(time(nullptr)), TimeSource::Ntp);
    }
}

}  // namespace roro