package magna

import (
	"fmt"
	"math/rand"
)

// Decode decodes a DNS packet.
func (m *Message) Decode(buf []byte) (err error) {
	// gets checked when parsing additional records
	m.HasEDNS = false

	offset, err := m.Header.Decode(buf, 0)
	if err != nil {
		return fmt.Errorf("failed to decode message header: %w", err)
	}

	m.Question = make([]Question, 0, m.Header.QDCount)
	for i := range m.Header.QDCount {
		var question Question
		offset, err = question.Decode(buf, offset)
		if err != nil {
			return fmt.Errorf("failed to decode question #%d: %w", i+1, err)
		}

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

	m.Answer = make([]ResourceRecord, 0, m.Header.ANCount)
	for i := range m.Header.ANCount {
		var rr ResourceRecord
		offset, err = rr.Decode(buf, offset)
		if err != nil {
			return fmt.Errorf("failed to decode answer record #%d: %w", i+1, err)
		}

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

	m.Authority = make([]ResourceRecord, 0, m.Header.NSCount)
	for i := range m.Header.NSCount {
		var rr ResourceRecord
		offset, err = rr.Decode(buf, offset)
		if err != nil {
			return fmt.Errorf("failed to decode authority record #%d: %w", i+1, err)
		}

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

	m.Additional = make([]ResourceRecord, 0, m.Header.ARCount)
	for i := range m.Header.ARCount {
		var rr ResourceRecord
		offset, err = rr.Decode(buf, offset)
		if err != nil {
			return fmt.Errorf("failed to decode additional record #%d: %w", i+1, err)
		}

		if rr.RType == OPTType {
			opt, ok := rr.RData.(*OPT)
			if !ok {
				// this should never fail
				return fmt.Errorf("unable to parse RData as OPT")
			}

			m.HasEDNS = true
			m.ExtendedRCode = uint8(rr.TTL >> 24 & 0xFF000000)
			m.EDNSVersion = uint8((rr.TTL & 0x00FF0000) >> 16)
			m.EDNSFlags = uint16(rr.TTL & 0x0000FFFF)
			m.EDNSOptions = opt.Options
			m.UDPSize = uint16(rr.RClass)
		}

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

	if !m.HasEDNS {
		m.ExtendedRCode = 0
		m.EDNSVersion = 0
		m.EDNSFlags = 0
		m.EDNSOptions = make([]EDNSOption, 0)
		m.UDPSize = 512 // default in rfc-1035 section 2.3.4.
	}

	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, error) {
	m.offsets = make(map[string]uint16)
	bytes := make([]byte, 0, 512)

	headerBytes := m.Header.Encode()
	bytes = append(bytes, headerBytes...)

	var err error

	for i, question := range m.Question {
		bytes, err = question.Encode(bytes, &m.offsets)
		if err != nil {
			return nil, fmt.Errorf("failed to encode question #%d (%s): %w", i+1, question.QName, err)
		}
	}

	for i, answer := range m.Answer {
		bytes, err = answer.Encode(bytes, &m.offsets)
		if err != nil {
			return nil, fmt.Errorf("failed to encode answer record #%d (%s): %w", i+1, answer.Name, err)
		}
	}

	for i, authority := range m.Authority {
		bytes, err = authority.Encode(bytes, &m.offsets)
		if err != nil {
			return nil, fmt.Errorf("failed to encode authority record #%d (%s): %w", i+1, authority.Name, err)
		}
	}

	for i, additional := range m.Additional {
		bytes, err = additional.Encode(bytes, &m.offsets)
		if err != nil {
			return nil, fmt.Errorf("failed to encode additional record #%d (%s): %w", i+1, additional.Name, err)
		}
	}

	return bytes, nil
}

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
}