#include <doctest/doctest.h>

#include <vector>

#include "packeteer/l7/tls.hpp"

using namespace packeteer::net;

namespace {

void append_be16(std::vector<unsigned char>& out, std::uint16_t v) {
    out.push_back(static_cast<unsigned char>(v >> 8));
    out.push_back(static_cast<unsigned char>(v & 0xFF));
}

// Builds a real, well-formed TLS record containing a ClientHello with
// (optionally) a single SNI host_name extension. Every length field is
// computed from the actual bytes assembled, not hand-counted - the
// same lesson from this session's earlier UDP-length test typo.
//
// `corrupt_sni_ext_len`, when set, writes an oversized SNI extension
// length instead of the real one (computed here, not via post-hoc
// offset math into the finished buffer - equally fragile).
std::vector<unsigned char> build_client_hello(const std::string& sni,
                                               bool corrupt_sni_ext_len = false,
                                               std::vector<std::string> alpn_protocols = {}) {
    std::vector<unsigned char> body;
    body.push_back(0x03);
    body.push_back(0x03);                       // client_version: TLS 1.2
    body.insert(body.end(), 32, 0x00);           // random
    body.push_back(0x00);                        // session_id length: 0
    append_be16(body, 2);                         // cipher_suites length
    body.push_back(0x00);
    body.push_back(0x2F);                         // one arbitrary cipher suite
    body.push_back(0x01);                         // compression_methods length: 1
    body.push_back(0x00);                          // null compression

    std::vector<unsigned char> extensions;
    if (!sni.empty()) {
        std::vector<unsigned char> server_name_list;
        server_name_list.push_back(0x00);  // name_type: host_name
        append_be16(server_name_list, static_cast<std::uint16_t>(sni.size()));
        server_name_list.insert(server_name_list.end(), sni.begin(), sni.end());

        std::vector<unsigned char> sni_ext_data;
        append_be16(sni_ext_data, static_cast<std::uint16_t>(server_name_list.size()));
        sni_ext_data.insert(sni_ext_data.end(), server_name_list.begin(), server_name_list.end());

        append_be16(extensions, kTlsExtensionServerName);
        std::uint16_t ext_len = corrupt_sni_ext_len
                                     ? static_cast<std::uint16_t>(0xFFFF)
                                     : static_cast<std::uint16_t>(sni_ext_data.size());
        append_be16(extensions, ext_len);
        extensions.insert(extensions.end(), sni_ext_data.begin(), sni_ext_data.end());
    }
    if (!alpn_protocols.empty()) {
        std::vector<unsigned char> protocol_list;
        for (const auto& proto : alpn_protocols) {
            protocol_list.push_back(static_cast<unsigned char>(proto.size()));
            protocol_list.insert(protocol_list.end(), proto.begin(), proto.end());
        }
        std::vector<unsigned char> alpn_ext_data;
        append_be16(alpn_ext_data, static_cast<std::uint16_t>(protocol_list.size()));
        alpn_ext_data.insert(alpn_ext_data.end(), protocol_list.begin(), protocol_list.end());

        append_be16(extensions, kTlsExtensionAlpn);
        append_be16(extensions, static_cast<std::uint16_t>(alpn_ext_data.size()));
        extensions.insert(extensions.end(), alpn_ext_data.begin(), alpn_ext_data.end());
    }
    append_be16(body, static_cast<std::uint16_t>(extensions.size()));
    body.insert(body.end(), extensions.begin(), extensions.end());

    std::vector<unsigned char> handshake;
    handshake.push_back(kTlsHandshakeTypeClientHello);
    std::uint32_t hs_len = static_cast<std::uint32_t>(body.size());
    handshake.push_back(static_cast<unsigned char>((hs_len >> 16) & 0xFF));
    handshake.push_back(static_cast<unsigned char>((hs_len >> 8) & 0xFF));
    handshake.push_back(static_cast<unsigned char>(hs_len & 0xFF));
    handshake.insert(handshake.end(), body.begin(), body.end());

    std::vector<unsigned char> record;
    record.push_back(kTlsContentTypeHandshake);
    record.push_back(0x03);
    record.push_back(0x01);  // record-layer version (legacy compat value)
    append_be16(record, static_cast<std::uint16_t>(handshake.size()));
    record.insert(record.end(), handshake.begin(), handshake.end());

    return record;
}

// Builds a real, well-formed TLS record containing a ServerHello.
// `negotiated_version`, when set, adds a supported_versions extension
// carrying that single 2-byte version - how a TLS 1.3 ServerHello
// actually signals its real version, since legacy_version is always
// 0x0303 there for middlebox compatibility.
std::vector<unsigned char> build_server_hello(std::uint16_t cipher_suite,
                                               std::optional<std::uint16_t> negotiated_version = {}) {
    std::vector<unsigned char> body;
    append_be16(body, 0x0303);            // legacy_version: always TLS 1.2 on the wire
    body.insert(body.end(), 32, 0x00);    // random
    body.push_back(0x00);                  // legacy_session_id_echo length: 0
    append_be16(body, cipher_suite);
    body.push_back(0x00);                  // legacy_compression_method: null

    std::vector<unsigned char> extensions;
    if (negotiated_version) {
        append_be16(extensions, kTlsExtensionSupportedVersions);
        append_be16(extensions, 2);
        append_be16(extensions, *negotiated_version);
    }
    append_be16(body, static_cast<std::uint16_t>(extensions.size()));
    body.insert(body.end(), extensions.begin(), extensions.end());

    std::vector<unsigned char> handshake;
    handshake.push_back(kTlsHandshakeTypeServerHello);
    std::uint32_t hs_len = static_cast<std::uint32_t>(body.size());
    handshake.push_back(static_cast<unsigned char>((hs_len >> 16) & 0xFF));
    handshake.push_back(static_cast<unsigned char>((hs_len >> 8) & 0xFF));
    handshake.push_back(static_cast<unsigned char>(hs_len & 0xFF));
    handshake.insert(handshake.end(), body.begin(), body.end());

    std::vector<unsigned char> record;
    record.push_back(kTlsContentTypeHandshake);
    record.push_back(0x03);
    record.push_back(0x03);
    append_be16(record, static_cast<std::uint16_t>(handshake.size()));
    record.insert(record.end(), handshake.begin(), handshake.end());

    return record;
}

}  // namespace

