package magna

import (
	"errors"
	"testing"

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

func TestQuestionDecode(t *testing.T) {
	tests := []struct {
		name           string
		input          []byte
		expectedOffset int
		expected       Question
		expectedErr    error
		wantErrMsg     string
	}{
		{
			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 - label too long",
			input:          []byte{64, 'i', 'n', 'v', 'a', 'l', 'i', 'd', 0, 0, 1, 0, 1},
			expectedOffset: 13,
			expected:       Question{},
			expectedErr:    &InvalidLabelError{},
			wantErrMsg:     "question decode: failed to decode QName: invalid domain label: length 64 exceeds maximum 63",
		},
		{
			name:           "Invalid domain name - compression loop",
			input:          []byte{0xC0, 0x00, 0, 1, 0, 1},
			expectedOffset: 6,
			expected:       Question{},
			expectedErr:    &DomainCompressionError{},
			wantErrMsg:     "question decode: failed to decode QName: invalid domain compression: pointer loop detected",
		},
		{
			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{},
			wantErrMsg:     "question decode: failed to decode QType for example.com: buffer overflow",
		},
		{
			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{},
			wantErrMsg:     "question decode: failed to decode QClass for example.com: buffer overflow",
		},
	}

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

			if tt.expectedErr != nil {
				assert.Error(t, err, "Expected an error but got nil")
				assert.True(t, errors.Is(err, tt.expectedErr), "Error type mismatch. Got %T, expected %T", err, tt.expectedErr)
				if tt.wantErrMsg != "" {
					assert.ErrorContains(t, err, tt.wantErrMsg, "Wrapped error message mismatch")
				}
				assert.Equal(t, tt.expectedOffset, offset, "Offset mismatch on error")
			} else {
				assert.NoError(t, err, "Expected no error but got one")
				assert.Equal(t, tt.expected, *q, "Decoded question mismatch")
				assert.Equal(t, tt.expectedOffset, offset, "Offset mismatch on success")
			}
		})
	}
}

func TestQuestionEncode(t *testing.T) {
	tests := []struct {
		name        string
		question    Question
		initialBuf  []byte
		offsets     map[string]uint16
		expected    []byte
		expectedErr error
		wantErrMsg  string
		newOffsets  map[string]uint16
	}{
		{
			name: "Simple domain - example.com A IN",
			question: Question{
				QName:  "example.com",
				QType:  DNSType(1),
				QClass: DNSClass(1),
			},
			initialBuf: nil,
			offsets:    make(map[string]uint16),
			expected:   []byte{7, 'e', 'x', 'a', 'm', 'p', 'l', 'e', 3, 'c', 'o', 'm', 0, 0, 1, 0, 1},
			newOffsets: map[string]uint16{"example.com": 0, "com": 8},
		},
		{
			name: "Subdomain - subdomain.example.com AAAA IN",
			question: Question{
				QName:  "subdomain.example.com",
				QType:  DNSType(28),
				QClass: DNSClass(1),
			},
			initialBuf: nil,
			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},
			newOffsets: map[string]uint16{"subdomain.example.com": 0, "example.com": 10, "com": 18},
		},
		{
			name: "Different class - example.com MX CH",
			question: Question{
				QName:  "example.com",
				QType:  DNSType(15),
				QClass: DNSClass(3),
			},
			initialBuf: nil,
			offsets:    make(map[string]uint16),
			expected:   []byte{7, 'e', 'x', 'a', 'm', 'p', 'l', 'e', 3, 'c', 'o', 'm', 0, 0, 15, 0, 3},
			newOffsets: map[string]uint16{"example.com": 0, "com": 8},
		},
		{
			name: "Domain compression - example.com after subdomain.example.com",
			question: Question{
				QName:  "example.com",
				QType:  DNSType(1),
				QClass: DNSClass(1),
			},
			initialBuf: nil,
			offsets: map[string]uint16{
				"subdomain.example.com": 0,
				"example.com":           10,
				"com":                   18,
			},
			expected: []byte{0xC0, 0x0a, 0x00, 0x01, 0x00, 0x01},
			newOffsets: map[string]uint16{
				"subdomain.example.com": 0,
				"example.com":           10,
				"com":                   18,
			},
		},
		{
			name: "Encode with initial buffer",
			question: Question{
				QName:  "test.org",
				QType:  AType,
				QClass: IN,
			},
			initialBuf: []byte{0xAA, 0xBB},
			offsets:    make(map[string]uint16),
			expected:   []byte{0xAA, 0xBB, 4, 't', 'e', 's', 't', 3, 'o', 'r', 'g', 0, 0, 1, 0, 1},
			newOffsets: map[string]uint16{"test.org": 2, "org": 7},
		},
		{
			name: "Encode invalid domain - label too long",
			question: Question{
				QName:  "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa.com",
				QType:  AType,
				QClass: IN,
			},
			initialBuf:  nil,
			offsets:     make(map[string]uint16),
			expected:    nil,
			expectedErr: &InvalidLabelError{},
			wantErrMsg:  "question encode: failed to encode QName",
			newOffsets:  map[string]uint16{},
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			currentOffsets := make(map[string]uint16)
			for k, v := range tt.offsets {
				currentOffsets[k] = v
			}

			result, err := tt.question.Encode(tt.initialBuf, &currentOffsets)

			if tt.expectedErr != nil {
				assert.Error(t, err, "Expected an error but got nil")
				assert.True(t, errors.Is(err, tt.expectedErr), "Error type mismatch. Got %T, expected %T", err, tt.expectedErr)
				if tt.wantErrMsg != "" {
					assert.ErrorContains(t, err, tt.wantErrMsg, "Wrapped error message mismatch")
				}
			} else {
				assert.NoError(t, err, "Expected no error but got one")
				assert.Equal(t, tt.expected, result, "Encoded question mismatch")
				assert.Equal(t, tt.newOffsets, currentOffsets, "Final offsets mismatch")
			}
		})
	}
}

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, err := tt.question.Encode(nil, &offsets)
			require.NoError(t, err, "Encoding failed")

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

			assert.NoError(t, err, "Decoding failed")
			assert.Equal(t, len(encoded), offset, "Offset after decoding should match encoded length")
			assert.Equal(t, tt.question, *decodedQuestion, "Decoded question does not match original")
		})
	}
}