package checker

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

	"github.com/miekg/dns"
)

// year68 mirrors the constant from miekg/dns used to wrap RRSIG validity
// periods around 2^32 seconds (≈68 years), as in the adlin checker.
const year68 = int64(1 << 31)

// dnsTimeout is the per-query deadline used by every helper here.
const dnsTimeout = 5 * time.Second

// dnsExchange sends a single query to the given server using the requested
// transport ("" for UDP, "tcp"). The server address must already include a
// port. RecursionDesired 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, 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)
	}

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

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

// hostPort returns "host:port", correctly bracketing IPv6 literals.
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
}

// resolveHost resolves an NS hostname to its A and AAAA addresses using the
// system resolver. It is used as a fallback when no glue is provided by the
// parent for an out-of-bailiwick NS.
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
}

// findParentZone walks up the labels of fqdn until it finds the closest
// enclosing zone (the one that has its own SOA), and returns the FQDN of
// that zone along with its authoritative server addresses (resolved from
// its NS RRset). The walk stops as soon as a SOA query at the system
// resolver returns NOERROR with an answer.
//
// If hintParent is non-empty, it is used as the assumed parent and we only
// resolve its NS, this matches happyDomain's data model where the parent
// zone is known.
func findParentZone(ctx context.Context, fqdn, hintParent string) (zone string, servers []string, err error) {
	zone = dns.Fqdn(hintParent)
	if zone == "" || zone == "." {
		// Walk up.
		labels := dns.SplitDomainName(fqdn)
		if len(labels) == 0 {
			return "", nil, fmt.Errorf("cannot derive parent of %q", fqdn)
		}
		zone = dns.Fqdn(strings.Join(labels[1:], "."))
	}

	servers, err = resolveZoneNSAddrs(ctx, zone)
	if err != nil {
		return "", nil, fmt.Errorf("resolving NS of parent zone %q: %w", zone, err)
	}
	if len(servers) == 0 {
		return "", nil, fmt.Errorf("parent zone %q has no resolvable NS", zone)
	}
	return zone, servers, nil
}

// resolveZoneNSAddrs returns the list of "host:53" entries for every NS of
// the given zone, as seen by the system resolver. It is used to discover the
// parent's authoritative servers.
func resolveZoneNSAddrs(ctx context.Context, zone string) ([]string, error) {
	var resolver net.Resolver
	nss, err := resolver.LookupNS(ctx, strings.TrimSuffix(zone, "."))
	if err != nil {
		return nil, err
	}

	var out []string
	for _, ns := range nss {
		addrs, err := resolver.LookupHost(ctx, strings.TrimSuffix(ns.Host, "."))
		if err != nil || len(addrs) == 0 {
			continue
		}
		for _, a := range addrs {
			out = append(out, hostPort(a, "53"))
		}
	}
	return out, nil
}

// queryDelegation queries the given parent server for the NS RRset of fqdn
// and extracts the advertised NS names plus any glue records found in the
// Additional section. The query is sent without RD; the response is the
// classical "referral" packet.
func queryDelegation(ctx context.Context, parentServer, fqdn string) (ns []string, glue map[string][]string, msg *dns.Msg, err error) {
	q := dns.Question{Name: dns.Fqdn(fqdn), Qtype: dns.TypeNS, Qclass: dns.ClassINET}

	msg, err = dnsExchange(ctx, "", parentServer, q, true)
	if err != nil {
		return nil, nil, nil, err
	}
	if msg.Rcode != dns.RcodeSuccess {
		return nil, nil, msg, fmt.Errorf("parent answered %s", dns.RcodeToString[msg.Rcode])
	}

	glue = map[string][]string{}

	collect := func(records []dns.RR) {
		for _, rr := range records {
			switch t := rr.(type) {
			case *dns.NS:
				if strings.EqualFold(strings.TrimSuffix(t.Header().Name, "."), strings.TrimSuffix(fqdn, ".")) {
					ns = append(ns, strings.ToLower(dns.Fqdn(t.Ns)))
				}
			case *dns.A:
				name := strings.ToLower(dns.Fqdn(t.Header().Name))
				glue[name] = append(glue[name], t.A.String())
			case *dns.AAAA:
				name := strings.ToLower(dns.Fqdn(t.Header().Name))
				glue[name] = append(glue[name], t.AAAA.String())
			}
		}
	}
	collect(msg.Answer)
	collect(msg.Ns)
	collect(msg.Extra)
	return
}

