package magna

import (
	"encoding/binary"
	"fmt"
	"net"
	"strings"
)

func (a *A) Decode(buf []byte, offset int, rdlength int) (int, error) {
	bytes, offset, err := getSlice(buf, offset, rdlength)
	if err != nil {
		return offset, fmt.Errorf("A record: failed to read address data: %w", err)
	}

	a.Address = net.IP(bytes)
	if a.Address.To4() == nil {
		return offset, fmt.Errorf("A record: decoded data is not a valid IPv4 address: %v", bytes)
	}
	return offset, nil
}

func (a *A) Encode(bytes []byte, offsets *map[string]uint16) ([]byte, error) {
	ipv4 := a.Address.To4()
	if ipv4 == nil {
		return nil, fmt.Errorf("A record: cannot encode non-IPv4 address %s", a.Address.String())
	}

	return append(bytes, a.Address.To4()...), nil
}

func (a A) String() string {
	return a.Address.String()
}

func (ns *NS) Decode(buf []byte, offset int, rdlength int) (int, error) {
	var err error
	ns.NSDName, offset, err = decodeDomain(buf, offset)
	if err != nil {
		return offset, fmt.Errorf("NS record: failed to decode NSDName: %w", err)
	}

	return offset, nil
}

func (ns *NS) Encode(bytes []byte, offsets *map[string]uint16) ([]byte, error) {
	var err error
	bytes, err = encodeDomain(bytes, ns.NSDName, offsets)
	if err != nil {
		return nil, fmt.Errorf("NS record: failed to encode NSDName %s: %w", ns.NSDName, err)
	}

	return bytes, nil
}

func (ns NS) String() string {
	return ns.NSDName
}

func (md *MD) Decode(buf []byte, offset int, rdlength int) (int, error) {
	var err error
	md.MADName, offset, err = decodeDomain(buf, offset)
	if err != nil {
		return offset, fmt.Errorf("MD record: failed to decode MADName %s: %w", md.MADName, err)
	}

	return offset, nil
}

func (md *MD) Encode(bytes []byte, offsets *map[string]uint16) ([]byte, error) {
	var err error
	bytes, err = encodeDomain(bytes, md.MADName, offsets)
	if err != nil {
		return nil, fmt.Errorf("MD record: failed to encode MADName %s: %w", md.MADName, err)
	}

	return bytes, nil
}

func (md MD) String() string {
	return md.MADName
}

func (mf *MF) Decode(buf []byte, offset int, rdlength int) (int, error) {
	var err error
	mf.MADName, offset, err = decodeDomain(buf, offset)
	if err != nil {
		return offset, fmt.Errorf("MF record: failed to decode MADName: %w", err)
	}

	return offset, nil
}

func (mf *MF) Encode(bytes []byte, offsets *map[string]uint16) ([]byte, error) {
	var err error
	bytes, err = encodeDomain(bytes, mf.MADName, offsets)
	if err != nil {
		return nil, fmt.Errorf("MF record: failed to encode MADName %s: %w", mf.MADName, err)
	}

	return bytes, nil
}

func (mf MF) String() string {
	return mf.MADName
}

func (c *CNAME) Decode(buf []byte, offset int, rdlength int) (int, error) {
	var err error
	c.CName, offset, err = decodeDomain(buf, offset)
	if err != nil {
		return offset, fmt.Errorf("CNAME record: failed to decode CNAME: %w", err)
	}

	return offset, nil
}

func (c *CNAME) Encode(bytes []byte, offsets *map[string]uint16) ([]byte, error) {
	var err error
	bytes, err = encodeDomain(bytes, c.CName, offsets)
	if err != nil {
		return nil, fmt.Errorf("CNAME record: failed to encode CNAME %s: %w", c.CName, err)
	}

	return bytes, nil
}

func (c CNAME) String() string {
	return c.CName
}

func (soa *SOA) Decode(buf []byte, offset int, rdlength int) (int, error) {
	var err error
	soa.MName, offset, err = decodeDomain(buf, offset)
	if err != nil {
		return offset, fmt.Errorf("SOA record: failed to decode MName: %w", err)
	}

	soa.RName, offset, err = decodeDomain(buf, offset)
	if err != nil {
		return offset, fmt.Errorf("SOA record: failed to decode RName: %w", err)
	}

	soa.Serial, offset, err = getU32(buf, offset)
	if err != nil {
		return offset, fmt.Errorf("SOA record: failed to decode Serial: %w", err)
	}

	soa.Refresh, offset, err = getU32(buf, offset)
	if err != nil {
		return offset, fmt.Errorf("SOA record: failed to decode Refresh: %w", err)
	}

	soa.Retry, offset, err = getU32(buf, offset)
	if err != nil {
		return offset, fmt.Errorf("SOA record: failed to decode Retry: %w", err)
	}

	soa.Expire, offset, err = getU32(buf, offset)
	if err != nil {
		return offset, fmt.Errorf("SOA record: failed to decode Expire: %w", err)
	}

	soa.Minimum, offset, err = getU32(buf, offset)
	if err != nil {
		return offset, fmt.Errorf("SOA record: failed to decode Minimum: %w", err)
	}

	return offset, nil
}

