package magna

import (
	"math/rand"
)

// Decode decodes a DNS packet.
func (m *Message) Decode(buf []byte) (err error) {
	offset, err := m.Header.Decode(buf, 0)
	if err != nil {
		return err
	}

	for x := 0; x < int(m.Header.QDCount); x++ {
		var question Question
		offset, err = question.Decode(buf, offset)
		if err != nil {
			return err
		}

		m.Question = append(m.Question, question)
	}

	for x := 0; x < int(m.Header.ANCount); x++ {
		var rr ResourceRecord
		offset, err = rr.Decode(buf, offset)
		if err != nil {
			return err
		}

		m.Answer = append(m.Answer, rr)
	}

	for x := 0; x < int(m.Header.NSCount); x++ {
		var rr ResourceRecord
		offset, err = rr.Decode(buf, offset)
		if err != nil {
			return err
		}

		m.Authority = append(m.Authority, rr)
	}

	for x := 0; x < int(m.Header.ARCount); x++ {
		var rr ResourceRecord
		offset, err = rr.Decode(buf, offset)
		if err != nil {
			return err
		}

		m.Additional = append(m.Additional, rr)
	}

	return nil
}

// Encode encodes a message to a DNS packet.
// TODO: set truncation bit if over 512 and udp is protocol
func (m *Message) Encode() []byte {
	m.offsets = make(map[string]uint16)
	bytes := make([]byte, 0, 512)
	bytes = append(bytes, m.Header.Encode()...)

	for _, question := range m.Question {
		bytes = question.Encode(bytes, &m.offsets)
	}

	for _, answer := range m.Answer {
		bytes = answer.Encode(bytes, &m.offsets)
	}

	for _, authority := range m.Authority {
		bytes = authority.Encode(bytes, &m.offsets)
	}

	for _, additional := range m.Additional {
		bytes = additional.Encode(bytes, &m.offsets)
	}

	return bytes
}

func CreateRequest(op OPCode, rd bool) *Message {
	return &Message{
		Header: Header{
			ID:      uint16(rand.Intn(65534) + 1),
			QR:      false,
			OPCode:  op,
			AA:      false,
			TC:      false,
			RD:      rd,
			RA:      false,
			Z:       0,
			RCode:   NOERROR,
			QDCount: 0,
			ARCount: 0,
			NSCount: 0,
			ANCount: 0,
		},
		Question:   make([]Question, 0),
		Answer:     make([]ResourceRecord, 0),
		Additional: make([]ResourceRecord, 0),
		Authority:  make([]ResourceRecord, 0),
	}
}

func (m *Message) CreateReply(req *Message) *Message {
	m.Header.ID = req.Header.ID
	m.Header.QR = true
	m.Header.OPCode = req.Header.OPCode

	return m
}

func (m *Message) AddQuestion(q Question) *Message {
	m.Header.QDCount += 1
	m.Question = append(m.Question, q)

	return m
}

func (m *Message) SetRCode(rc RCode) *Message {
	m.Header.RCode = rc

	return m
}