// tls13_selftest.cpp
#include <vector>
#include <cstdint>
#include <cstring>
#include <iostream>
#include <stdexcept>

#include "../src/socket.h"
#include "../src/crypto/aes.h"
#include "../src/exception.h"

static void dumpHex(const char* tag, const uint8_t* p, size_t n, size_t max=96) {
    std::cout << tag << " (" << n << " bytes)\n";
    size_t m = (n > max) ? max : n;
    for (size_t i=0;i<m;i++) {
        printf("%02x ", p[i]);
        if ((i % 16) == 15) printf("\n");
    }
    if (m % 16) printf("\n");
    if (n > max) std::cout << "...\n";
}

namespace netplus {

// In ssl-Klasse hinzufügen:
//   friend class TestSSL;

class TestSSL {
public:
    explicit TestSSL(netplus::ssl& s) : S(s) {}

    void run_tls13_record_loopback() {
        std::cout << "\n=== TLS1.3 AEAD Record Loopback Selftest ===\n";

        // ---- deterministic test key
        const key128 KEY = {
            0x2b,0x7e,0x15,0x16,0x28,0xae,0xd2,0xa6,
            0xab,0xf7,0x15,0x88,0x09,0xcf,0x4f,0x3c
        };

        uint8_t iv[12];
        for (int i=0;i<12;i++) iv[i] = uint8_t(0xA0 + i);

        // install HS keys (loopback: same key/iv both dirs)
        S._tls.aes13_hs_send.reset(new aes128(KEY));
        S._tls.aes13_hs_recv.reset(new aes128(KEY));
        std::memcpy(S._tls.tls13_hs_iv_s2c, iv, 12);
        std::memcpy(S._tls.tls13_hs_iv_c2s, iv, 12);
        S._tls.tls13_hs_send_seq = 0;
        S._tls.tls13_hs_recv_seq = 0;

        // record #1
        std::vector<uint8_t> msg1 = { 0x08,0x00,0x00,0x02,0x00,0x00 };
        auto rec1 = send_record_capture(/*inner_type=*/0x16, msg1, /*handshake_keys=*/true);
        dumpHex("Captured rec1", rec1.data(), rec1.size());

        auto dec1 = recv_record_parse(rec1, /*handshake_keys=*/true);
        if (dec1.inner_type != 0x16) throw std::runtime_error("rec1 inner_type mismatch");
        if (dec1.plain != msg1)      throw std::runtime_error("rec1 plaintext mismatch");

        if (S._tls.tls13_hs_send_seq != 1) throw std::runtime_error("send seq not 1 after rec1");
        if (S._tls.tls13_hs_recv_seq != 1) throw std::runtime_error("recv seq not 1 after rec1");
        std::cout << "✅ PASS: rec1 loopback OK\n";

        // record #2 (anders lang => triggert GHASH partial differently)
        std::vector<uint8_t> msg2 = { 0x14,0x00,0x00,0x00 };
        auto rec2 = send_record_capture(0x16, msg2, true);
        dumpHex("Captured rec2", rec2.data(), rec2.size());

        auto dec2 = recv_record_parse(rec2, true);
        if (dec2.inner_type != 0x16) throw std::runtime_error("rec2 inner_type mismatch");
        if (dec2.plain != msg2)      throw std::runtime_error("rec2 plaintext mismatch");

        if (S._tls.tls13_hs_send_seq != 2) throw std::runtime_error("send seq not 2 after rec2");
        if (S._tls.tls13_hs_recv_seq != 2) throw std::runtime_error("recv seq not 2 after rec2");
        std::cout << "✅ PASS: rec2 loopback OK\n";

        // tag corruption test
        auto rec3 = send_record_capture(0x16, msg1, true);
        rec3.back() ^= 0x01;

        bool rejected = false;
        try {
            (void)recv_record_parse(rec3, true);
        } catch (...) {
            rejected = true;
        }
        if (!rejected) throw std::runtime_error("corrupted tag was NOT rejected");
        std::cout << "✅ PASS: corrupted tag rejected\n";
    }

private:
    netplus::ssl& S;

    struct Dec {
        std::vector<uint8_t> plain;
        uint8_t inner_type = 0;
    };

    void clear_send_queue() {
        S._tls.send_queue.clear();
    }

    std::vector<uint8_t> drain_send_queue_bytes() {
        std::vector<uint8_t> out;
        for (auto& chunk : S._tls.send_queue) {
            out.insert(out.end(), chunk.begin(), chunk.end());
        }
        S._tls.send_queue.clear();
        return out;
    }

    std::vector<uint8_t> send_record_capture(uint8_t inner_type,
                                             const std::vector<uint8_t>& content,
                                             bool handshake_keys)
    {
        clear_send_queue();
        S._tls.tls13_send_record(inner_type, content, handshake_keys);
        auto bytes = drain_send_queue_bytes();
        if (bytes.size() < 5 + 16) throw std::runtime_error("captured record too small");
        return bytes;
    }

    Dec recv_record_parse(const std::vector<uint8_t>& rec, bool handshake_keys) {
        if (rec.size() < 5 + 16) throw std::runtime_error("record too short");

        uint8_t  rtype = rec[0];
        uint16_t ver   = (uint16_t(rec[1])<<8) | rec[2];
        uint16_t rlen  = (uint16_t(rec[3])<<8) | rec[4];
        if (5 + size_t(rlen) != rec.size()) throw std::runtime_error("record length mismatch");

        Dec d;
        std::vector<uint8_t> plain;
        uint8_t inner = 0;

        S._tls.tls13_recv_record(
            rtype, ver,
            rec.data() + 5, rlen,
            plain, inner,
            handshake_keys
        );

        d.plain = std::move(plain);
        d.inner_type = inner;
        return d;
    }
};

} // namespace netplus

int main() {
    try {
        // Erstelle eine SSL-Instanz mit einem Zertifikat-Map für SNI-Support
        netplus::x509cert cert;
        std::map<std::string, netplus::ssl::CertificateBundle> certs;
        netplus::ssl::CertificateBundle bundle;
        bundle.cert = cert;
        bundle.privateKeyDer = {};
        certs["localhost"] = bundle;
        
        netplus::ssl ssl(certs);

        netplus::TestSSL T(ssl);
        T.run_tls13_record_loopback();

        std::cout << "\n✅ ALL TLS13 SELFTESTS PASSED\n";
        return 0;
    } catch (const netplus::NetException& e) {
        std::cerr << "\n❌ TLS13 SELFTEST FAILED: " << e.what() << "\n";
        return 1;
    }
}