func (soa *SOA) Encode(bytes []byte, offsets *map[string]uint16) ([]byte, error) {
	var err error
	bytes, err = encodeDomain(bytes, soa.MName, offsets)
	if err != nil {
		return nil, fmt.Errorf("SOA record: failed to encode MName %s: %w", soa.MName, err)
	}

	bytes, err = encodeDomain(bytes, soa.RName, offsets)
	if err != nil {
		return nil, fmt.Errorf("SOA record: failed to encode RName %s: %w", soa.RName, err)
	}

	bytes = binary.BigEndian.AppendUint32(bytes, soa.Serial)
	bytes = binary.BigEndian.AppendUint32(bytes, soa.Refresh)
	bytes = binary.BigEndian.AppendUint32(bytes, soa.Retry)
	bytes = binary.BigEndian.AppendUint32(bytes, soa.Expire)
	bytes = binary.BigEndian.AppendUint32(bytes, soa.Minimum)

	return bytes, nil
}

func (soa SOA) String() string {
	return fmt.Sprintf("%s %s %d %d %d %d %d", soa.MName, soa.RName, soa.Serial, soa.Refresh, soa.Retry, soa.Expire, soa.Minimum)
}

func (mb *MB) Decode(buf []byte, offset int, rdlength int) (int, error) {
	madname, offset, err := decodeDomain(buf, offset)
	if err != nil {
		return offset, fmt.Errorf("MB record: failed to decode MADName: %w", err)
	}

	mb.MADName = string(madname)
	return offset, nil
}

func (mb *MB) Encode(bytes []byte, offsets *map[string]uint16) ([]byte, error) {
	var err error
	bytes, err = encodeDomain(bytes, mb.MADName, offsets)
	if err != nil {
		return nil, fmt.Errorf("MB record: failed to encode MADName %s: %w", mb.MADName, err)
	}

	return bytes, nil
}

func (mb MB) String() string {
	return mb.MADName
}

func (mg *MG) Decode(buf []byte, offset int, rdlength int) (int, error) {
	var err error
	mg.MGMName, offset, err = decodeDomain(buf, offset)
	if err != nil {
		return offset, fmt.Errorf("MG record: failed to decode MGMName: %w", err)
	}

	return offset, nil
}

func (mg *MG) Encode(bytes []byte, offsets *map[string]uint16) ([]byte, error) {
	var err error
	bytes, err = encodeDomain(bytes, mg.MGMName, offsets)
	if err != nil {
		return nil, fmt.Errorf("MG record: failed to encode MGMName %s: %w", mg.MGMName, err)
	}
	return bytes, nil
}

func (mg MG) String() string {
	return mg.MGMName
}

func (mr *MR) Decode(buf []byte, offset int, rdlength int) (int, error) {
	var err error
	mr.NEWName, offset, err = decodeDomain(buf, offset)
	if err != nil {
		return offset, fmt.Errorf("MR record: failed to decode NEWName: %w", err)
	}

	return offset, nil
}

func (mr *MR) Encode(bytes []byte, offsets *map[string]uint16) ([]byte, error) {
	var err error
	bytes, err = encodeDomain(bytes, mr.NEWName, offsets)
	if err != nil {
		return nil, fmt.Errorf("MR record: failed to encode NEWName: %w", err)
	}

	return bytes, nil
}

func (mr MR) String() string {
	return mr.NEWName
}

func (null *NULL) Decode(buf []byte, offset int, rdlength int) (int, error) {
	var err error
	null.Anything, offset, err = getSlice(buf, offset, int(rdlength))
	if err != nil {
		return offset, fmt.Errorf("NULL record: failed to read data: %w", err)
	}

	return offset, nil
}

func (null *NULL) Encode(bytes []byte, offsets *map[string]uint16) ([]byte, error) {
	return append(bytes, null.Anything...), nil
}

func (null NULL) String() string {
	return string(null.Anything)
}