// queryDS asks the parent server for the DS RRset of fqdn and returns the
// DS records plus any RRSIGs found in the same section.
func queryDS(ctx context.Context, parentServer, fqdn string) (ds []*dns.DS, sigs []*dns.RRSIG, err error) {
	q := dns.Question{Name: dns.Fqdn(fqdn), Qtype: dns.TypeDS, Qclass: dns.ClassINET}

	r, err := dnsExchange(ctx, "tcp", parentServer, q, true)
	if err != nil {
		return nil, nil, err
	}
	if r.Rcode != dns.RcodeSuccess {
		return nil, nil, fmt.Errorf("parent answered %s for DS", dns.RcodeToString[r.Rcode])
	}

	for _, rr := range r.Answer {
		switch t := rr.(type) {
		case *dns.DS:
			ds = append(ds, t)
		case *dns.RRSIG:
			sigs = append(sigs, t)
		}
	}
	return
}

// querySOA asks the given authoritative server for the SOA of fqdn and
// returns the SOA record plus the AA flag from the response header.
func querySOA(ctx context.Context, proto, server, fqdn string) (soa *dns.SOA, aa bool, err error) {
	q := dns.Question{Name: dns.Fqdn(fqdn), Qtype: dns.TypeSOA, Qclass: dns.ClassINET}
	r, err := dnsExchange(ctx, proto, server, q, false)
	if err != nil {
		return nil, false, err
	}
	if r.Rcode != dns.RcodeSuccess {
		return nil, r.Authoritative, 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, nil
		}
	}
	return nil, r.Authoritative, fmt.Errorf("no SOA in answer section")
}

// queryNSAt asks the given authoritative server for the NS RRset of fqdn.
func queryNSAt(ctx context.Context, server, fqdn string) ([]string, error) {
	q := dns.Question{Name: dns.Fqdn(fqdn), 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
}

// queryAddrsAt asks an authoritative server for the A and AAAA records of
// host (typically an in-bailiwick NS hostname).
func queryAddrsAt(ctx context.Context, server, host string) ([]string, error) {
	var out []string
	for _, qt := range []uint16{dns.TypeA, dns.TypeAAAA} {
		r, err := dnsExchange(ctx, "", server, dns.Question{Name: dns.Fqdn(host), Qtype: qt, Qclass: dns.ClassINET}, false)
		if err != nil {
			continue
		}
		if r.Rcode != dns.RcodeSuccess {
			continue
		}
		for _, rr := range r.Answer {
			switch t := rr.(type) {
			case *dns.A:
				out = append(out, t.A.String())
			case *dns.AAAA:
				out = append(out, t.AAAA.String())
			}
		}
	}
	return out, nil
}

// queryDNSKEY asks the given child server for the DNSKEY RRset of fqdn.
func queryDNSKEY(ctx context.Context, server, fqdn string) ([]*dns.DNSKEY, error) {
	q := dns.Question{Name: dns.Fqdn(fqdn), Qtype: dns.TypeDNSKEY, Qclass: dns.ClassINET}
	r, err := dnsExchange(ctx, "tcp", server, q, true)
	if err != nil {
		return nil, err
	}
	if r.Rcode != dns.RcodeSuccess {
		return nil, fmt.Errorf("server answered %s for DNSKEY", dns.RcodeToString[r.Rcode])
	}
	var out []*dns.DNSKEY
	for _, rr := range r.Answer {
		if t, ok := rr.(*dns.DNSKEY); ok {
			out = append(out, t)
		}
	}
	return out, nil
}

// dsEqual returns true when two DS records refer to the same key material.
func dsEqual(a, b *dns.DS) bool {
	return a.KeyTag == b.KeyTag &&
		a.Algorithm == b.Algorithm &&
		a.DigestType == b.DigestType &&
		strings.EqualFold(a.Digest, b.Digest)
}

// validityWindow returns a human-readable explanation of why a signature is
// outside its validity period, mirroring the year68 logic from the adlin
// checker.
func validityWindow(sig *dns.RRSIG) string {
	utc := time.Now().UTC().Unix()

	modi := (int64(sig.Inception) - utc) / year68
	ti := int64(sig.Inception) + modi*year68

	mode := (int64(sig.Expiration) - utc) / year68
	te := int64(sig.Expiration) + mode*year68

	if ti > utc {
		return "signature not yet valid"
	} else if utc > te {
		return "signature expired"
	}
	return "signature outside its validity window"
}