TEST_CASE("parse_tls_client_hello extracts a real SNI extension") {
    auto record = build_client_hello("example.com");
    auto hello = parse_tls_client_hello(record);
    REQUIRE(hello.has_value());
    REQUIRE(hello->server_name.has_value());
    CHECK(*hello->server_name == "example.com");
}

TEST_CASE("parse_tls_client_hello succeeds with no SNI when there's no extensions block") {
    auto record = build_client_hello("");
    auto hello = parse_tls_client_hello(record);
    REQUIRE(hello.has_value());
    CHECK_FALSE(hello->server_name.has_value());
}

TEST_CASE("parse_tls_client_hello rejects a non-Handshake record") {
    auto record = build_client_hello("example.com");
    record[0] = 0x17;  // application_data, not handshake
    CHECK_FALSE(parse_tls_client_hello(record).has_value());
}

TEST_CASE("parse_tls_client_hello rejects a non-ClientHello handshake type") {
    auto record = build_client_hello("example.com");
    record[5] = 0x02;  // ServerHello, not ClientHello
    CHECK_FALSE(parse_tls_client_hello(record).has_value());
}

TEST_CASE("parse_tls_client_hello rejects a truncated record") {
    auto record = build_client_hello("example.com");
    record.resize(record.size() - 5);  // claims more than it has
    CHECK_FALSE(parse_tls_client_hello(record).has_value());
}

TEST_CASE("parse_tls_client_hello rejects a buffer shorter than the record header") {
    std::vector<unsigned char> bytes(4, 0);
    CHECK_FALSE(parse_tls_client_hello(bytes).has_value());
}

