package internal

import (
	"reflect"
	"testing"

	"git.difuse.io/Difuse/Mellaris/analyzer/utils"
)

func buildClientHelloMsg(t *testing.T) *utils.ByteBuffer {
	t.Helper()
	// Bytes taken from the standard TLS test vector, starting after the record
	// header (5 bytes) and handshake header (4 bytes).
	body := []byte{
		0x03, 0x03, // version TLS 1.2
		0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07,
		0x08, 0x09, 0x0a, 0x0b, 0x0c, 0x0d, 0x0e, 0x0f,
		0x10, 0x11, 0x12, 0x13, 0x14, 0x15, 0x16, 0x17,
		0x18, 0x19, 0x1a, 0x1b, 0x1c, 0x1d, 0x1e, 0x1f, // random
		0x00, // session ID length = 0
		0x00, 0x20, // cipher suites length = 32 bytes = 16 suites
		0xcc, 0xa8, 0xcc, 0xa9, 0xc0, 0x2f, 0xc0, 0x30,
		0xc0, 0x2b, 0xc0, 0x2c, 0xc0, 0x13, 0xc0, 0x09,
		0xc0, 0x14, 0xc0, 0x0a, 0x00, 0x9c, 0x00, 0x9d,
		0x00, 0x2f, 0x00, 0x35, 0xc0, 0x12, 0x00, 0x0a, // ciphers
		0x01, // compression methods length = 1
		0x00, // compression method = null
		0x00, 0x58, // extensions length = 88
		// extension: server_name
		0x00, 0x00, // type = server_name
		0x00, 0x18, // length = 24
		0x00, 0x16, // server name list length = 22
		0x00, // name type = hostname
		0x00, 0x13, // name length = 19
		'e', 'x', 'a', 'm', 'p', 'l', 'e', '.', 'u', 'l', 'f', 'h', 'e', 'i', 'm', '.', 'n', 'e', 't',
		// extension: status_request
		0x00, 0x05, // type = status_request
		0x00, 0x05, // length = 5
		0x01, 0x00, 0x00, 0x00, 0x00,
		// extension: supported_groups
		0x00, 0x0a, // type = supported_groups
		0x00, 0x0a, // length = 10
		0x00, 0x08, // list length = 8
		0x00, 0x1d, 0x00, 0x17, 0x00, 0x18, 0x00, 0x19,
		// extension: ec_point_formats
		0x00, 0x0b, // type = ec_point_formats
		0x00, 0x02, // length = 2
		0x01, 0x00,
		// extension: signature_algorithms
		0x00, 0x0d, // type = signature_algorithms
		0x00, 0x12, // length = 18
		0x00, 0x10, // list length = 16
		0x04, 0x01, 0x04, 0x03, 0x05, 0x01, 0x05, 0x03,
		0x06, 0x01, 0x06, 0x03, 0x02, 0x01, 0x02, 0x03,
		// extension: renegotiation_info
		0xff, 0x01, // type = renegotiation_info
		0x00, 0x01, // length = 1
		0x00,
		// extension: extended_master_secret (empty)
		0x00, 0x17, // type = extended_master_secret
		0x00, 0x00, // length = 0
		// extension: session_ticket (empty)
		0x00, 0x23, // type = session_ticket
		0x00, 0x00, // length = 0
	}
	return &utils.ByteBuffer{Buf: body}
}

func TestParseTLSClientHelloMsgData(t *testing.T) {
	chBuf := buildClientHelloMsg(t)
	m := ParseTLSClientHelloMsgData(chBuf)
	if m == nil {
		t.Fatal("ParseTLSClientHelloMsgData returned nil")
	}

	wantVersion := uint16(0x0303)
	if v, ok := m["version"].(uint16); !ok || v != wantVersion {
		t.Errorf("version = %v, want %v", m["version"], wantVersion)
	}

	wantCiphers := []uint16{52392, 52393, 49199, 49200, 49195, 49196, 49171, 49161, 49172, 49162, 156, 157, 47, 53, 49170, 10}
	if c, ok := m["ciphers"].([]uint16); ok {
		if !reflect.DeepEqual(c, wantCiphers) {
			t.Errorf("ciphers = %v, want %v", c, wantCiphers)
		}
	} else {
		t.Errorf("ciphers missing or wrong type: %T", m["ciphers"])
	}

	if sni, ok := m["sni"].(string); !ok || sni != "example.ulfheim.net" {
		t.Errorf("sni = %q, want %q", m["sni"], "example.ulfheim.net")
	}

	if _, ok := m["compression"]; !ok {
		t.Error("compression key missing")
	}

	if _, ok := m["session"]; !ok {
		t.Error("session key missing")
	}

	if _, ok := m["random"]; !ok {
		t.Error("random key missing")
	}
}

