// tls_roundtrip_sha256_test.cpp
//
// Integrity test: fill a file with random bytes, hash it with SHA-256,
// then bounce its contents back and forth several times over a TLS
// echo server (real handshake + record layer, in-process), and finally
// verify the SHA-256 checksum is still identical to the original.

#include <iostream>
#include <string>
#include <vector>
#include <map>
#include <fstream>
#include <cstdio>
#include <cstring>
#include <thread>
#include <atomic>
#include <chrono>
#include <stdexcept>

#include "connection.h"
#include "eventapi.h"
#include "socket.h"
#include "exception.h"
#include "random.h"
#include "crypto/sha.h"

#include "https_certs.h"
#include "https_ca_cert.h"

using namespace netplus;

static int g_passed = 0, g_failed = 0;

static void check(bool ok, const std::string& name) {
    if (ok) { std::cout << "  PASS: " << name << std::endl; g_passed++; }
    else    { std::cout << "  FAIL: " << name << std::endl; g_failed++; }
}

// ============================================================================
// File / hash helpers
// ============================================================================

static std::string toHex(const std::vector<uint8_t>& v) {
    static const char* hexd = "0123456789abcdef";
    std::string s;
    s.reserve(v.size() * 2);
    for (uint8_t b : v) {
        s.push_back(hexd[b >> 4]);
        s.push_back(hexd[b & 0x0F]);
    }
    return s;
}

static std::vector<uint8_t> writeRandomFile(const std::string& path, size_t len) {
    std::vector<uint8_t> buf(len);
    fillRandomBytes(buf.data(), buf.size());

    std::ofstream out(path, std::ios::binary | std::ios::trunc);
    if (!out) throw std::runtime_error("cannot open " + path + " for writing");
    out.write(reinterpret_cast<const char*>(buf.data()), static_cast<std::streamsize>(buf.size()));
    if (!out) throw std::runtime_error("failed writing " + path);
    return buf;
}

static std::vector<uint8_t> readFile(const std::string& path) {
    std::ifstream in(path, std::ios::binary | std::ios::ate);
    if (!in) throw std::runtime_error("cannot open " + path + " for reading");
    std::streamsize sz = in.tellg();
    in.seekg(0, std::ios::beg);
    std::vector<uint8_t> buf(static_cast<size_t>(sz));
    if (sz > 0 && !in.read(reinterpret_cast<char*>(buf.data()), sz))
        throw std::runtime_error("failed reading " + path);
    return buf;
}

struct TempFileGuard {
    std::string path;
    explicit TempFileGuard(std::string p) : path(std::move(p)) {}
    ~TempFileGuard() { std::remove(path.c_str()); }
};

// ============================================================================
// TLS echo server
// ============================================================================

class EchoServer : public event {
public:
    EchoServer(std::vector<netplus::socket*> socks, int timeout = 500)
        : event(socks, timeout) {}

    void RequestEvent(con& curcon, const int tid, ULONG_PTR args) override {
        if (!curcon.RecvData.empty()) {
            curcon.SendData.append(curcon.RecvData.data(), curcon.RecvData.size());
            curcon.RecvData.clear();
        }
    }
    void ResponseEvent(con&, const int, ULONG_PTR) override {}
    void ConnectEvent(con&, const int, ULONG_PTR) override {}
    void DisconnectEvent(con&, const int, ULONG_PTR) override {}
    void CreateConnection(std::shared_ptr<con>& res) override {
        res = std::make_shared<con>(this);
    }
};

static std::atomic<bool> g_tls_ready(false);

static void waitReady(std::atomic<bool>& ready) {
    while (!ready.load())
        std::this_thread::sleep_for(std::chrono::milliseconds(10));
    std::this_thread::sleep_for(std::chrono::milliseconds(200));
}

static void runTlsServer(std::map<std::string, ssl::CertificateBundle>& certs, int port) {
    try {
        ssl serverSock(certs, "127.0.0.1", port, 64, -1);
        EchoServer srv({&serverSock});
        g_tls_ready.store(true);
        srv.runEventloop();
    } catch (NetException& e) {
        std::cerr << "[TLS server] NetException: " << e.what() << std::endl;
        g_tls_ready.store(true);
    } catch (std::exception& e) {
        std::cerr << "[TLS server] Exception: " << e.what() << std::endl;
        g_tls_ready.store(true);
    }
}

// ============================================================================
// TLS round trip: connect() blocks until the handshake is done, then drive
// sendData()/recvData() directly (blocking socket) — retrying on
// NetException::Note, which the TLS layer throws when a read produced only
// non-application records (e.g. a NewSessionTicket) and no app data yet.
// ============================================================================

