package magna

import (
	"errors"
	"testing"

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

func BenchmarkDecodeDomainSimple(b *testing.B) {
	input := []byte{3, 'w', 'w', 'w', 7, 'e', 'x', 'a', 'm', 'p', 'l', 'e', 3, 'c', 'o', 'm', 0}
	for i := 0; i < b.N; i++ {
		_, _, _ = decodeDomain(input, 0)
	}
}

func BenchmarkDecodeDomainCompressed(b *testing.B) {
	input := []byte{
		0x07, 'e', 'x', 'a', 'm', 'p', 'l', 'e', 0x03, 'c', 'o', 'm', 0x00,
		0x03, 'f', 'o', 'o', 0xc0, 0x00,
	}
	offset := 13
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		_, _, _ = decodeDomain(input, offset)
	}
}

func BenchmarkEncodeDomainSimple(b *testing.B) {
	domain := "www.example.com"
	offsets := make(map[string]uint16)
	out := make([]byte, 0, 64)
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		_, _ = encodeDomain(out[:0], domain, &offsets)
		for k := range offsets {
			delete(offsets, k)
		}
	}
}

func BenchmarkEncodeDomainWithCompression(b *testing.B) {
	domain1 := "www.example.com"
	domain2 := "mail.example.com"
	offsets := make(map[string]uint16)
	out := make([]byte, 0, 128)
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		tempOut, _ := encodeDomain(out[:0], domain1, &offsets)
		_, _ = encodeDomain(tempOut, domain2, &offsets)
		for k := range offsets {
			delete(offsets, k)
		}
	}
}

func TestDecodeDomain(t *testing.T) {
	tests := []struct {
		name           string
		offset         int
		input          []byte
		expectedDomain string
		expectedOffset int
		expectedError  error
		errorCheck     func(t *testing.T, err error)
	}{
		{
			name:           "Simple domain",
			input:          []byte{3, 'w', 'w', 'w', 7, 'e', 'x', 'a', 'm', 'p', 'l', 'e', 3, 'c', 'o', 'm', 0},
			expectedDomain: "www.example.com",
			expectedOffset: 17,
			expectedError:  nil,
		},
		{
			name:           "Domain with compression",
			offset:         17,
			input:          []byte{3, 'w', 'w', 'w', 7, 'e', 'x', 'a', 'm', 'p', 'l', 'e', 3, 'c', 'o', 'm', 0, 3, 'f', 'o', 'o', 0xC0, 0},
			expectedDomain: "foo.www.example.com",
			expectedOffset: 23,
			expectedError:  nil,
		},
		{
			name:           "Invalid label length",
			input:          []byte{64, 'x'},
			expectedDomain: "",
			expectedOffset: 2,
			expectedError:  &InvalidLabelError{Length: 64},
			errorCheck: func(t *testing.T, err error) {
				var target *InvalidLabelError
				assert.True(t, errors.As(err, &target))
				assert.Equal(t, 64, target.Length)
			},
		},
		{
			name:           "Compression loop",
			input:          []byte{0xC0, 0, 0xC0, 0},
			expectedDomain: "",
			expectedOffset: 4,
			expectedError:  &DomainCompressionError{},
			errorCheck: func(t *testing.T, err error) {
				assert.IsType(t, &DomainCompressionError{}, err)
			},
		},
		{
			name:           "Truncated input",
			input:          []byte{3, 'w', 'w'},
			expectedDomain: "",
			expectedOffset: 3,
			expectedError:  &BufferOverflowError{Length: 3, Offset: 4},
			errorCheck: func(t *testing.T, err error) {
				var target *BufferOverflowError
				assert.True(t, errors.As(err, &target), "Expected BufferOverflowError")
				if target != nil {
					assert.Equal(t, 3, target.Length)
					assert.Equal(t, 1+3, target.Offset)
				}
				assert.Contains(t, err.Error(), "failed to read domain label data")
			},
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			domain, offset, err := decodeDomain(tt.input, tt.offset)

			t.Logf("Test: %s, Input: %x, OffsetIn: %d => Domain: '%s', OffsetOut: %d, Err: %v", tt.name, tt.input, tt.offset, domain, offset, err)

			if tt.expectedError != nil {
				assert.Error(t, err, "Expected an error but got nil")
				if tt.errorCheck != nil {
					tt.errorCheck(t, err)
				} else {
					assert.IsType(t, tt.expectedError, err, "Error type mismatch")
				}
			} else {
				assert.NoError(t, err, "Expected no error but got one")
			}

			assert.Equal(t, tt.expectedDomain, domain, "Domain mismatch")
			if tt.expectedError == nil {
				assert.Equal(t, tt.expectedOffset, offset, "Offset mismatch")
			}
		})
	}
}