func TestParseTLSServerHelloMsgData(t *testing.T) {
	// ServerHello message body (after record header + handshake header)
	body := []byte{
		0x03, 0x03, // version TLS 1.2
		0x70, 0x71, 0x72, 0x73, 0x74, 0x75, 0x76, 0x77,
		0x78, 0x79, 0x7a, 0x7b, 0x7c, 0x7d, 0x7e, 0x7f,
		0x80, 0x81, 0x82, 0x83, 0x84, 0x85, 0x86, 0x87,
		0x88, 0x89, 0x8a, 0x8b, 0x8c, 0x8d, 0x8e, 0x8f, // random
		0x00, // session ID length = 0
		0xc0, 0x13, // cipher suite
		0x00, // compression method
		0x00, 0x05, // extensions length = 5
		// extension: renegotiation_info
		0xff, 0x01, // type
		0x00, 0x01, // length = 1
		0x00,
	}

	m := ParseTLSServerHelloMsgData(&utils.ByteBuffer{Buf: body})
	if m == nil {
		t.Fatal("ParseTLSServerHelloMsgData returned nil")
	}

	wantCipher := uint16(0xc013)
	if c, ok := m["cipher"].(uint16); !ok || c != wantCipher {
		t.Errorf("cipher = %v, want %v", m["cipher"], wantCipher)
	}

	wantCompression := uint8(0)
	if c, ok := m["compression"].(uint8); !ok || c != wantCompression {
		t.Errorf("compression = %v, want %v", m["compression"], wantCompression)
	}
}

func TestParseTLSServerHelloMsgData_NoExtensions(t *testing.T) {
	// ServerHello message without extensions
	body := []byte{
		0x03, 0x03, // version
		0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07,
		0x08, 0x09, 0x0a, 0x0b, 0x0c, 0x0d, 0x0e, 0x0f,
		0x10, 0x11, 0x12, 0x13, 0x14, 0x15, 0x16, 0x17,
		0x18, 0x19, 0x1a, 0x1b, 0x1c, 0x1d, 0x1e, 0x1f, // random
		0x00, // session ID length
		0x00, 0xff, // cipher suite
		0x00, // compression method
	}

	m := ParseTLSServerHelloMsgData(&utils.ByteBuffer{Buf: body})
	if m == nil {
		t.Fatal("ParseTLSServerHelloMsgData returned nil for no-extensions case")
	}
	if _, ok := m["cipher"]; !ok {
		t.Error("cipher key missing")
	}
}

func TestParseTLSClientHelloMsgData_Truncated(t *testing.T) {
	tests := []struct {
		name string
		buf  []byte
	}{
		{"too short for session id", []byte{0x03, 0x03, 0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0a, 0x0b, 0x0c, 0x0d, 0x0e, 0x0f, 0x10, 0x11, 0x12, 0x13, 0x14, 0x15, 0x16, 0x17, 0x18, 0x19, 0x1a, 0x1b, 0x1c, 0x1d, 0x1e, 0x1f}},
		{"odd cipher suites length", []byte{
			0x03, 0x03, // version
			0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0a, 0x0b, 0x0c, 0x0d, 0x0e, 0x0f,
			0x10, 0x11, 0x12, 0x13, 0x14, 0x15, 0x16, 0x17, 0x18, 0x19, 0x1a, 0x1b, 0x1c, 0x1d, 0x1e, 0x1f, // random
			0x00, // session ID length
			0x00, 0x03, // odd cipher suites length
		}},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			m := ParseTLSClientHelloMsgData(&utils.ByteBuffer{Buf: tt.buf})
			if m != nil {
				t.Error("expected nil for truncated input")
			}
		})
	}
}

func TestParseTLSClientHelloMsgData_ECH(t *testing.T) {
	// ClientHello with ECH extension
	body := []byte{
		0x03, 0x03, // version
		0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07,
		0x08, 0x09, 0x0a, 0x0b, 0x0c, 0x0d, 0x0e, 0x0f,
		0x10, 0x11, 0x12, 0x13, 0x14, 0x15, 0x16, 0x17,
		0x18, 0x19, 0x1a, 0x1b, 0x1c, 0x1d, 0x1e, 0x1f, // random
		0x00, // session ID length
		0x00, 0x02, // cipher suites length = 2
		0x13, 0x01, // TLS_AES_128_GCM_SHA256
		0x01, // compression methods length = 1
		0x00, // null compression
		0x00, 0x07, // extensions length = 7
		// ECH extension
		0xfe, 0x0d, // type = encrypted_client_hello
		0x00, 0x03, // length = 3
		0x01, 0x02, 0x03, // some ECH data
	}

	m := ParseTLSClientHelloMsgData(&utils.ByteBuffer{Buf: body})
	if m == nil {
		t.Fatal("ParseTLSClientHelloMsgData returned nil")
	}
	ech, ok := m["ech"].(bool)
	if !ok || !ech {
		t.Errorf("ech = %v, want true", m["ech"])
	}
}

