package daemonsetdns

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

	guuid "github.com/google/uuid"
	"github.com/miekg/dns"

	"edge-infra.dev/pkg/lib/fog"
	"edge-infra.dev/pkg/sds/k8s/iplookup"
)

const dnsTimeout = 30 * time.Second

var (
	errDNSLen  = errors.New("fewer DNS labels than expected")
	errDNSName = errors.New("unexpected DNS name format")
)

// Extracts the node name, daemonset name and namespace of a given namespace
// from the DNS request following the <nodeName>.<daemonsetName>.<Namespace>.pod-locator.
// format. domainName must be a syntactically valid domain name.
func extractName(domainName string) (nodeName string, daemonset string, namespace string, err error) {
	nameParts := dns.SplitDomainName(domainName)
	labelCount := len(nameParts)
	// DNS request must have at least 4 labels. The node name can optionally be more than one label
	if labelCount < 4 {
		return "", "", "", errDNSLen
	}

	// top-level domain for the daemonset dns should always be pod-locator
	if nameParts[labelCount-1] != "pod-locator" {
		return "", "", "", errDNSName
	}

	namespace = nameParts[labelCount-2]
	daemonset = nameParts[labelCount-3]
	nodeNameParts := nameParts[:(labelCount - 3)]
	nodeName = strings.Join(nodeNameParts, ".")

	return nodeName, daemonset, namespace, nil
}

func handleResponse(ipclient *iplookup.IPLookup) func(w dns.ResponseWriter, r *dns.Msg) {
	return func(w dns.ResponseWriter, r *dns.Msg) {
		ctx, cancelFunc := context.WithTimeout(context.Background(), dnsTimeout)
		defer cancelFunc()

		log := fog.New().
			WithName("daemonsetdns").
			WithName("resolver").
			WithValues(fog.OperationFields(guuid.New().String())...)

		ctx = fog.IntoContext(ctx, log)

		var rr dns.RR

		m := new(dns.Msg)
		m.SetReply(r)

		dnsName := r.Question[0].Name

		log.Info("Incoming DNS request", "Name", dnsName, "type", r.Question[0].Qtype)

		nodeName, daemonset, namespace, err := extractName(dnsName)
		if err != nil {
			log.Error(err, "Unable to extract daemonset properties from dns request")
			if err := w.WriteMsg(m); err != nil {
				log.Error(err, "Failed to write response")
			}
			return
		}

		query := iplookup.Query{
			Namespace: namespace,
			Daemonset: daemonset,
		}

		log.Info("Discovering IP address", "Node Name", nodeName, "Namespace", namespace, "Daemonset", daemonset)

		ip, err := ipclient.GetDaemonsetPodIP(ctx, nodeName, query)
		if err != nil {
			log.Error(err, "Error occurred when discovering IP")
			if err := w.WriteMsg(m); err != nil {
				log.Error(err, "Failed to write response")
			}
			return
		}

		// Create DNS answer section
		if ip.To4() != nil {
			rr = &dns.A{
				Hdr: dns.RR_Header{Name: dnsName, Rrtype: dns.TypeA, Class: dns.ClassINET, Ttl: 0},
				A:   ip,
			}
		} else {
			rr = &dns.AAAA{
				Hdr:  dns.RR_Header{Name: dnsName, Rrtype: dns.TypeAAAA, Class: dns.ClassINET, Ttl: 0},
				AAAA: ip,
			}
		}

		// Apply DNS answer section when question and answer record types match
		switch r.Question[0].Qtype {
		case dns.TypeA:
			if ip.To4() == nil {
				log.Error(nil, "Record type and returned ip version do not match")
				break
			}
			m.Answer = append(m.Answer, rr)
		default:
			log.Info("Unsupported request type")
		}

		// Write the DNS response
		if err := w.WriteMsg(m); err != nil {
			log.Error(err, "Failed to write response")
		}
	}
}

// Return a dns.Server for a given net protocol which can be used to find the
// ipv4 address of a pod in a daemonset
func Dnsserver(net string, ipclient *iplookup.IPLookup) *dns.Server {
	log := fog.New().WithName("daemonsetdns").WithName("serve")
	log.Info("Starting server", "protocol", net)
	dns.HandleFunc("pod-locator.", handleResponse(ipclient))

	server := &dns.Server{
		Addr: "[::]:5353", Net: net, TsigSecret: nil, ReusePort: false,
	}
	return server
}