static std::vector<uint8_t> tlsRoundTrip(const std::vector<uint8_t>& payload,
                                          std::map<std::string, ssl::CertificateBundle>& certs,
                                          const std::string& host, int port) {
    ssl client(certs);
    client.connect(host, port);

    size_t sent_total = 0;
    while (sent_total < payload.size()) {
        size_t chunk = std::min<size_t>(payload.size() - sent_total, 16384);
        buffer snd(reinterpret_cast<const char*>(payload.data() + sent_total), chunk);
        size_t sent;
        try {
            sent = client.sendData(snd, 0);
        } catch (NetException& e) {
            if (e.getErrorType() == NetException::Note) continue;
            throw;
        }
        if (sent == 0) throw std::runtime_error("tlsRoundTrip: sendData stalled");
        sent_total += sent;
    }

    std::vector<char> scratch(65536);
    std::vector<uint8_t> received(payload.size());
    size_t recv_total = 0;
    while (recv_total < received.size()) {
        size_t want = std::min(scratch.size(), received.size() - recv_total);
        buffer rcv(scratch.data(), want);
        size_t got;
        try {
            got = client.recvData(rcv, 0);
        } catch (NetException& e) {
            if (e.getErrorType() == NetException::Note) continue;
            throw;
        }
        if (got == 0) throw std::runtime_error("tlsRoundTrip: connection closed early");
        std::memcpy(received.data() + recv_total, scratch.data(), got);
        recv_total += got;
    }
    return received;
}

// ============================================================================
// main
// ============================================================================

int main() {
    std::cout << "=== TLS Round-Trip SHA-256 Integrity Test ===" << std::endl;

    const std::string kTempFile = "tls_roundtrip_sha256_test.bin";
    const size_t kFileSize = 4096 * 1024;
    const int kRounds = 5;
    const int kTlsPort = 19544;

    TempFileGuard tempGuard(kTempFile);
    int rc = 0;

    try {
        // 1) Fill a file with random bytes and record its SHA-256
        std::vector<uint8_t> original = writeRandomFile(kTempFile, kFileSize);
        std::vector<uint8_t> originalHash = sha256_hash(original);
        std::cout << "Generated " << kFileSize << " random bytes, sha256=" << toHex(originalHash) << std::endl;

        check(sha256_hash(readFile(kTempFile)) == originalHash,
              "file-on-disk sha256 matches generated buffer");

        // 2) Certificate/key material for the TLS server
        x509cert cert;
        if (!cert.loadFromBuffer(test_cert_der)) {
            std::cerr << "Failed to load certificate" << std::endl;
            return 1;
        }
        std::map<std::string, ssl::CertificateBundle> certs;
        ssl::CertificateBundle bundle;
        bundle.cert = cert;
        bundle.privateKeyDer = std::vector<uint8_t>(test_key_der.begin(), test_key_der.end());
        bundle.rsa_key = rsa(bundle.privateKeyDer);
        bundle.chain.push_back(std::vector<uint8_t>(MKCERT_ROOT_CA_DER,
            MKCERT_ROOT_CA_DER + MKCERT_ROOT_CA_DER_LEN));
        certs["localhost"] = bundle;
        certs["127.0.0.1"] = bundle;

        // 3) Start the TLS echo server
        std::thread tlsServerThread(runTlsServer, std::ref(certs), kTlsPort);
        waitReady(g_tls_ready);

        // 4) Bounce the file content back and forth over TLS a few times
        std::vector<uint8_t> current = original;
        for (int round = 1; round <= kRounds; ++round) {
            std::cout << "\n--- TLS round " << round << "/" << kRounds << " ---" << std::endl;
            current = tlsRoundTrip(current, certs, "127.0.0.1", kTlsPort);
            check(current == original, "TLS round trip preserves bytes");
            check(sha256_hash(current) == originalHash, "TLS round trip preserves sha256");
        }

        event::Running = false;
        tlsServerThread.join();

        // 5) Final SHA-256 check against the original file content
        std::vector<uint8_t> finalHash = sha256_hash(current);
        check(finalHash == originalHash, "final sha256 matches original after all round trips");
        std::cout << "Final sha256=" << toHex(finalHash) << std::endl;

    } catch (NetException& e) {
        std::cerr << "NetException: " << e.what() << std::endl;
        rc = 1;
    } catch (std::exception& e) {
        std::cerr << "Exception: " << e.what() << std::endl;
        rc = 1;
    }

    std::cout << "\n==============================" << std::endl;
    std::cout << "Results: " << g_passed << " passed, " << g_failed << " failed" << std::endl;
    std::cout << "==============================" << std::endl;

    return (rc != 0 || g_failed > 0) ? 1 : 0;
}