func TestParseTLSClientHelloMsgData_ALPN(t *testing.T) {
	// ClientHello with ALPN extension
	body := []byte{
		0x03, 0x03, // version
		0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07,
		0x08, 0x09, 0x0a, 0x0b, 0x0c, 0x0d, 0x0e, 0x0f,
		0x10, 0x11, 0x12, 0x13, 0x14, 0x15, 0x16, 0x17,
		0x18, 0x19, 0x1a, 0x1b, 0x1c, 0x1d, 0x1e, 0x1f, // random
		0x00, // session ID length
		0x00, 0x02, // cipher suites length = 2
		0x13, 0x01, // cipher suite
		0x01, // compression methods length = 1
		0x00, // compression method
		0x00, 0x12, // extensions length = 18
		// ALPN extension
		0x00, 0x10, // type = ALPN
		0x00, 0x0e, // length = 14
		0x00, 0x0c, // list length = 12
		0x08, 'h', 't', 't', 'p', '/', '1', '.', '1',
		0x02, 'h', '2',
	}

	m := ParseTLSClientHelloMsgData(&utils.ByteBuffer{Buf: body})
	if m == nil {
		t.Fatal("ParseTLSClientHelloMsgData returned nil")
	}
	alpn, ok := m["alpn"].([]string)
	if !ok {
		t.Fatalf("alpn missing or wrong type: %T", m["alpn"])
	}
	if len(alpn) != 2 || alpn[0] != "http/1.1" || alpn[1] != "h2" {
		t.Errorf("alpn = %v, want [http/1.1 h2]", alpn)
	}
}

func TestParseTLSClientHelloMsgData_SupportedVersionsClient(t *testing.T) {
	// ClientHello with supported_versions extension (client format - list)
	body := []byte{
		0x03, 0x03, // version
		0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07,
		0x08, 0x09, 0x0a, 0x0b, 0x0c, 0x0d, 0x0e, 0x0f,
		0x10, 0x11, 0x12, 0x13, 0x14, 0x15, 0x16, 0x17,
		0x18, 0x19, 0x1a, 0x1b, 0x1c, 0x1d, 0x1e, 0x1f, // random
		0x00, // session ID length
		0x00, 0x02, // cipher suites length = 2
		0x13, 0x01, // cipher suite
		0x01, // compression methods length = 1
		0x00, // compression method
		0x00, 0x0b, // extensions length = 11
		// supported_versions (client list format)
		0x00, 0x2b, // type = supported_versions
		0x00, 0x07, // length = 7
		0x06, // list length = 6
		0x03, 0x04, // TLS 1.3
		0x03, 0x03, // TLS 1.2
		0x03, 0x01, // TLS 1.0
	}

	m := ParseTLSClientHelloMsgData(&utils.ByteBuffer{Buf: body})
	if m == nil {
		t.Fatal("ParseTLSClientHelloMsgData returned nil")
	}
	versions, ok := m["supported_versions"].([]uint16)
	if !ok {
		t.Fatalf("supported_versions missing or wrong type: %T", m["supported_versions"])
	}
	want := []uint16{0x0304, 0x0303, 0x0301}
	if !reflect.DeepEqual(versions, want) {
		t.Errorf("supported_versions = %v, want %v", versions, want)
	}
}

func TestParseTLSServerHelloMsgData_SupportedVersionsServer(t *testing.T) {
	// ServerHello with supported_versions extension (server format - single value)
	body := []byte{
		0x03, 0x03, // version
		0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07,
		0x08, 0x09, 0x0a, 0x0b, 0x0c, 0x0d, 0x0e, 0x0f,
		0x10, 0x11, 0x12, 0x13, 0x14, 0x15, 0x16, 0x17,
		0x18, 0x19, 0x1a, 0x1b, 0x1c, 0x1d, 0x1e, 0x1f, // random
		0x00, // session ID length
		0x13, 0x01, // cipher suite
		0x00, // compression method
		0x00, 0x06, // extensions length = 6
		// supported_versions (server format - single value)
		0x00, 0x2b, // type
		0x00, 0x02, // length = 2
		0x03, 0x04, // TLS 1.3
	}

	m := ParseTLSServerHelloMsgData(&utils.ByteBuffer{Buf: body})
	if m == nil {
		t.Fatal("ParseTLSServerHelloMsgData returned nil")
	}
	v, ok := m["supported_versions"].(uint16)
	if !ok {
		t.Fatalf("supported_versions missing or wrong type: %T", m["supported_versions"])
	}
	if v != 0x0304 {
		t.Errorf("supported_versions = 0x%04x, want 0x0304", v)
	}
}