TEST_CASE("parse_tls_client_hello stops gracefully on a malformed extension length") {
    auto record = build_client_hello("example.com", /*corrupt_sni_ext_len=*/true);
    auto hello = parse_tls_client_hello(record);
    REQUIRE(hello.has_value());  // still a structurally valid ClientHello otherwise
    CHECK_FALSE(hello->server_name.has_value());  // SNI extension was malformed, so skipped
}

TEST_CASE("TlsSniDissector claims port 443 and its summary matches parse_tls_client_hello") {
    TlsSniDissector dissector;
    CHECK(dissector.port() == kTlsPort);

    auto record = build_client_hello("packeteer.test");
    auto summary = dissector.summarize(record);
    REQUIRE(summary.has_value());
    CHECK(summary->substr(0, 3) == "TLS");
    CHECK(summary->find("SNI=packeteer.test") != std::string::npos);
}

TEST_CASE("parse_tls_client_hello extracts ALPN protocols in order") {
    auto record = build_client_hello("example.com", false, {"h2", "http/1.1"});
    auto hello = parse_tls_client_hello(record);
    REQUIRE(hello.has_value());
    REQUIRE(hello->alpn_protocols.has_value());
    REQUIRE(hello->alpn_protocols->size() == 2);
    CHECK((*hello->alpn_protocols)[0] == "h2");
    CHECK((*hello->alpn_protocols)[1] == "http/1.1");
}

TEST_CASE("parse_tls_client_hello leaves alpn_protocols unset when the extension is absent") {
    auto record = build_client_hello("example.com");
    auto hello = parse_tls_client_hello(record);
    REQUIRE(hello.has_value());
    CHECK_FALSE(hello->alpn_protocols.has_value());
}

TEST_CASE("parse_tls_server_hello extracts the negotiated cipher suite") {
    auto record = build_server_hello(0x1301);  // TLS_AES_128_GCM_SHA256
    auto hello = parse_tls_server_hello(record);
    REQUIRE(hello.has_value());
    CHECK(hello->cipher_suite == 0x1301);
}

TEST_CASE("parse_tls_server_hello uses legacy_version when there's no supported_versions extension") {
    auto record = build_server_hello(0xC02F);  // a TLS 1.2-era cipher suite
    auto hello = parse_tls_server_hello(record);
    REQUIRE(hello.has_value());
    CHECK(hello->negotiated_version == 0x0303);  // TLS 1.2
}

TEST_CASE("parse_tls_server_hello prefers supported_versions over legacy_version for TLS 1.3") {
    auto record = build_server_hello(0x1301, std::uint16_t{0x0304});
    auto hello = parse_tls_server_hello(record);
    REQUIRE(hello.has_value());
    CHECK(hello->negotiated_version == 0x0304);  // TLS 1.3, not the legacy 0x0303 on the wire
}

TEST_CASE("parse_tls_server_hello rejects a ClientHello handshake type") {
    auto record = build_client_hello("example.com");
    CHECK_FALSE(parse_tls_server_hello(record).has_value());
}

TEST_CASE("tls_cipher_suite_name names all five TLS 1.3 suites and falls back to hex otherwise") {
    CHECK(tls_cipher_suite_name(0x1301) == "TLS_AES_128_GCM_SHA256");
    CHECK(tls_cipher_suite_name(0x1305) == "TLS_AES_128_CCM_8_SHA256");
    CHECK(tls_cipher_suite_name(0xC02F) == "0xc02f");
}

TEST_CASE("TlsSniDissector formats a ServerHello with version and cipher") {
    TlsSniDissector dissector;
    auto record = build_server_hello(0x1301, std::uint16_t{0x0304});
    auto summary = dissector.summarize(record);
    REQUIRE(summary.has_value());
    CHECK(*summary == "TLS ServerHello version=TLS1.3 cipher=TLS_AES_128_GCM_SHA256");
}