package checker

import (
	"bytes"
	"context"
	"crypto/tls"
	"encoding/base64"
	"fmt"
	"io"
	"net"
	"net/http"
	"net/url"
	"sort"
	"strings"
	"time"

	"github.com/miekg/dns"
)

// Default per-query deadline. Public resolvers are expected to answer in far
// less; anything slower is either unreachable or too flaky to be useful for
// a propagation check.
const dnsTimeout = 5 * time.Second

// ednsUDPSize is the EDNS0 advertised UDP buffer size (RFC 6891). 4096 is the
// commonly accepted ceiling for unfragmented responses across the public
// Internet.
const ednsUDPSize = 4096

// maxDoHResponseBytes caps how much we read from a DoH response body before
// giving up. 64 KiB comfortably exceeds any well-behaved DNS message and
// keeps a hostile server from streaming us junk.
const maxDoHResponseBytes = 64 * 1024

// dohClient is a package-wide HTTP client reused across DoH probes. Setting
// the tls config and timeouts once keeps the allocations down when we
// probe dozens of resolvers concurrently.
var dohClient = &http.Client{
	Timeout: dnsTimeout + 2*time.Second,
	Transport: &http.Transport{
		TLSClientConfig: &tls.Config{
			MinVersion: tls.VersionTLS12,
		},
		TLSHandshakeTimeout:   dnsTimeout,
		ResponseHeaderTimeout: dnsTimeout,
		ExpectContinueTimeout: 1 * time.Second,
		DisableKeepAlives:     false,
		MaxIdleConnsPerHost:   4,
	},
}

// queryResult is what a single DNS exchange returns: deliberately flatter
// than *dns.Msg so the collector can stay protocol-agnostic.
type queryResult struct {
	Rcode   int
	Answer  []dns.RR
	AD      bool
	Latency time.Duration
}

// queryResolver dispatches a question to a single public resolver over the
// requested transport. It forces RD=1 (we want the resolver to recurse) and
// CD=0 (we want the resolver to validate DNSSEC when it can). AD=1 is
// requested so AD-capable resolvers signal validation in the response.
func queryResolver(ctx context.Context, r Resolver, tr Transport, name string, qtype uint16) (*queryResult, error) {
	q := dns.Question{Name: dns.Fqdn(name), Qtype: qtype, Qclass: dns.ClassINET}

	m := new(dns.Msg)
	m.Id = dns.Id()
	m.Question = []dns.Question{q}
	m.RecursionDesired = true
	m.CheckingDisabled = false
	m.AuthenticatedData = true
	m.SetEdns0(ednsUDPSize, true)

	switch tr {
	case TransportUDP:
		return exchangeUDPOrTCP(ctx, m, r.IP+":53", "udp")
	case TransportTCP:
		return exchangeUDPOrTCP(ctx, m, r.IP+":53", "tcp")
	case TransportDoT:
		if r.DoTHost == "" {
			return nil, fmt.Errorf("no DoT endpoint for %s", r.ID)
		}
		return exchangeDoT(ctx, m, r.IP, r.DoTHost)
	case TransportDoH:
		if r.DoHURL == "" {
			return nil, fmt.Errorf("no DoH endpoint for %s", r.ID)
		}
		return exchangeDoH(ctx, m, r.DoHURL)
	default:
		return nil, fmt.Errorf("unknown transport %q", tr)
	}
}

// exchangeUDPOrTCP performs the exchange over plain DNS. miekg/dns handles
// truncation and retries transparently when we say "udp" + EDNS; for larger
// answers we switch to TCP explicitly.
func exchangeUDPOrTCP(ctx context.Context, m *dns.Msg, server, proto string) (*queryResult, error) {
	client := dns.Client{Net: proto, Timeout: dnsTimeout}
	if deadline, ok := ctx.Deadline(); ok {
		if d := time.Until(deadline); d > 0 && d < client.Timeout {
			client.Timeout = d
		}
	}

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

	// Truncated UDP answers force a retry over TCP per RFC 5966.
	if proto == "udp" && r.Truncated {
		tcpClient := dns.Client{Net: "tcp", Timeout: dnsTimeout}
		if r2, rtt2, err2 := tcpClient.ExchangeContext(ctx, m, server); err2 == nil && r2 != nil {
			return &queryResult{
				Rcode: r2.Rcode, Answer: r2.Answer,
				AD: r2.AuthenticatedData, Latency: rtt2,
			}, nil
		}
	}

	return &queryResult{
		Rcode: r.Rcode, Answer: r.Answer,
		AD: r.AuthenticatedData, Latency: rtt,
	}, nil
}