func (wks *WKS) Decode(buf []byte, offset int, rdlength int) (int, error) {
	if rdlength < 5 {
		return len(buf), fmt.Errorf("WKS record: RDLENGTH %d is too short, minimum 5 required", rdlength)
	}

	addressBytes, nextOffset, err := getSlice(buf, offset, 4)
	if err != nil {
		return len(buf), fmt.Errorf("WKS record: failed to read address: %w", err)
	}
	offset = nextOffset
	wks.Address = net.IP(addressBytes)

	protocol, nextOffset, err := getU8(buf, offset)
	if err != nil {
		return len(buf), fmt.Errorf("WKS record: failed to read protocol: %w", err)
	}
	offset = nextOffset
	wks.Protocol = protocol

	bitmapLength := rdlength - 5
	wks.BitMap, nextOffset, err = getSlice(buf, offset, bitmapLength)
	if err != nil {
		return len(buf), fmt.Errorf("WKS record: failed to read bitmap: %w", err)
	}
	offset = nextOffset

	return offset, nil
}

func (wks *WKS) Encode(bytes []byte, offsets *map[string]uint16) ([]byte, error) {
	bytes = append(bytes, wks.Address.To4()...)
	bytes = append(bytes, wks.Protocol)
	bytes = append(bytes, wks.BitMap...)

	return bytes, nil
}

func (wks WKS) String() string {
	return fmt.Sprintf("%s %d %x", wks.Address.String(), wks.Protocol, wks.BitMap)
}

func (ptr *PTR) Decode(buf []byte, offset int, rdlength int) (int, error) {
	var err error
	ptr.PTRDName, offset, err = decodeDomain(buf, offset)
	if err != nil {
		return offset, fmt.Errorf("PTR record: failed to decode PTRDName: %w", err)
	}

	return offset, nil
}

func (ptr *PTR) Encode(bytes []byte, offsets *map[string]uint16) ([]byte, error) {
	var err error
	bytes, err = encodeDomain(bytes, ptr.PTRDName, offsets)
	if err != nil {
		return nil, fmt.Errorf("PTR record: failed to encode PTRD %s: %w", ptr.PTRDName, err)
	}

	return bytes, nil
}

func (ptr PTR) String() string {
	return ptr.PTRDName
}

func (hinfo *HINFO) Decode(buf []byte, offset int, rdlength int) (int, error) {
	startOffset := offset
	endOffset := offset + rdlength
	if endOffset > len(buf) {
		return len(buf), &BufferOverflowError{Length: len(buf), Offset: endOffset}
	}

	currentOffset := offset
	var err error

	cpuLen, nextOffset, err := getU8(buf, currentOffset)
	if err != nil {
		return len(buf), fmt.Errorf("HINFO record: failed to read CPU length: %w", err)
	}
	currentOffset = nextOffset
	if currentOffset+int(cpuLen) > endOffset {
		return len(buf), &BufferOverflowError{Length: len(buf), Offset: currentOffset + int(cpuLen)}
	}
	cpuBytes, nextOffset, err := getSlice(buf, currentOffset, int(cpuLen))
	if err != nil {
		return len(buf), fmt.Errorf("HINFO record: failed to read CPU data: %w", err)
	}
	currentOffset = nextOffset
	hinfo.CPU = string(cpuBytes)

	osLen, nextOffset, err := getU8(buf, currentOffset)
	if err != nil {
		if currentOffset == endOffset {
			return len(buf), fmt.Errorf("HINFO record: missing OS length byte at offset %d (expected end: %d)", currentOffset, endOffset)
		}
		return len(buf), fmt.Errorf("HINFO record: failed to read OS length: %w", err)
	}
	currentOffset = nextOffset
	if currentOffset+int(osLen) > endOffset {
		return len(buf), &BufferOverflowError{Length: len(buf), Offset: currentOffset + int(osLen)}
	}
	osBytes, nextOffset, err := getSlice(buf, currentOffset, int(osLen))
	if err != nil {
		return len(buf), fmt.Errorf("HINFO record: failed to read OS data: %w", err)
	}
	currentOffset = nextOffset
	hinfo.OS = string(osBytes)

	if currentOffset != endOffset {
		return len(buf), fmt.Errorf("HINFO record: RDATA length mismatch, consumed %d bytes, expected %d", currentOffset-startOffset, rdlength)
	}

	return currentOffset, nil
}

