package checker

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

	"github.com/miekg/dns"
)

const dnsTimeout = 5 * time.Second

// FallbackResolver is the resolver used when /etc/resolv.conf is missing or
// empty. It can be overridden at startup (e.g. via a CLI flag) so operators
// don't silently leak lookups to a third party.
var FallbackResolver = net.JoinHostPort("1.1.1.1", "53")

// dnsExchange sends a single query. proto="" uses UDP and retries over TCP on
// truncation; recursion controls the RD flag.
func dnsExchange(ctx context.Context, proto, server string, q dns.Question, recursion 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 = recursion
	m.SetEdns0(4096, true)

	if deadline, ok := ctx.Deadline(); 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)
	}
	if r.Truncated && proto == "" {
		tcpClient := dns.Client{Net: "tcp", Timeout: client.Timeout}
		if r2, _, err2 := tcpClient.Exchange(m, server); err2 == nil && r2 != nil {
			return r2, nil
		}
	}
	return r, nil
}

// systemResolver returns the first configured resolver of the local system.
func systemResolver() string {
	cfg, err := dns.ClientConfigFromFile("/etc/resolv.conf")
	if err != nil || len(cfg.Servers) == 0 {
		return FallbackResolver
	}
	return net.JoinHostPort(cfg.Servers[0], cfg.Port)
}

// hostPort returns "host:port", correctly bracketing IPv6 literals.
func hostPort(host, port string) string {
	host = strings.TrimSuffix(host, ".")
	return net.JoinHostPort(host, port)
}

// findReverseZone walks up the labels of fqdn until it finds a zone cut
// (SOA). Returns the apex FQDN and the list of "host:53" authoritative
// servers. The walk stops at the reverse-arpa apex (in-addr.arpa or
// ip6.arpa) so we never accept a non-reverse zone (or the root) as a match.
func findReverseZone(ctx context.Context, fqdn string) (apex string, servers []string, err error) {
	resolver := systemResolver()
	labels := dns.SplitDomainName(fqdn)
	for i := range labels {
		candidate := dns.Fqdn(strings.Join(labels[i:], "."))
		if !isReverseArpa(candidate) {
			break
		}
		q := dns.Question{Name: candidate, Qtype: dns.TypeSOA, Qclass: dns.ClassINET}
		r, rerr := dnsExchange(ctx, "", resolver, q, true)
		if rerr != nil {
			continue
		}
		if r.Rcode != dns.RcodeSuccess {
			continue
		}
		hasSOA := false
		for _, rr := range r.Answer {
			if _, ok := rr.(*dns.SOA); ok {
				hasSOA = true
				break
			}
		}
		if !hasSOA {
			continue
		}
		apex = candidate
		// NS resolution failures are non-fatal: we still located the zone,
		// and queryPTR will fall back to the system resolver. Returning an
		// error here would make reverseZoneRule wrongly report "zone not
		// found".
		servers, _ := resolveZoneNSAddrs(ctx, apex)
		return apex, servers, nil
	}
	return "", nil, fmt.Errorf("could not locate reverse zone of %s", fqdn)
}

// resolveZoneNSAddrs returns "host:53" entries for every NS of the zone.
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
}

// queryAtAuth sends q to the first reachable server of the list.
func queryAtAuth(ctx context.Context, servers []string, q dns.Question) (*dns.Msg, string, error) {
	var lastErr error
	for _, s := range servers {
		r, err := dnsExchange(ctx, "", s, q, false)
		if err != nil {
			lastErr = err
			continue
		}
		return r, s, nil
	}
	if lastErr == nil {
		lastErr = fmt.Errorf("no servers provided")
	}
	return nil, "", lastErr
}

// rcodeText returns the textual name of an rcode or a fallback string.
func rcodeText(r int) string {
	if s, ok := dns.RcodeToString[r]; ok {
		return s
	}
	return fmt.Sprintf("RCODE(%d)", r)
}

// lowerFQDN returns the canonical lowercase FQDN form of name.
func lowerFQDN(name string) string {
	return strings.ToLower(dns.Fqdn(name))
}

// reverseNameToIP decodes a reverse-arpa name back to a net.IP. It accepts
// both in-addr.arpa (IPv4) and ip6.arpa (IPv6). Returns nil if the name is
// malformed.
func reverseNameToIP(name string) net.IP {
	n := strings.ToLower(strings.TrimSuffix(dns.Fqdn(name), "."))

	switch {
	case strings.HasSuffix(n, ".in-addr.arpa"):
		labels := strings.Split(strings.TrimSuffix(n, ".in-addr.arpa"), ".")
		if len(labels) != 4 {
			return nil
		}
		// Reverse: "4.3.2.1" -> "1.2.3.4"
		out := make([]string, 4)
		for i, l := range labels {
			if _, err := strconv.Atoi(l); err != nil {
				return nil
			}
			out[3-i] = l
		}
		return net.ParseIP(strings.Join(out, "."))

	case strings.HasSuffix(n, ".ip6.arpa"):
		labels := strings.Split(strings.TrimSuffix(n, ".ip6.arpa"), ".")
		if len(labels) != 32 {
			return nil
		}
		// Reverse the nibbles and regroup into 8 × 4-hex blocks.
		var sb strings.Builder
		for i := len(labels) - 1; i >= 0; i-- {
			if len(labels[i]) != 1 {
				return nil
			}
			sb.WriteString(labels[i])
			if i > 0 && (len(labels)-i)%4 == 0 {
				sb.WriteByte(':')
			}
		}
		return net.ParseIP(sb.String())
	}
	return nil
}

// isReverseArpa reports whether name lies inside in-addr.arpa or ip6.arpa.
func isReverseArpa(name string) bool {
	n := lowerFQDN(name)
	return strings.HasSuffix(n, ".in-addr.arpa.") || strings.HasSuffix(n, ".ip6.arpa.")
}