package checker

import (
	"context"
	"fmt"
	"net"
	"strings"
	"time"

	"github.com/miekg/dns"
)

const dnsTimeout = 5 * time.Second

// proto is "" for UDP or "tcp"; server must already include a port.
// RD is forced off: this checker only talks to authoritative servers.
func dnsExchange(ctx context.Context, proto, server string, q dns.Question, edns bool) (*dns.Msg, time.Duration, error) {
	client := dns.Client{Net: proto, Timeout: dnsTimeout}

	m := new(dns.Msg)
	m.Id = dns.Id()
	m.Question = []dns.Question{q}
	m.RecursionDesired = false
	if edns {
		m.SetEdns0(4096, true)
	}

	if deadline, ok := ctx.Deadline(); ok {
		if d := time.Until(deadline); d > 0 && d < client.Timeout {
			client.Timeout = d
		}
	}

	r, rtt, err := client.Exchange(m, server)
	if err != nil {
		return nil, rtt, err
	}
	if r == nil {
		return nil, rtt, fmt.Errorf("nil response from %s", server)
	}
	return r, rtt, nil
}

// Brackets IPv6 literals so the result is dialable by net.Dial.
func hostPort(host, port string) string {
	if ip := net.ParseIP(host); ip != nil && ip.To4() == nil {
		return "[" + host + "]:" + port
	}
	host = strings.TrimSuffix(host, ".")
	return host + ":" + port
}

func resolveHost(ctx context.Context, host string) ([]string, error) {
	var resolver net.Resolver
	addrs, err := resolver.LookupHost(ctx, strings.TrimSuffix(host, "."))
	if err != nil {
		return nil, err
	}
	return addrs, nil
}

func querySOA(ctx context.Context, proto, server, zone string) (soa *dns.SOA, aa bool, rtt time.Duration, err error) {
	q := dns.Question{Name: dns.Fqdn(zone), Qtype: dns.TypeSOA, Qclass: dns.ClassINET}
	r, rtt, err := dnsExchange(ctx, proto, server, q, false)
	if err != nil {
		return nil, false, rtt, err
	}
	if r.Rcode != dns.RcodeSuccess {
		return nil, r.Authoritative, rtt, fmt.Errorf("server answered %s", dns.RcodeToString[r.Rcode])
	}
	for _, rr := range r.Answer {
		if t, ok := rr.(*dns.SOA); ok {
			return t, r.Authoritative, rtt, nil
		}
	}
	// Some servers place the SOA in the Authority section instead of Answer.
	for _, rr := range r.Ns {
		if t, ok := rr.(*dns.SOA); ok {
			return t, r.Authoritative, rtt, nil
		}
	}
	return nil, r.Authoritative, rtt, fmt.Errorf("no SOA in answer section")
}

func queryNSAt(ctx context.Context, server, zone string) ([]string, error) {
	q := dns.Question{Name: dns.Fqdn(zone), Qtype: dns.TypeNS, Qclass: dns.ClassINET}
	r, _, err := dnsExchange(ctx, "", server, q, false)
	if err != nil {
		return nil, err
	}
	if r.Rcode != dns.RcodeSuccess {
		return nil, fmt.Errorf("server answered %s", dns.RcodeToString[r.Rcode])
	}
	var out []string
	for _, rr := range r.Answer {
		if t, ok := rr.(*dns.NS); ok {
			out = append(out, strings.ToLower(dns.Fqdn(t.Ns)))
		}
	}
	return out, nil
}

// Falls back to TCP on UDP failure: some middleboxes drop large UDP packets
// carrying OPT but let TCP/53 through (RFC 7766 mandates TCP fallback).
func probeEDNS0(ctx context.Context, server, zone string) error {
	q := dns.Question{Name: dns.Fqdn(zone), Qtype: dns.TypeSOA, Qclass: dns.ClassINET}
	r, _, err := dnsExchange(ctx, "", server, q, true)
	if err != nil {
		rt, _, terr := dnsExchange(ctx, "tcp", server, q, true)
		if terr != nil {
			return fmt.Errorf("EDNS0 query failed over UDP (%v) and TCP (%w)", err, terr)
		}
		r = rt
	}
	if r.Rcode == dns.RcodeFormatError {
		return fmt.Errorf("server returned FORMERR on EDNS0 query")
	}
	if r.Rcode != dns.RcodeSuccess {
		return fmt.Errorf("server answered %s on EDNS0 query", dns.RcodeToString[r.Rcode])
	}
	// RFC 6891 requires the OPT pseudo-RR to be echoed in the response.
	if r.IsEdns0() == nil {
		return fmt.Errorf("server stripped the EDNS0 OPT record from its response")
	}
	return nil
}

// First parent server returning a non-empty referral wins.
func parentReferral(ctx context.Context, zone string) ([]string, error) {
	zone = dns.Fqdn(zone)
	labels := dns.SplitDomainName(zone)
	if len(labels) < 2 {
		return nil, fmt.Errorf("zone %q has no parent", zone)
	}
	parent := dns.Fqdn(strings.Join(labels[1:], "."))

	resolver := net.Resolver{}
	nss, err := resolver.LookupNS(ctx, strings.TrimSuffix(parent, "."))
	if err != nil {
		return nil, fmt.Errorf("resolving NS of parent zone %q: %w", parent, err)
	}

	var lastErr error
	seen := map[string]bool{}
	var out []string
	for _, ns := range nss {
		addrs, rerr := resolver.LookupHost(ctx, strings.TrimSuffix(ns.Host, "."))
		if rerr != nil || len(addrs) == 0 {
			lastErr = rerr
			continue
		}
		for _, a := range addrs {
			srv := hostPort(a, "53")
			q := dns.Question{Name: zone, Qtype: dns.TypeNS, Qclass: dns.ClassINET}
			r, _, qerr := dnsExchange(ctx, "", srv, q, true)
			if qerr != nil {
				lastErr = qerr
				continue
			}
			if r.Rcode != dns.RcodeSuccess {
				lastErr = fmt.Errorf("parent %s answered %s", ns.Host, dns.RcodeToString[r.Rcode])
				continue
			}
			collect := func(records []dns.RR) {
				for _, rr := range records {
					if t, ok := rr.(*dns.NS); ok {
						if strings.EqualFold(strings.TrimSuffix(t.Header().Name, "."), strings.TrimSuffix(zone, ".")) {
							name := strings.ToLower(dns.Fqdn(t.Ns))
							if !seen[name] {
								seen[name] = true
								out = append(out, name)
							}
						}
					}
				}
			}
			collect(r.Answer)
			collect(r.Ns)
			if len(out) > 0 {
				return out, nil
			}
		}
	}
	if lastErr != nil {
		return nil, lastErr
	}
	return nil, fmt.Errorf("no parent server returned a delegation for %s", zone)
}