#include "services/vpn_service.h"

#include <Arduino.h>
#include <SD.h>

#include <esp_wireguard.h>
#include <lwip/dns.h>
#include <lwip/tcpip.h>

#include "ipv4.h"
#include "platform/console.h"

namespace roro {

namespace {
constexpr uint32_t kRetryMs = 10000;
constexpr size_t kMaxConf = 4096;

// The library calls lwIP's raw functions and takes no lock; this build checks that the lock is
// held (CONFIG_LWIP_CHECK_THREAD_SAFETY) and stops the device when it isn't.
struct LwipLock {
    LwipLock() { LOCK_TCPIP_CORE(); }
    ~LwipLock() { UNLOCK_TCPIP_CORE(); }
};
}  // namespace

struct VpnService::Tunnel {
    wireguard_config_t config = ESP_WIREGUARD_CONFIG_DEFAULT();
    wireguard_ctx_t ctx = ESP_WIREGUARD_CONTEXT_DEFAULT();
    std::string address, netmask;  // what `config` points into, with config_'s own strings
    bool inited = false, connected = false, isDefault = false, dnsIn = false;
    ip_addr_t dnsBefore[2];
};

const char* VpnService::stateText() const {
    switch (state_) {
        case State::NoConfig: return "not set";
        case State::Off: return "off";
        case State::WaitingWifi: return "waiting for Wi-Fi";
        case State::WaitingClock: return "waiting for the clock";
        case State::Resolving: return "looking up the server";
        case State::Trying: return "no answer yet";
        case State::Up: return "up";
    }
    return "";
}

void VpnService::loadConfig() {
    const std::string& stored = settings_.getString(Setting::VpnConfig);
    configured_ = !stored.empty() && net::parseWgConf(stored, config_).empty();
    if (!configured_) config_ = net::WgConfig();
}

void VpnService::start() {
    loadConfig();
    wanted_ = configured_ && settings_.getBool(Setting::VpnAuto);
    state_ = !configured_ ? State::NoConfig : State::Off;
}

void VpnService::want(bool on, uint32_t seconds) {
    wanted_ = on && configured_;
    timed_ = wanted_ && seconds > 0;
    untilMs_ = millis() + seconds * 1000;
    retryMs_ = 0;
    if (!wanted_) takeDown();
}

std::string VpnService::import(const std::string& confText) {
    net::WgConfig fresh;
    std::string why = net::parseWgConf(confText, fresh);
    if (!why.empty()) return why;
    if (!settings_.setString(Setting::VpnConfig, net::toWgConf(fresh))) return "it couldn't be stored";
    takeDown();  // it comes back up by itself with the new one, if it was wanted
    loadConfig();
    state_ = State::Off;
    return "";
}

std::string VpnService::importFile(StorageService& storage, const std::string& path) {
    std::string text, why;
    bool ran = storage.runAndWait([&]() {
        File f = SD.open(path.c_str(), FILE_READ);
        if (!f || f.isDirectory()) {
            why = "there is no " + path;
            return;
        }
        size_t size = f.size();
        if (size > kMaxConf) why = "that file is too big to be a .conf";
        else {
            text.resize(size);
            if (size && f.read(reinterpret_cast<uint8_t*>(&text[0]), size) != static_cast<int>(size)) why = "the card refused to read it";
        }
        f.close();
    });
    if (!ran) return "no SD card";
    if (!why.empty()) return why;
    return import(text);
}

void VpnService::forget() {
    takeDown();
    wanted_ = false;
    settings_.setString(Setting::VpnConfig, "");
    settings_.setBool(Setting::VpnAuto, false);
    loadConfig();
    state_ = State::NoConfig;
}

void VpnService::bringUp() {
    if (!tunnel_) tunnel_ = new Tunnel();
    Tunnel& t = *tunnel_;
    net::WgRouting routing = net::routingOf(config_);
    t.address = net::formatIpv4(config_.address);
    t.netmask = net::formatIpv4(net::maskOf(routing.full ? config_.prefix : routing.prefix));
    t.config.private_key = config_.privateKey.c_str();
    t.config.public_key = config_.peerKey.c_str();
    t.config.preshared_key = config_.presharedKey.empty() ? nullptr : config_.presharedKey.c_str();
    t.config.address = t.address.c_str();
    t.config.netmask = t.netmask.c_str();
    t.config.endpoint = config_.endpointHost.c_str();
    t.config.port = config_.endpointPort;
    t.config.listen_port = config_.listenPort;
    t.config.persistent_keepalive = static_cast<uint16_t>(config_.keepalive);

    LwipLock lock;
    esp_err_t err = ESP_OK;
    if (!t.inited) {
        err = esp_wireguard_init(&t.config, &t.ctx);
        t.inited = err == ESP_OK;
    }
    if (err == ESP_OK) err = esp_wireguard_connect(&t.ctx);
    if (err == ESP_ERR_RETRY) {  // the server's name isn't resolved yet: asked again at the next tick
        state_ = State::Resolving;
        return;
    }
    if (err == ESP_OK) {
        // What may come out of the tunnel, and with "everything", where every packet now goes. The
        // tunnel's own packets don't: the library sends them on the interface it started on.
        for (int i = 0; i < config_.allowedCount && err == ESP_OK; i++) {
            std::string address = net::formatIpv4(config_.allowed[i].address), mask = net::formatIpv4(net::maskOf(config_.allowed[i].prefix));
            err = esp_wireguard_add_allowed_ip(&t.ctx, address.c_str(), mask.c_str());
        }
    }
    if (err != ESP_OK) {
        error_ = std::string("the tunnel couldn't start (") + esp_err_to_name(err) + ")";
        console.printf("vpn: error %s\n", error_.c_str());
        esp_wireguard_disconnect(&t.ctx);
        t = Tunnel();
        retryMs_ = millis() + kRetryMs;
        state_ = State::Trying;
        return;
    }
    if (routing.full) t.isDefault = esp_wireguard_set_default(&t.ctx) == ESP_OK;
    if (config_.mtu && t.ctx.netif) t.ctx.netif->mtu = static_cast<u16_t>(config_.mtu);
    // The file's DNS servers, if they can be reached through the tunnel at all.
    if (config_.dns[0] && net::wgReaches(config_, config_.dns[0])) {
        t.dnsIn = true;
        for (int i = 0; i < 2; i++) ip_addr_set_any(false, &t.dnsBefore[i]);
        keepDns();
    }
    t.connected = true;
    error_.clear();
    announced_ = false;
    lastHandshake_ = 0;
    state_ = State::Trying;
    console.printf("vpn: started, %s:%u, through it %s\n", config_.endpointHost.c_str(), (unsigned)config_.endpointPort, net::describeWgRouting(config_).c_str());
}

bool VpnService::dnsThroughIt() const { return tunnel_ && tunnel_->dnsIn; }

// With lwIP's lock held. What is found in the two slots, if it isn't the tunnel's, is what goes
// back when the tunnel stops: so a DHCP renewal while it is up is not lost.
void VpnService::keepDns() {
    Tunnel& t = *tunnel_;
    for (int i = 0; i < 2; i++) {
        ip_addr_t wanted;
        ip_addr_set_zero_ip4(&wanted);
        if (config_.dns[i]) ip_addr_set_ip4_u32(&wanted, lwip_htonl(config_.dns[i]));
        const ip_addr_t* now = dns_getserver(static_cast<u8_t>(i));
        if (ip_addr_cmp(now, &wanted)) continue;
        t.dnsBefore[i] = *now;
        dns_setserver(static_cast<u8_t>(i), &wanted);
    }
}

void VpnService::takeDown() {
    if (tunnel_) {
        Tunnel& t = *tunnel_;
        bool dns = t.dnsIn;
        {
            LwipLock lock;
            if (t.dnsIn)
                for (int i = 0; i < 2; i++) dns_setserver(static_cast<u8_t>(i), &t.dnsBefore[i]);
            if (t.isDefault) esp_wireguard_restore_default(&t.ctx);
            if (t.inited) esp_wireguard_disconnect(&t.ctx);
        }
        delete tunnel_;
        tunnel_ = nullptr;
        if (dns) wifi_.holdDns(false);
        console.println("vpn: stopped");
    }
    lastHandshake_ = 0;
    state_ = !configured_ ? State::NoConfig : State::Off;
}

void VpnService::tick(uint32_t nowMs) {
    if (timed_ && static_cast<int32_t>(nowMs - untilMs_) >= 0) want(false);
    if (!configured_ || !wanted_) {
        if (tunnel_) takeDown();
        return;
    }
    // A tunnel doesn't outlive the network it was started on: the next one starts it afresh.
    if (wifi_.state() != WifiController::State::Connected) {
        if (tunnel_) takeDown();
        state_ = State::WaitingWifi;
        return;
    }
    if (clock_.utcNow() < 0) {
        state_ = State::WaitingClock;
        return;
    }
    if (!tunnel_ || !tunnel_->connected) {
        if (retryMs_ && static_cast<int32_t>(nowMs - retryMs_) < 0) return;
        retryMs_ = 0;
        bringUp();
        if (tunnel_ && tunnel_->dnsIn) wifi_.holdDns(true);
        return;
    }
    bool up;
    time_t last = 0;
    {
        LwipLock lock;
        up = esp_wireguard_peer_is_up(&tunnel_->ctx) == ESP_OK;
        esp_wireguard_latest_handshake(&tunnel_->ctx, &last);
        if (tunnel_->dnsIn) keepDns();
    }
    if (last > 0) lastHandshake_ = static_cast<int64_t>(last);
    State was = state_;
    state_ = up ? State::Up : State::Trying;
    if (state_ == State::Up && !announced_) {
        announced_ = true;
        bus_.publish(Event::withText(EventType::Notification, ("VPN up: " + config_.endpointHost).c_str(), static_cast<int32_t>(NotificationLevel::Info)));
    } else if (was == State::Up && state_ == State::Trying) {
        announced_ = false;
        bus_.publish(Event::withText(EventType::Notification, "VPN: the server stopped answering", static_cast<int32_t>(NotificationLevel::Warning)));
    }
}

}  // namespace roro