func (hinfo *HINFO) Encode(bytes []byte, offsets *map[string]uint16) ([]byte, error) {
	if len(hinfo.CPU) > 255 {
		return nil, fmt.Errorf("HINFO record: CPU string length %d exceeds maximum 255", len(hinfo.CPU))
	}
	if len(hinfo.OS) > 255 {
		return nil, fmt.Errorf("HINFO record: OS string length %d exceeds maximum 255", len(hinfo.OS))
	}

	bytes = append(bytes, byte(len(hinfo.CPU)))
	bytes = append(bytes, []byte(hinfo.CPU)...)
	bytes = append(bytes, byte(len(hinfo.OS)))
	bytes = append(bytes, []byte(hinfo.OS)...)
	return bytes, nil
}

func (hinfo HINFO) String() string {
	return fmt.Sprintf("%q %q", hinfo.CPU, hinfo.OS)
}

func (minfo *MINFO) Decode(buf []byte, offset int, rdlength int) (int, error) {
	var err error

	minfo.RMailBx, offset, err = decodeDomain(buf, offset)
	if err != nil {
		return offset, fmt.Errorf("MINFO record: failed to decode RMailBx: %w", err)
	}

	minfo.EMailBx, offset, err = decodeDomain(buf, offset)
	if err != nil {
		return offset, fmt.Errorf("MINFO record: failed to decode EMailBx: %w", err)
	}

	return offset, nil
}

func (minfo *MINFO) Encode(bytes []byte, offsets *map[string]uint16) ([]byte, error) {
	var err error
	bytes, err = encodeDomain(bytes, minfo.RMailBx, offsets)
	if err != nil {
		return nil, fmt.Errorf("MINFO record: failed to encode RMailBx %s: %w", minfo.RMailBx, err)
	}

	bytes, err = encodeDomain(bytes, minfo.EMailBx, offsets)
	if err != nil {
		return nil, fmt.Errorf("MINFO record: failed to encode EMailBx %s: %w", minfo.EMailBx, err)
	}

	return bytes, nil
}

func (minfo MINFO) String() string {
	return minfo.RMailBx + " " + minfo.EMailBx
}

func (mx *MX) Decode(buf []byte, offset int, rdlength int) (int, error) {
	var err error
	mx.Preference, offset, err = getU16(buf, offset)
	if err != nil {
		return offset, fmt.Errorf("MX record: failed to decode Preference: %w", err)
	}

	mx.Exchange, offset, err = decodeDomain(buf, offset)
	if err != nil {
		return offset, fmt.Errorf("MX record: failed to decode Exchange: %w", err)
	}

	return offset, nil
}

func (mx *MX) Encode(bytes []byte, offsets *map[string]uint16) ([]byte, error) {
	var err error
	bytes = binary.BigEndian.AppendUint16(bytes, mx.Preference)
	bytes, err = encodeDomain(bytes, mx.Exchange, offsets)
	if err != nil {
		return nil, fmt.Errorf("MX record: failed to encode Exchange %s: %w", mx.Exchange, err)
	}

	return bytes, nil
}

func (mx MX) String() string {
	return fmt.Sprintf("%d %s", mx.Preference, mx.Exchange)
}

func (txt *TXT) Decode(buf []byte, offset int, rdlength int) (int, error) {
	txt.TxtData = make([]string, 0, 1)
	endOffset := offset + rdlength
	if endOffset > len(buf) {
		return len(buf), &BufferOverflowError{Length: len(buf), Offset: endOffset}
	}

	currentOffset := offset
	for currentOffset < endOffset {
		strLen, nextOffsetAfterLen, err := getU8(buf, currentOffset)
		if err != nil {
			return len(buf), fmt.Errorf("TXT record: failed to read string length byte: %w", err)
		}

		nextOffsetAfterData := nextOffsetAfterLen + int(strLen)
		if nextOffsetAfterData > endOffset {
			return len(buf), fmt.Errorf("TXT record: string segment length %d exceeds RDLENGTH boundary %d", strLen, endOffset)
		}

		strBytes, actualNextOffsetAfterData, err := getSlice(buf, nextOffsetAfterLen, int(strLen))
		if err != nil {
			return len(buf), fmt.Errorf("TXT record: failed to read string data (length %d): %w", strLen, err)
		}

		txt.TxtData = append(txt.TxtData, string(strBytes))
		currentOffset = actualNextOffsetAfterData
	}

	if currentOffset != endOffset {
		return len(buf), fmt.Errorf("TXT record: RDATA parsing finished at offset %d, but expected end at %d", currentOffset, endOffset)
	}

	return currentOffset, nil
}