func TestEncodeDomain(t *testing.T) {
	tests := []struct {
		name        string
		input       string
		initialBuf  []byte
		offsets     map[string]uint16
		expected    []byte
		expectedErr error
		newOffsets  map[string]uint16
	}{
		{
			name:       "Simple domain",
			input:      "example.com",
			initialBuf: []byte{},
			offsets:    make(map[string]uint16),
			expected:   []byte{7, 'e', 'x', 'a', 'm', 'p', 'l', 'e', 3, 'c', 'o', 'm', 0},
			newOffsets: map[string]uint16{"example.com": 0, "com": 8},
		},
		{
			name:       "Domain with existing offset for compression",
			input:      "test.example.com",
			initialBuf: []byte{},
			offsets:    map[string]uint16{"example.com": 10},
			expected:   []byte{4, 't', 'e', 's', 't', 0xC0, 0x0A},
			newOffsets: map[string]uint16{"test.example.com": 0, "example.com": 10},
		},
		{
			name:       "Multiple subdomains",
			input:      "a.b.c.d",
			initialBuf: []byte{},
			offsets:    make(map[string]uint16),
			expected:   []byte{1, 'a', 1, 'b', 1, 'c', 1, 'd', 0},
			newOffsets: map[string]uint16{"a.b.c.d": 0, "b.c.d": 2, "c.d": 4, "d": 6},
		},
		{
			name:       "Root domain",
			input:      ".",
			initialBuf: []byte{},
			offsets:    make(map[string]uint16),
			expected:   []byte{0},
			newOffsets: map[string]uint16{},
		},
		{
			name:       "Empty domain",
			input:      "",
			initialBuf: []byte{},
			offsets:    make(map[string]uint16),
			expected:   []byte{0},
			newOffsets: map[string]uint16{},
		},
		{
			name:        "Label too long",
			input:       "labeltoolonglabeltoolonglabeltoolonglabeltoolonglabeltoolonglabeltoolong.com",
			initialBuf:  []byte{},
			offsets:     make(map[string]uint16),
			expected:    nil,
			expectedErr: &InvalidLabelError{Length: 72},
			newOffsets:  map[string]uint16{},
		},
		{
			name:        "Empty label inside domain",
			input:       "example..com",
			initialBuf:  []byte{},
			offsets:     make(map[string]uint16),
			expected:    nil,
			expectedErr: &InvalidLabelError{Length: 0},
			newOffsets:  map[string]uint16{},
		},
		{
			name:       "Append to existing buffer",
			input:      "example.com",
			initialBuf: []byte{0xAA, 0xBB},
			offsets:    make(map[string]uint16),
			expected:   []byte{0xAA, 0xBB, 7, 'e', 'x', 'a', 'm', 'p', 'l', 'e', 3, 'c', 'o', 'm', 0},
			newOffsets: map[string]uint16{"example.com": 2, "com": 10},
		},
	}

	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 := encodeDomain(tt.initialBuf, tt.input, &currentOffsets)

			if tt.expectedErr != nil {
				assert.Error(t, err, "Expected an error but got nil")
				assert.IsType(t, tt.expectedErr, err, "Error type mismatch")
				if expectedILE, ok := tt.expectedErr.(*InvalidLabelError); ok {
					actualILE := &InvalidLabelError{}
					if assert.True(t, errors.As(err, &actualILE)) {
						assert.Equal(t, expectedILE.Length, actualILE.Length)
					}
				}
			} else {
				assert.NoError(t, err, "Expected no error but got one")
				assert.Equal(t, tt.expected, result, "Encoded domain does not match expected output")
				assert.Equal(t, tt.newOffsets, currentOffsets, "Offsets map does not match expected state")
			}
		})
	}
}

func FuzzDecodeDomain(f *testing.F) {
	testcases := [][]byte{
		{
			0x03, 0x63, 0x6f, 0x6d, 0x00,
		},
		{
			0x03, 0x63, 0x6f, 0x6d, 0x00, 0x01, 0x63, 0xc0, 0x00,
		},
		{
			0x03, 0x63, 0x6f, 0x6d, 0xc0, 0x00,
		},
		{
			0xc0, 0x00,
		},
		{
			0xc0, 0xff,
		},
		{
			0x40,
		},
		{
			0x03, 0x63, 0x6f,
		},
		{
			0xc0,
		},
	}
	for _, tc := range testcases {
		f.Add(tc)
	}
	f.Fuzz(func(t *testing.T, msg []byte) {
		_, _, err := decodeDomain(msg, 0)
		if err != nil {
			var bufErr *BufferOverflowError
			var labelErr *InvalidLabelError
			var compErr *DomainCompressionError

			if !(errors.As(err, &bufErr) || errors.As(err, &labelErr) || errors.As(err, &compErr)) {
				t.Errorf("Fuzzing decodeDomain: unexpected error type %T: %v for input %x", err, err, msg)
			}
		}
	})
}