package magna

import (
	"bytes"
	"encoding/binary"
	"net"
	"testing"

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

func TestMessageDecode(t *testing.T) {
	tests := []struct {
		name     string
		input    []byte
		expected Message
		wantErr  bool
	}{
		{
			name: "Valid DNS message with one question",
			input: func() []byte {
				buf := new(bytes.Buffer)
				binary.Write(buf, binary.BigEndian, uint16(1234))
				binary.Write(buf, binary.BigEndian, uint16(0x0100))
				binary.Write(buf, binary.BigEndian, uint16(1))
				binary.Write(buf, binary.BigEndian, uint16(0))
				binary.Write(buf, binary.BigEndian, uint16(0))
				binary.Write(buf, binary.BigEndian, uint16(0))
				buf.Write([]byte{3, 'w', 'w', 'w', 7, 'e', 'x', 'a', 'm', 'p', 'l', 'e', 3, 'c', 'o', 'm', 0})
				binary.Write(buf, binary.BigEndian, uint16(1))
				binary.Write(buf, binary.BigEndian, uint16(1))
				return buf.Bytes()
			}(),
			expected: Message{
				Header: Header{
					ID:      1234,
					QR:      false,
					RD:      true,
					OPCode:  0,
					QDCount: 1,
				},
				Question: []Question{
					{
						QName:  "www.example.com",
						QType:  1,
						QClass: 1,
					},
				},
			},
			wantErr: false,
		},
		{
			name: "Valid DNS message with one answer",
			input: func() []byte {
				buf := new(bytes.Buffer)
				binary.Write(buf, binary.BigEndian, uint16(5678))
				binary.Write(buf, binary.BigEndian, uint16(0x8180))
				binary.Write(buf, binary.BigEndian, uint16(0))
				binary.Write(buf, binary.BigEndian, uint16(1))
				binary.Write(buf, binary.BigEndian, uint16(0))
				binary.Write(buf, binary.BigEndian, uint16(0))
				buf.Write([]byte{3, 'w', 'w', 'w', 7, 'e', 'x', 'a', 'm', 'p', 'l', 'e', 3, 'c', 'o', 'm', 0})
				binary.Write(buf, binary.BigEndian, uint16(1))
				binary.Write(buf, binary.BigEndian, uint16(1))
				binary.Write(buf, binary.BigEndian, uint32(3600))
				binary.Write(buf, binary.BigEndian, uint16(4))
				binary.Write(buf, binary.BigEndian, uint32(0x0A000001))
				return buf.Bytes()
			}(),
			expected: Message{
				Header: Header{
					ID:      5678,
					QR:      true,
					OPCode:  0,
					AA:      false,
					RD:      true,
					RA:      true,
					RCode:   0,
					ANCount: 1,
				},
				Answer: []ResourceRecord{
					{
						Name:     "www.example.com",
						RType:    1,
						RClass:   1,
						TTL:      3600,
						RDLength: 4,
						RData:    &A{net.IP([]byte{10, 0, 0, 1})},
					},
				},
			},
			wantErr: false,
		},
		{
			name:    "Invalid input - empty buffer",
			input:   []byte{},
			wantErr: true,
		},
	}

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

			if tt.wantErr {
				assert.Error(t, err)
			} else {
				assert.NoError(t, err)
				assert.Equal(t, tt.expected.Header, m.Header)
				assert.Equal(t, tt.expected.Question, m.Question)
				assert.Equal(t, tt.expected.Answer, m.Answer)
				assert.Equal(t, tt.expected.Authority, m.Authority)
				assert.Equal(t, tt.expected.Additional, m.Additional)
			}
		})
	}
}

func FuzzDecodeMessage(f *testing.F) {
	testcases := [][]byte{
		{
			0x8e, 0x19, 0x81, 0x80, 0x00, 0x01, 0x00, 0x01, 0x00, 0x00, 0x00, 0x00, 0x04, 0x6e, 0x65,
			0x77, 0x73, 0x0b, 0x79, 0x63, 0x6f, 0x6d, 0x62, 0x69, 0x6e, 0x61, 0x74, 0x6f, 0x72, 0x03,
			0x63, 0x6f, 0x6d, 0x00, 0x00, 0x01, 0x00, 0x01, 0xc0, 0x0c, 0x00, 0x01, 0x00, 0x01, 0x00,
			0x00, 0x00, 0x01, 0x00, 0x04, 0xd1, 0xd8, 0xe6, 0xcf,
		},
	}
	for _, tc := range testcases {
		f.Add(tc)
	}
	f.Fuzz(func(t *testing.T, msg []byte) {
		var m Message
		m.Decode(msg)
	})
}