#include "engine_crypto.h"
#include <iostream>
#include <string>
#include <fstream>
#include <vector>
#include <chrono>
#include <unistd.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <sys/wait.h>
#include <ifaddrs.h>
#include <net/if.h>
#include <pwd.h>
#include <curl/curl.h>
#include <sodium.h>

const std::string NODE_ID = "669427";
const std::string RELAY_IP = "172.28.47.1";      // Vault (869377) via compass tunnel per NET-02
const std::string RELAY_PORT = "8302";             // Vault keyserv port
const std::string KEY_PATH = "/dev/shm/vault.key";
const std::string HMAC_KEY_FILE = "/dev/shm/compass_hmac.key"; // compass tunnel HMAC key

bool check_tun_mesh() {
    // 1. Check if tun0 exists
    bool tun0_found = false;
    struct ifaddrs *ifaddr, *ifa;
    if (getifaddrs(&ifaddr) == -1) {
        std::cerr << "ERROR: getifaddrs failed\n";
        return false;
    }
    for (ifa = ifaddr; ifa != NULL; ifa = ifa->ifa_next) {
        if (ifa->ifa_name != NULL && std::string(ifa->ifa_name) == "compass") {
            tun0_found = true;
            break;
        }
    }
    freeifaddrs(ifaddr);

    if (!tun0_found) {
        std::cerr << "ERROR: tunnel interface not found\n";
        return false;
    }

    // 2. Read latest handshakes using wg safely
    int pipefd[2];
    if (pipe(pipefd) == -1) {
        std::cerr << "ERROR: pipe failed\n";
        return false;
    }

    pid_t pid = fork();
    if (pid == -1) {
        std::cerr << "ERROR: fork failed\n";
        return false;
    }

    if (pid == 0) {
        // Child
        close(pipefd[0]);
        dup2(pipefd[1], STDOUT_FILENO);
        close(pipefd[1]);
        
        char* args[] = {(char*)"wg", (char*)"show", (char*)"compass", (char*)"latest-handshakes", NULL};
        execvp("wg", args);
        exit(1);
    }

    // Parent
    close(pipefd[1]);
    std::string wg_output;
    char buffer[128];
    ssize_t bytes;
    while ((bytes = read(pipefd[0], buffer, sizeof(buffer) - 1)) > 0) {
        buffer[bytes] = '\0';
        wg_output += buffer;
    }
    close(pipefd[0]);
    waitpid(pid, NULL, 0);

    // Parse wg show tun0 latest-handshakes output (Peer \t Timestamp)
    uint64_t latest = 0;
    size_t tab_pos = wg_output.find('\t');
    if (tab_pos != std::string::npos) {
        try {
            latest = std::stoull(wg_output.substr(tab_pos + 1));
        } catch (...) {}
    }

    if (latest == 0) {
        std::cerr << "ERROR: tunnel mesh not established\n";
        return false;
    }

    auto now = std::chrono::duration_cast<std::chrono::seconds>(std::chrono::system_clock::now().time_since_epoch()).count();
    auto age = now - latest;

    if (age > 180) {
        std::cerr << "ERROR: compass handshake stale (" << age << "s)\n";
        return false;
    }

    std::cout << "tunnel mesh verified (handshake " << age << "s ago)\n";
    return true;
}

void verify_key_not_on_disk() {
    struct stat st;
    if (stat("/opt/apache/engine/vault.key", &st) == 0) {
        std::cout << "WARNING: key found on persistent disk, removing\n";
        unlink("/opt/apache/engine/vault.key");
    }
    if (stat("/tmp/vault.key", &st) == 0) {
        std::cout << "WARNING: key found on persistent disk, removing\n";
        unlink("/tmp/vault.key");
    }
}

size_t curlWriteCb(void* contents, size_t size, size_t nmemb, std::string* userp) {
    userp->append((char*)contents, size * nmemb);
    return size * nmemb;
}

