package magna

import (
	"testing"

	"github.com/stretchr/testify/assert"
)

func TestQuestionDecode(t *testing.T) {
	tests := []struct {
		name           string
		input          []byte
		expectedOffset int
		expected       Question
		expectedErr    error
	}{
		{
			name:           "Valid question - example.com A IN",
			input:          []byte{7, 'e', 'x', 'a', 'm', 'p', 'l', 'e', 3, 'c', 'o', 'm', 0, 0, 1, 0, 1},
			expectedOffset: 17,
			expected: Question{
				QName:  "example.com",
				QType:  DNSType(1),
				QClass: DNSClass(1),
			},
			expectedErr: nil,
		},
		{
			name:           "Valid question - example.com MX CH",
			input:          []byte{7, 'e', 'x', 'a', 'm', 'p', 'l', 'e', 3, 'c', 'o', 'm', 0, 0, 15, 0, 3},
			expectedOffset: 17,
			expected: Question{
				QName:  "example.com",
				QType:  DNSType(15),
				QClass: DNSClass(3),
			},
			expectedErr: nil,
		},
		{
			name:           "Invalid domain name",
			input:          []byte{255, 'i', 'n', 'v', 'a', 'l', 'i', 'd', 0, 0, 1, 0, 1},
			expectedOffset: 13,
			expected:       Question{},
			expectedErr:    &BufferOverflowError{},
		},
		{
			name:           "Insufficient buffer for QType",
			input:          []byte{7, 'e', 'x', 'a', 'm', 'p', 'l', 'e', 3, 'c', 'o', 'm', 0, 0},
			expectedOffset: 14,
			expected:       Question{QName: "example.com"},
			expectedErr:    &BufferOverflowError{},
		},
		{
			name:           "Insufficient buffer for QClass",
			input:          []byte{7, 'e', 'x', 'a', 'm', 'p', 'l', 'e', 3, 'c', 'o', 'm', 0, 0, 1, 0},
			expectedOffset: 16,
			expected:       Question{QName: "example.com", QType: DNSType(1)},
			expectedErr:    &BufferOverflowError{},
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			q := &Question{}
			offset, err := q.Decode(tt.input, 0)

			assert.Equal(t, tt.expectedOffset, offset)

			if tt.expectedErr != nil {
				assert.Error(t, err)
				assert.IsType(t, tt.expectedErr, err)
			} else {
				assert.NoError(t, err)
				assert.Equal(t, tt.expected, *q)
			}
		})
	}
}

func TestQuestionEncode(t *testing.T) {
	tests := []struct {
		name     string
		question Question
		offsets  map[string]uint16
		expected []byte
	}{
		{
			name: "Simple domain - example.com A IN",
			question: Question{
				QName:  "example.com",
				QType:  DNSType(1),
				QClass: DNSClass(1),
			},
			offsets:  make(map[string]uint16),
			expected: []byte{7, 'e', 'x', 'a', 'm', 'p', 'l', 'e', 3, 'c', 'o', 'm', 0, 0, 1, 0, 1},
		},
		{
			name: "Subdomain - subdomain.example.com AAAA IN",
			question: Question{
				QName:  "subdomain.example.com",
				QType:  DNSType(28),
				QClass: DNSClass(1),
			},
			offsets:  make(map[string]uint16),
			expected: []byte{9, 's', 'u', 'b', 'd', 'o', 'm', 'a', 'i', 'n', 7, 'e', 'x', 'a', 'm', 'p', 'l', 'e', 3, 'c', 'o', 'm', 0, 0, 28, 0, 1},
		},
		{
			name: "Different class - example.com MX CH",
			question: Question{
				QName:  "example.com",
				QType:  DNSType(15),
				QClass: DNSClass(3),
			},
			offsets:  make(map[string]uint16),
			expected: []byte{7, 'e', 'x', 'a', 'm', 'p', 'l', 'e', 3, 'c', 'o', 'm', 0, 0, 15, 0, 3},
		},
		{
			name: "Domain compression - example.com after subdomain.example.com",
			question: Question{
				QName:  "example.com",
				QType:  DNSType(1),
				QClass: DNSClass(1),
			},
			offsets: map[string]uint16{
				"com":         22,
				"example.com": 19,
			},
			expected: []byte{0xC0, 0x13, 0x00, 0x01, 0x00, 0x01},
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			result := tt.question.Encode(nil, &tt.offsets)
			assert.Equal(t, tt.expected, result)

			if len(tt.offsets) == 0 {
				expectedOffsets := map[string]uint16{
					tt.question.QName: 0,
				}
				for i := 0; i < len(tt.question.QName); i++ {
					if tt.question.QName[i] == '.' {
						expectedOffsets[tt.question.QName[i+1:]] = uint16(i + 1)
					}
				}
				assert.Equal(t, expectedOffsets, tt.offsets)
			}
		})
	}
}

func TestQuestionEncodeDecodeRoundTrip(t *testing.T) {
	tests := []struct {
		name     string
		question Question
	}{
		{
			name: "Simple domain - example.com A IN",
			question: Question{
				QName:  "example.com",
				QType:  DNSType(1),
				QClass: DNSClass(1),
			},
		},
		{
			name: "Subdomain - subdomain.example.com AAAA IN",
			question: Question{
				QName:  "subdomain.example.com",
				QType:  DNSType(28),
				QClass: DNSClass(1),
			},
		},
		{
			name: "Different class - example.com MX CH",
			question: Question{
				QName:  "example.com",
				QType:  DNSType(15),
				QClass: DNSClass(3),
			},
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			offsets := make(map[string]uint16)
			encoded := tt.question.Encode(nil, &offsets)

			decodedQuestion := &Question{}
			_, err := decodedQuestion.Decode(encoded, 0)

			assert.NoError(t, err)
			assert.Equal(t, tt.question, *decodedQuestion)
		})
	}
}

func TestQuestionEncodeWithExistingBuffer(t *testing.T) {
	question := Question{
		QName:  "example.com",
		QType:  DNSType(1),
		QClass: DNSClass(1),
	}

	existingBuffer := []byte{0xFF, 0xFF, 0xFF, 0xFF}
	offsets := make(map[string]uint16)

	result := question.Encode(existingBuffer, &offsets)

	expected := append(
		existingBuffer,
		[]byte{7, 'e', 'x', 'a', 'm', 'p', 'l', 'e', 3, 'c', 'o', 'm', 0, 0, 1, 0, 1}...,
	)

	assert.Equal(t, expected, result)
}

func TestQuestionEncodeLongDomainName(t *testing.T) {
	longLabel := make([]byte, 63)
	for i := range longLabel {
		longLabel[i] = 'a'
	}
	longDomainName := string(longLabel) + "." + string(longLabel) + "." + string(longLabel) + "." + string(longLabel[:61])

	question := Question{
		QName:  longDomainName,
		QType:  DNSType(1),
		QClass: DNSClass(1),
	}

	offsets := make(map[string]uint16)
	encoded := question.Encode(nil, &offsets)

	assert.Equal(t, 259, len(encoded))

	decodedQuestion := &Question{}
	_, err := decodedQuestion.Decode(encoded, 0)

	assert.NoError(t, err)
	assert.Equal(t, question, *decodedQuestion)
}