package dns

import (
	"fmt"
	"strings"
	"time"

	mdns "github.com/miekg/dns"
)

var defaultTypes = []string{"A", "AAAA", "MX", "TXT", "NS", "CNAME", "SOA", "CAA", "SRV"}

const defaultResolver = "8.8.8.8:53"

var typeMap = map[string]uint16{
	"A":     mdns.TypeA,
	"AAAA":  mdns.TypeAAAA,
	"MX":    mdns.TypeMX,
	"TXT":   mdns.TypeTXT,
	"NS":    mdns.TypeNS,
	"CNAME": mdns.TypeCNAME,
	"SOA":   mdns.TypeSOA,
	"CAA":   mdns.TypeCAA,
	"SRV":   mdns.TypeSRV,
}

func Lookup(req LookupRequest) (*LookupResponse, error) {
	domain := req.Domain
	if !strings.HasSuffix(domain, ".") {
		domain += "."
	}

	types := req.Types
	if len(types) == 0 {
		types = defaultTypes
	}

	start := time.Now()
	records := make(map[string][]Record)
	hasDNSSEC := false

	c := new(mdns.Client)
	c.Net = "udp"

	for _, t := range types {
		qtype, ok := typeMap[strings.ToUpper(t)]
		if !ok {
			continue
		}

		m := new(mdns.Msg)
		m.SetQuestion(domain, qtype)
		m.SetEdns0(4096, true)

		r, _, err := c.Exchange(m, defaultResolver)
		if err != nil {
			continue
		}

		if r.AuthenticatedData {
			hasDNSSEC = true
		}

		recs := parseRecords(r.Answer, strings.ToUpper(t))
		if len(recs) > 0 {
			records[strings.ToUpper(t)] = recs
		}
	}

	return &LookupResponse{
		Domain:      req.Domain,
		Records:     records,
		Resolver:    "8.8.8.8",
		QueryTimeMs: time.Since(start).Milliseconds(),
		DNSSEC:      hasDNSSEC,
		Timestamp:   time.Now().UTC(),
	}, nil
}

func parseRecords(answers []mdns.RR, qtype string) []Record {
	var records []Record
	for _, ans := range answers {
		rec := Record{TTL: ans.Header().Ttl}
		switch rr := ans.(type) {
		case *mdns.A:
			rec.Value = rr.A.String()
		case *mdns.AAAA:
			rec.Value = rr.AAAA.String()
		case *mdns.MX:
			rec.Value = strings.TrimSuffix(rr.Mx, ".")
			rec.Priority = rr.Preference
		case *mdns.TXT:
			rec.Value = strings.Join(rr.Txt, " ")
		case *mdns.NS:
			rec.Value = strings.TrimSuffix(rr.Ns, ".")
		case *mdns.CNAME:
			rec.Value = strings.TrimSuffix(rr.Target, ".")
		case *mdns.SOA:
			rec.Value = fmt.Sprintf("%s %s %d %d %d %d %d",
				rr.Ns, rr.Mbox, rr.Serial, rr.Refresh, rr.Retry, rr.Expire, rr.Minttl)
		case *mdns.CAA:
			rec.Value = fmt.Sprintf("%d %s \"%s\"", rr.Flag, rr.Tag, rr.Value)
		case *mdns.SRV:
			rec.Value = fmt.Sprintf("%d %d %d %s", rr.Priority, rr.Weight, rr.Port, strings.TrimSuffix(rr.Target, "."))
			rec.Priority = rr.Priority
		default:
			continue
		}
		records = append(records, rec)
	}
	return records
}