bool fetch_key_from_shield() {
    struct stat st;
    if (stat(KEY_PATH.c_str(), &st) == 0) {
        std::cout << "Vault key already in RAM\n";
        return true;
    }

    node669427::CryptoEngine crypto;
    auto now = std::chrono::duration_cast<std::chrono::seconds>(std::chrono::system_clock::now().time_since_epoch()).count();
    std::string timestamp = std::to_string(now);
    std::string nonce = crypto.generateNonce(16);
    std::string payload = "FETCH_VAULT_KEY|" + timestamp + "|" + nonce;

    std::ifstream kf(HMAC_KEY_FILE);
    if (!kf.is_open()) {
        std::cerr << "ERROR: HMAC key not available\n";
        return false;
    }
    std::string hmac_key((std::istreambuf_iterator<char>(kf)), std::istreambuf_iterator<char>());
    while (!hmac_key.empty() && (hmac_key.back() == '\n' || hmac_key.back() == '\r')) hmac_key.pop_back();

    crypto.loadKeyFromBuffer(hmac_key);
    std::string hmac_sig = crypto.hmacSha256(payload);

    // Convert HMAC to Base64 (to match openssl dgst -binary | base64)
    std::string hmac_b64;
    {
        unsigned char const* bytes_to_encode = reinterpret_cast<const unsigned char*>(hmac_sig.data());
        size_t in_len = hmac_sig.length();
        size_t b64_len = sodium_base64_ENCODED_LEN(in_len, sodium_base64_VARIANT_ORIGINAL);
        char* b64_out = new char[b64_len];
        sodium_bin2base64(b64_out, b64_len, bytes_to_encode, in_len, sodium_base64_VARIANT_ORIGINAL);
        hmac_b64 = b64_out;
        delete[] b64_out;
    }

    CURL* curl = curl_easy_init();
    if (!curl) return false;

    std::string response;
    struct curl_slist* headers = NULL;
    headers = curl_slist_append(headers, ("x-engine-node: " + NODE_ID).c_str());
    headers = curl_slist_append(headers, ("x-engine-ts: " + timestamp).c_str());
    headers = curl_slist_append(headers, ("x-engine-nonce: " + nonce).c_str());
    headers = curl_slist_append(headers, ("x-engine-sig: " + hmac_b64).c_str());

    std::string url = "http://" + RELAY_IP + ":" + RELAY_PORT + "/vault/key";

    curl_easy_setopt(curl, CURLOPT_URL, url.c_str());
    curl_easy_setopt(curl, CURLOPT_HTTPHEADER, headers);
    curl_easy_setopt(curl, CURLOPT_WRITEFUNCTION, curlWriteCb);
    curl_easy_setopt(curl, CURLOPT_WRITEDATA, &response);
    curl_easy_setopt(curl, CURLOPT_TIMEOUT, 10L);

    CURLcode res = curl_easy_perform(curl);
    curl_slist_free_all(headers);
    curl_easy_cleanup(curl);

    if (res != CURLE_OK || response.empty() || response == "DENIED") {
        std::cerr << "ERROR: relay denied key fetch\n";
        return false;
    }

    std::ofstream out(KEY_PATH);
    if (!out.is_open()) {
        std::cerr << "ERROR: cannot write to " << KEY_PATH << "\n";
        return false;
    }
    out << response;
    out.close();

    chmod(KEY_PATH.c_str(), 0600);
    struct passwd *pw = getpwnam("mail-svc");
    if (pw != NULL) {
        chown(KEY_PATH.c_str(), pw->pw_uid, pw->pw_gid);
    }

    std::cout << "Vault key loaded to RAM (" << KEY_PATH << ")\n";
    return true;
}

void start_engine(const std::string& script_dir) {
    std::cout << "Starting engine...\n";
    std::string engine_path = script_dir + "/engine";
    
    char* args[] = {(char*)engine_path.c_str(), (char*)"--start-engine", NULL};
    execv(engine_path.c_str(), args);
    std::cerr << "ERROR: engine binary not found or could not be executed\n";
    exit(1);
}

int main(int argc, char* argv[]) {
    if (sodium_init() < 0) {
        std::cerr << "Failed to initialize libsodium\n";
        return 1;
    }

    std::string cmd = (argc > 1) ? argv[1] : "";

    std::string script_dir = ".";
    char buffer[4096];
    ssize_t len = readlink("/proc/self/exe", buffer, sizeof(buffer)-1);
    if (len != -1) {
        buffer[len] = '\0';
        std::string full_path = buffer;
        size_t last_slash = full_path.find_last_of('/');
        if (last_slash != std::string::npos) {
            script_dir = full_path.substr(0, last_slash);
        }
    }

    if (cmd == "--start-engine") {
        if (!check_tun_mesh()) return 1;
        verify_key_not_on_disk();
        if (!fetch_key_from_shield()) return 1;
        start_engine(script_dir);
    } else if (cmd == "--fetch-only") {
        if (!check_tun_mesh()) return 1;
        verify_key_not_on_disk();
        if (!fetch_key_from_shield()) return 1;
    } else if (cmd == "--verify-tunnel") {
        if (!check_tun_mesh()) return 1;
    } else {
        std::cerr << "Usage: " << argv[0] << " [--start-engine|--fetch-only|--verify-tunnel]\n";
        return 1;
    }

    return 0;
}
