package exchanger

import (
	"context"
	"errors"
	"net"
	"strings"
	"time"

	"github.com/miekg/dns"
)

type systemExchanger struct {
	rawAddr  string
	timeout  time.Duration
	lookupIP func(ctx context.Context, network, host string) ([]net.IP, error)
}

func newSystemExchanger(rawAddr string, opts Options) Exchanger {
	if opts.timeout <= 0 {
		opts.timeout = 5 * time.Second
	}
	ex := &systemExchanger{
		rawAddr: rawAddr,
		timeout: opts.timeout,
	}
	ex.lookupIP = net.DefaultResolver.LookupIP
	return ex
}

func (ex *systemExchanger) Exchange(ctx context.Context, msg []byte) ([]byte, error) {
	mq := new(dns.Msg)
	if err := mq.Unpack(msg); err != nil {
		return nil, err
	}

	mr := new(dns.Msg)
	mr.SetReply(mq)
	mr.RecursionAvailable = true

	if len(mq.Question) == 0 {
		return mr.Pack()
	}

	q := mq.Question[0]
	network := "ip"
	switch q.Qtype {
	case dns.TypeA:
		network = "ip4"
	case dns.TypeAAAA:
		network = "ip6"
	default:
		return mr.Pack()
	}

	if ex.timeout > 0 {
		var cancel context.CancelFunc
		ctx, cancel = context.WithTimeout(ctx, ex.timeout)
		defer cancel()
	}

	host := strings.TrimSuffix(q.Name, ".")
	ips, err := ex.lookupIP(ctx, network, host)
	if err != nil {
		var dnsErr *net.DNSError
		if !errors.As(err, &dnsErr) || !dnsErr.IsNotFound {
			return nil, err
		}
	}

	for _, ip := range ips {
		switch q.Qtype {
		case dns.TypeA:
			if ip4 := ip.To4(); ip4 != nil {
				mr.Answer = append(mr.Answer, &dns.A{
					Hdr: dns.RR_Header{
						Name:   q.Name,
						Rrtype: dns.TypeA,
						Class:  q.Qclass,
					},
					A: ip4,
				})
			}
		case dns.TypeAAAA:
			if ip.To4() == nil {
				if ip16 := ip.To16(); ip16 != nil {
					mr.Answer = append(mr.Answer, &dns.AAAA{
						Hdr: dns.RR_Header{
							Name:   q.Name,
							Rrtype: dns.TypeAAAA,
							Class:  q.Qclass,
						},
						AAAA: ip16,
					})
				}
			}
		}
	}

	return mr.Pack()
}

func (ex *systemExchanger) String() string {
	if ex.rawAddr == "" {
		return "system"
	}
	return ex.rawAddr
}