#pragma once

#include <cstdint>
#include <optional>
#include <span>
#include <string>

#include "wireframe/byteio.hpp"
#include "wireframe/l7/dissector.hpp"

// TLS ClientHello -> SNI extension parsing. Most web traffic is TLS
// today, so HTTP alone covers a shrinking fraction of it - SNI is what
// makes a packet analyzer useful against that traffic without
// decrypting anything: the server name is sent in cleartext in the
// ClientHello, before any encryption starts, in every TLS version this
// parses (the ClientHello/extension wire format hasn't changed across
// versions - only what happens after it has).
//
// Same scope as the other L7 dissectors: single-segment, best-effort.
// A ClientHello padded across multiple TCP segments (large cookie/PSK
// extensions, unusual but possible) is only partially visible here.
// Every length field is bounds-checked against what's actually left in
// the buffer before use - this is exactly the kind of nested,
// attacker-influenced TLV structure the project's decoders are meant
// to get right.
namespace wireframe::net {

inline constexpr std::uint16_t kTlsPort = 443;
inline constexpr std::uint8_t kTlsContentTypeHandshake = 0x16;
inline constexpr std::uint8_t kTlsHandshakeTypeClientHello = 0x01;
inline constexpr std::uint16_t kTlsExtensionServerName = 0x0000;

struct TlsClientHello {
    std::optional<std::string> server_name;  // SNI, if the extension was present and well-formed
};

inline std::optional<TlsClientHello> parse_tls_client_hello(std::span<const unsigned char> bytes) {
    // Record header: ContentType(1) ProtocolVersion(2) Length(2)
    if (bytes.size() < 5) return std::nullopt;
    if (bytes[0] != kTlsContentTypeHandshake) return std::nullopt;
    std::uint16_t record_len = read_be16(bytes, 3);
    if (bytes.size() < static_cast<std::size_t>(5) + record_len) return std::nullopt;

    std::span<const unsigned char> handshake = bytes.subspan(5);

    // Handshake header: HandshakeType(1) Length(3, 24-bit BE)
    if (handshake.size() < 4) return std::nullopt;
    if (handshake[0] != kTlsHandshakeTypeClientHello) return std::nullopt;
    std::uint32_t hs_len = (static_cast<std::uint32_t>(handshake[1]) << 16) |
                            (static_cast<std::uint32_t>(handshake[2]) << 8) |
                            static_cast<std::uint32_t>(handshake[3]);

    std::span<const unsigned char> body = handshake.subspan(4);
    if (body.size() < hs_len) return std::nullopt;
    body = body.first(hs_len);  // never read past the declared handshake body

    std::size_t offset = 0;

    // client_version(2) + random(32)
    if (body.size() < offset + 34) return std::nullopt;
    offset += 34;

    // legacy_session_id: length(1) + data
    if (body.size() < offset + 1) return std::nullopt;
    std::uint8_t session_id_len = body[offset];
    offset += 1;
    if (body.size() < offset + session_id_len) return std::nullopt;
    offset += session_id_len;

    // cipher_suites: length(2) + data
    if (body.size() < offset + 2) return std::nullopt;
    std::uint16_t cipher_suites_len = read_be16(body, offset);
    offset += 2;
    if (body.size() < static_cast<std::size_t>(offset) + cipher_suites_len) return std::nullopt;
    offset += cipher_suites_len;

    // legacy_compression_methods: length(1) + data
    if (body.size() < offset + 1) return std::nullopt;
    std::uint8_t compression_len = body[offset];
    offset += 1;
    if (body.size() < offset + compression_len) return std::nullopt;
    offset += compression_len;

    TlsClientHello hello;
    if (offset == body.size()) return hello;  // no extensions block: no SNI, still a valid hello

    // extensions: length(2) + data
    if (body.size() < offset + 2) return std::nullopt;
    std::uint16_t extensions_len = read_be16(body, offset);
    offset += 2;
    if (body.size() < static_cast<std::size_t>(offset) + extensions_len) return std::nullopt;
    std::size_t extensions_end = offset + extensions_len;

    while (offset + 4 <= extensions_end) {
        std::uint16_t ext_type = read_be16(body, offset);
        std::uint16_t ext_len = read_be16(body, offset + 2);
        std::size_t ext_data_start = offset + 4;
        std::size_t ext_data_end = ext_data_start + ext_len;
        if (ext_data_end > extensions_end) break;  // malformed: stop, keep what we have

        if (ext_type == kTlsExtensionServerName && ext_len >= 2) {
            // ServerNameList: list_len(2) + entries; only the first
            // entry is used, matching every real client's behavior of
            // sending exactly one host_name entry.
            std::uint16_t list_len = read_be16(body, ext_data_start);
            std::size_t list_start = ext_data_start + 2;
            std::size_t list_end = list_start + list_len;
            if (list_end <= ext_data_end && list_start + 3 <= list_end) {
                std::uint8_t name_type = body[list_start];
                std::uint16_t name_len = read_be16(body, list_start + 1);
                std::size_t name_start = list_start + 3;
                if (name_type == 0 && name_start + name_len <= list_end) {
                    hello.server_name = std::string(
                        reinterpret_cast<const char*>(body.data() + name_start), name_len);
                }
            }
        }

        offset = ext_data_end;
    }

    return hello;
}

class TlsSniDissector : public L7Dissector {
public:
    std::uint16_t port() const override { return kTlsPort; }

    std::optional<std::string> summarize(std::span<const unsigned char> payload) const override {
        auto hello = parse_tls_client_hello(payload);
        if (!hello) return std::nullopt;

        std::string out = "TLS ClientHello";
        if (hello->server_name) out += " SNI=" + *hello->server_name;
        return out;
    }
};

}  // namespace wireframe::net