package dns_api

import "github.com/miekg/dns"

type RRSet []dns.RR

func (s RRSet) Get(rrt uint16, name string) (rrset RRSet) {
	for _, rr := range s {
		if rr != nil && (rr.Header().Rrtype == rrt || rrt == dns.TypeANY) && (rr.Header().Name == name || name == "") {
			rrset = append(rrset, rr)
		}
	}

	return
}

func (s RRSet) Delete(rrt uint16, name string) {
	for i, rr := range s {
		if rr != nil && rr.Header().Rrtype == rrt && rr.Header().Name == name {
			s[i] = nil
		}
	}
}

func (s RRSet) SplitByType() (result []RRSet) {

	var types []uint16

	inTypes := func(i uint16) bool {
		for _, j := range types {
			if i == j {
				return true
			}
		}
		return false
	}

	for _, rr := range s {
		if rr != nil && !inTypes(rr.Header().Rrtype) {
			result = append(result, s.Get(rr.Header().Rrtype, ""))
			types = append(types, rr.Header().Rrtype)
		}
	}

	return
}

func (s RRSet) SplitByName() (result []RRSet) {
	var names []string

	inNames := func(i string) bool {
		for _, j := range names {
			if i == j {
				return true
			}
		}
		return false
	}

	for _, rr := range s {
		if rr != nil && !inNames(rr.Header().Name) {
			result = append(result, s.Get(dns.TypeANY, rr.Header().Name))
			names = append(names, rr.Header().Name)
		}
	}

	return
}

func (s RRSet) NotNil() (result RRSet) {
	result = make(RRSet, 0, len(s))
	for _, rr := range s {
		if rr != nil {
			result = append(result, rr)
		}
	}

	return
}

func (s RRSet) RemoveTypes(disallow ...uint16) (result RRSet) {

	inList := func(i uint16) bool {

		for _, j := range disallow {
			if i == j {
				return true
			}
		}
		return false
	}

	for _, rr := range s {
		if rr != nil && !inList(rr.Header().Rrtype) {
			result = append(result, rr)
		}
	}

	return
}

func (s RRSet) OnlyTypes(allow ...uint16) (result RRSet) {

	inList := func(i uint16) bool {

		for _, j := range allow {
			if i == j {
				return true
			}
		}
		return false
	}

	for _, rr := range s {
		if rr != nil && inList(rr.Header().Rrtype) {
			result = append(result, rr)
		}
	}

	return
}

func (s RRSet) Encode() (result []byte, err error) {
	var off int
	for _, rr := range s {
		buf := make([]byte, dns.Len(rr))
		off, err = dns.PackRR(rr, buf, 0, nil, false)
		if err != nil {
			return nil, err
		}
		result = append(result, buf[:off]...)
	}

	return result, nil
}

func DecodeRRSet(buf []byte) (result RRSet, err error) {
	for len(buf) > 0 {
		rr, off, err := dns.UnpackRR(buf, 0)
		if err != nil {
			return nil, err
		}
		buf = buf[off:]
		result = append(result, rr)
	}

	return result, nil
}