// exchangeDoT opens a TLS connection to ip:853 and speaks plain DNS over it,
// validating the certificate against sni. miekg/dns' Client does TLS natively
// when Net="tcp-tls".
func exchangeDoT(ctx context.Context, m *dns.Msg, ip, sni string) (*queryResult, error) {
	client := dns.Client{
		Net:     "tcp-tls",
		Timeout: dnsTimeout,
		TLSConfig: &tls.Config{
			ServerName: sni,
			MinVersion: tls.VersionTLS12,
		},
	}
	if deadline, ok := ctx.Deadline(); ok {
		if d := time.Until(deadline); d > 0 && d < client.Timeout {
			client.Timeout = d
		}
	}
	r, rtt, err := client.ExchangeContext(ctx, m, net.JoinHostPort(ip, "853"))
	if err != nil {
		return nil, err
	}
	if r == nil {
		return nil, fmt.Errorf("nil response from %s", ip)
	}
	return &queryResult{
		Rcode: r.Rcode, Answer: r.Answer,
		AD: r.AuthenticatedData, Latency: rtt,
	}, nil
}

// exchangeDoH sends the wire-format DNS message as the "dns" query parameter
// per RFC 8484, using GET for cache friendliness. Response is parsed with
// miekg/dns so we share the answer/rcode handling.
func exchangeDoH(ctx context.Context, m *dns.Msg, endpoint string) (*queryResult, error) {
	// DoH wants Id=0 so intermediate HTTP caches can merge equivalent
	// queries. Setting it after composing keeps other fields intact.
	m.Id = 0
	packed, err := m.Pack()
	if err != nil {
		return nil, fmt.Errorf("packing message: %w", err)
	}

	u, err := url.Parse(endpoint)
	if err != nil {
		return nil, fmt.Errorf("invalid DoH endpoint %q: %w", endpoint, err)
	}
	q := u.Query()
	q.Set("dns", base64.RawURLEncoding.EncodeToString(packed))
	u.RawQuery = q.Encode()

	req, err := http.NewRequestWithContext(ctx, http.MethodGet, u.String(), nil)
	if err != nil {
		return nil, err
	}
	req.Header.Set("Accept", "application/dns-message")
	req.Header.Set("User-Agent", "happyDomain-checker-resolver-propagation/"+Version)

	start := time.Now()
	resp, err := dohClient.Do(req)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		return nil, fmt.Errorf("DoH HTTP %d", resp.StatusCode)
	}
	ct := resp.Header.Get("Content-Type")
	if !strings.HasPrefix(ct, "application/dns-message") {
		return nil, fmt.Errorf("DoH unexpected content-type %q", ct)
	}

	var buf bytes.Buffer
	if _, err := io.Copy(&buf, io.LimitReader(resp.Body, maxDoHResponseBytes)); err != nil {
		return nil, err
	}
	latency := time.Since(start)

	r := new(dns.Msg)
	if err := r.Unpack(buf.Bytes()); err != nil {
		return nil, fmt.Errorf("unpacking DoH response: %w", err)
	}
	return &queryResult{
		Rcode: r.Rcode, Answer: r.Answer,
		AD: r.AuthenticatedData, Latency: latency,
	}, nil
}

// canonicalRR returns an RR's RDATA in canonical, TTL-stripped form. We use
// miekg/dns's zone-file representation, cut off the header fields (name,
// class, TTL, type), and trim: what remains is the RDATA.
func canonicalRR(rr dns.RR) string {
	if rr == nil {
		return ""
	}
	s := rr.String()
	// Presentation: "owner TTL class type rdata...". Split on tabs/spaces,
	// skip the first four fields. dns.RR always emits the header so the
	// shape is stable.
	fields := strings.Fields(s)
	if len(fields) <= 4 {
		return ""
	}
	rdata := strings.Join(fields[4:], " ")
	// TXT rdata retains its original quoting; lowercase hostnames to avoid
	// case-only drift appearing as disagreement.
	return strings.ToLower(strings.TrimSpace(rdata))
}

// signatureFromRRs returns a deterministic signature built from the RDATA of
// the RRs that match owner+qtype. It also returns the sorted presentation
// list used for display and the smallest TTL observed.
func signatureFromRRs(rrs []dns.RR, owner string, qtype uint16) (sig string, records []string, minTTL uint32) {
	ownerL := strings.ToLower(dns.Fqdn(owner))
	for _, rr := range rrs {
		h := rr.Header()
		if h == nil {
			continue
		}
		if !strings.EqualFold(dns.Fqdn(h.Name), ownerL) {
			continue
		}
		if h.Rrtype != qtype {
			continue
		}
		if c := canonicalRR(rr); c != "" {
			records = append(records, c)
			if minTTL == 0 || h.Ttl < minTTL {
				minTTL = h.Ttl
			}
		}
	}
	sort.Strings(records)
	sig = strings.Join(records, "|")
	return sig, records, minTTL
}

// rcodeToString is a small wrapper around dns.RcodeToString that falls back
// to a numeric representation when the rcode is unknown.
func rcodeToString(c int) string {
	if s, ok := dns.RcodeToString[c]; ok {
		return s
	}
	return fmt.Sprintf("RCODE%d", c)
}