package magna

import (
	"encoding/binary"
	"fmt"
)

// Decode decodes the header from the bytes.
func (h *Header) Decode(buf []byte, offset int) (int, error) {
	var err error
	h.ID, offset, err = getU16(buf, offset)
	if err != nil {
		return len(buf), fmt.Errorf("header decode: failed to read ID: %w", err)
	}

	flags, offset, err := getU16(buf, offset)
	if err != nil {
		return len(buf), fmt.Errorf("header decode: failed to read flags: %w", err)
	}

	h.QDCount, offset, err = getU16(buf, offset)
	if err != nil {
		return len(buf), fmt.Errorf("header decode: failed to read QDCount: %w", err)
	}

	h.ANCount, offset, err = getU16(buf, offset)
	if err != nil {
		return len(buf), fmt.Errorf("header decode: failed to read ANCount: %w", err)
	}

	h.NSCount, offset, err = getU16(buf, offset)
	if err != nil {
		return len(buf), fmt.Errorf("header decode: failed to read NSCount: %w", err)
	}

	h.ARCount, offset, err = getU16(buf, offset)
	if err != nil {
		return len(buf), fmt.Errorf("header decode: failed to read ARCount: %w", err)
	}

	h.QR = ((flags >> 15) & 0x01) == 1
	h.OPCode = OPCode((flags >> 11) & 0x0F)
	h.AA = ((flags >> 10) & 0x01) == 1
	h.TC = ((flags >> 9) & 0x01) == 1
	h.RD = ((flags >> 8) & 0x01) == 1
	h.RA = ((flags >> 7) & 0x01) == 1
	h.Z = uint8((flags >> 4) & 0x07)
	h.RCode = RCode((flags & 0x0F))

	return offset, nil
}

// Encode encodes the header packet to the byte representation.
func (h *Header) Encode() []byte {
	bytes := make([]byte, 0, 12)

	bytes = binary.BigEndian.AppendUint16(bytes, h.ID)

	var flags uint16
	if h.QR {
		flags |= 1 << 15
	}

	flags |= uint16(h.OPCode) << 11

	if h.AA {
		flags |= 1 << 10
	}

	if h.TC {
		flags |= 1 << 9
	}

	if h.RD {
		flags |= 1 << 8
	}

	if h.RA {
		flags |= 1 << 7
	}

	flags |= uint16(h.Z) << 4
	flags |= uint16(h.RCode)

	bytes = binary.BigEndian.AppendUint16(bytes, flags)
	bytes = binary.BigEndian.AppendUint16(bytes, h.QDCount)
	bytes = binary.BigEndian.AppendUint16(bytes, h.ANCount)
	bytes = binary.BigEndian.AppendUint16(bytes, h.NSCount)
	bytes = binary.BigEndian.AppendUint16(bytes, h.ARCount)

	return bytes
}