func (txt *TXT) Encode(bytes []byte, offsets *map[string]uint16) ([]byte, error) {
	for _, s := range txt.TxtData {
		if len(s) > 255 {
			return nil, fmt.Errorf("TXT record: string segment length %d exceeds maximum 255", len(s))
		}
		bytes = append(bytes, byte(len(s)))
		bytes = append(bytes, []byte(s)...)
	}
	return bytes, nil
}

func (txt TXT) String() string {
	quoted := make([]string, len(txt.TxtData))
	for i, s := range txt.TxtData {
		quoted[i] = fmt.Sprintf("%q", s)
	}
	return strings.Join(quoted, " ")
}

func (r *Reserved) Decode(buf []byte, offset int, rdlength int) (int, error) {
	var err error
	r.Bytes, offset, err = getSlice(buf, offset, int(rdlength))
	if err != nil {
		return offset, fmt.Errorf("reserved record: failed to read data: %w", err)
	}

	return offset, err
}

func (r *Reserved) Encode(bytes []byte, offsets *map[string]uint16) ([]byte, error) {
	return append(bytes, r.Bytes...), nil
}

func (r Reserved) String() string {
	return fmt.Sprintf("[Reserved Data: %x]", r.Bytes)
}

// Decode decodes a resource record from buf at the offset.
func (r *ResourceRecord) Decode(buf []byte, offset int) (int, error) {
	var err error
	r.Name, offset, err = decodeDomain(buf, offset)
	if err != nil {
		return offset, fmt.Errorf("rr decode: failed to decode record name: %w", err)
	}

	var rtype uint16
	rtype, offset, err = getU16(buf, offset)
	if err != nil {
		return offset, fmt.Errorf("rr decode: failed to decode RType for %s: %w", r.Name, err)
	}
	r.RType = DNSType(rtype)

	var rclass uint16
	rclass, offset, err = getU16(buf, offset)
	if err != nil {
		return offset, fmt.Errorf("rr decode: failed to decode RClass for %s: %w", r.Name, err)
	}
	r.RClass = DNSClass(rclass)

	r.TTL, offset, err = getU32(buf, offset)
	if err != nil {
		return offset, fmt.Errorf("rr decode: failed to decode TTL for %s: %w", r.Name, err)
	}

	r.RDLength, offset, err = getU16(buf, offset)
	if err != nil {
		return offset, fmt.Errorf("rr decode: failed to decode RDLength for %s: %w", r.Name, err)
	}

	switch r.RType {
	case 1:
		r.RData = &A{}
	case 2:
		r.RData = &NS{}
	case 3:
		r.RData = &MD{}
	case 4:
		r.RData = &MF{}
	case 5:
		r.RData = &CNAME{}
	case 6:
		r.RData = &SOA{}
	case 7:
		r.RData = &MB{}
	case 8:
		r.RData = &MG{}
	case 9:
		r.RData = &MR{}
	case 10:
		r.RData = &NULL{}
	case 11:
		r.RData = &WKS{}
	case 12:
		r.RData = &PTR{}
	case 13:
		r.RData = &HINFO{}
	case 14:
		r.RData = &MINFO{}
	case 15:
		r.RData = &MX{}
	case 16:
		r.RData = &TXT{}
	default:
		r.RData = &Reserved{}
	}

	if r.RData != nil {
		offset, err = r.RData.Decode(buf, offset, int(r.RDLength))
		if err != nil {
			return offset, fmt.Errorf("rr decode: failed to decode RData for %s (%s): %w", r.Name, r.RType.String(), err)
		}
	}

	return offset, nil
}

// Encode encdoes a resource record and returns the input bytes appened.
func (r *ResourceRecord) Encode(bytes []byte, offsets *map[string]uint16) ([]byte, error) {
	var err error
	bytes, err = encodeDomain(bytes, r.Name, offsets)
	if err != nil {
		return nil, fmt.Errorf("rr encode: failed to encode record name %s: %w", r.Name, err)
	}

	bytes = binary.BigEndian.AppendUint16(bytes, uint16(r.RType))
	bytes = binary.BigEndian.AppendUint16(bytes, uint16(r.RClass))
	bytes = binary.BigEndian.AppendUint32(bytes, r.TTL)

	rdata_start := len(bytes)
	bytes = binary.BigEndian.AppendUint16(bytes, 0)
	bytes, err = r.RData.Encode(bytes, offsets)
	if err != nil {
		return nil, fmt.Errorf("rr encode: failed to encode RData for %s (%s): %w", r.Name, r.RType.String(), err)
	}

	rdata_length := uint16(len(bytes) - rdata_start - 2)
	binary.BigEndian.PutUint16(bytes[rdata_start:rdata_start+2], rdata_length)

	return bytes, nil
}