package nscheck

import (
	"context"
	"fmt"
	"net/netip"
	"sort"
	"strings"
	"sync"

	"github.com/miekg/dns"
)

// Finding is one conclusion of a delegation check. Level is "ok",
// "info", "warning" or "error".
type Finding struct {
	Level   string `json:"level"`
	Message string `json:"message"`
}

// ParentAnswer is what one of the parent zone's servers says.
type ParentAnswer struct {
	Name    string `json:"name"`
	Address string `json:"address,omitempty"`
	// Status: "delegated", "not_delegated" (NXDOMAIN, or the name exists
	// without a delegation), "unreachable" or "error".
	Status string   `json:"status"`
	NS     []string `json:"ns,omitempty"`
	// Glue is the addresses the parent hands out for in-zone name
	// servers.
	Glue  map[string][]string `json:"glue,omitempty"`
	Error string              `json:"error,omitempty"`
}

// GlueCheck compares the parent's glue for an in-zone name server with
// Jarl's addresses for it.
type GlueCheck struct {
	Name   string   `json:"name"`
	Parent []string `json:"parent"`
	Jarl   []string `json:"jarl"`
	Match  bool     `json:"match"`
}

// DelegationReport is the result of Delegation.
type DelegationReport struct {
	Zone     string `json:"zone"`
	Parent   string `json:"parent,omitempty"`
	Resolver string `json:"resolver"`
	// Delegated: at least one parent server delegates the zone.
	Delegated bool           `json:"delegated"`
	ParentNS  []string       `json:"parent_ns"`
	JarlNS    []string       `json:"jarl_ns"`
	Parents   []ParentAnswer `json:"parent_servers"`
	Glue      []GlueCheck    `json:"glue"`
	// Servers are the delegated name servers' SOA answers, as resolvers
	// on the internet would reach them.
	Servers  []Result  `json:"servers"`
	Findings []Finding `json:"findings"`
}

// Delegation checks how the parent zone delegates zone: it finds the
// parent through the resolver, asks each of the parent's servers for the
// zone's NS records and glue, compares those with jarlNS and jarlAddrs
// (Jarl's NS records and its addresses for in-zone name servers), and
// asks every delegated server for the zone's serial (compared with ref).
func (c *Checker) Delegation(ctx context.Context, zone string, jarlNS []string, jarlAddrs map[string][]string, ref *uint32) *DelegationReport {
	zone = dns.CanonicalName(zone)
	rep := &DelegationReport{Zone: zone, JarlNS: canonAll(jarlNS), ParentNS: []string{}, Parents: []ParentAnswer{}, Glue: []GlueCheck{}, Servers: []Result{}, Findings: []Finding{}}
	sort.Strings(rep.JarlNS)
	add := func(level, format string, args ...any) {
		for i, a := range args {
			switch v := a.(type) {
			case string:
				args[i] = show(v)
			case []string:
				args[i] = showAll(v)
			}
		}
		rep.Findings = append(rep.Findings, Finding{level, fmt.Sprintf(format, args...)})
	}
	res, err := c.resolver()
	if err != nil {
		add("error", "No resolver to find the parent zone with: %v", err)
		return rep
	}
	rep.Resolver = res
	if zone == "." {
		add("info", "The root zone has no parent to delegate it.")
		return rep
	}

	// The parent zone: the zone that holds the name just above ours.
	parent, err := c.enclosingZone(ctx, res, parentName(zone))
	if err != nil {
		add("error", "Couldn't find the parent zone of %s through %s: %v", zone, res, err)
		return rep
	}
	rep.Parent = parent
	pns, err := c.nsNames(ctx, res, parent)
	if err != nil || len(pns) == 0 {
		add("error", "Couldn't find the name servers of the parent zone %s through %s: %v", parent, res, errOr(err, "none listed"))
		return rep
	}

	// Ask every parent server, in parallel.
	var wg sync.WaitGroup
	answers := make([]ParentAnswer, len(pns))
	for i, name := range pns {
		wg.Add(1)
		go func(i int, name string) {
			defer wg.Done()
			answers[i] = c.askParent(ctx, zone, name)
		}(i, name)
	}
	wg.Wait()
	rep.Parents = answers

	parentNS := map[string]bool{}
	glue := map[string]map[string]bool{}
	reached, delegated, notDelegated := 0, 0, 0
	nsSets := map[string]bool{}
	for _, a := range answers {
		switch a.Status {
		case "delegated":
			reached++
			delegated++
			nsSets[strings.Join(a.NS, " ")] = true
			for _, n := range a.NS {
				parentNS[n] = true
			}
			for n, addrs := range a.Glue {
				if glue[n] == nil {
					glue[n] = map[string]bool{}
				}
				for _, ad := range addrs {
					glue[n][ad] = true
				}
			}
		case "not_delegated":
			reached++
			notDelegated++
		}
	}
	if reached == 0 {
		add("error", "None of the parent zone's servers (%s) answered; this server may not be able to reach them.", pns)
		return rep
	}
	if delegated == 0 {
		rep.Delegated = false
		where := "the parent zone " + show(parent)
		if parent == "." {
			where = "the root zone"
		}
		add("info", "%s isn't delegated from %s, so the internet can't see it. That's expected for a zone used only on your own network. To publish it, register the domain and give your registrar these name servers: %s.", zone, where, rep.JarlNS)
		return rep
	}
	rep.Delegated = true
	if notDelegated > 0 {
		add("warning", "%d of the parent's %d servers don't know the delegation yet (a recent change at the registrar can take a while to reach them all).", notDelegated, reached)
	}
	if len(nsSets) > 1 {
		add("warning", "The parent's servers give different name server lists for %s.", zone)
	}
	rep.ParentNS = keys(parentNS)

	// The NS lists.
	jarlSet := map[string]bool{}
	for _, n := range rep.JarlNS {
		jarlSet[n] = true
	}
	var missing, extra []string
	for _, n := range rep.JarlNS {
		if !parentNS[n] {
			missing = append(missing, n)
		}
	}
	for _, n := range rep.ParentNS {
		if !jarlSet[n] {
			extra = append(extra, n)
		}
	}
	if len(missing) == 0 && len(extra) == 0 {
		add("ok", "The parent delegates %s to the same name servers as its NS records: %s.", zone, rep.ParentNS)
	}
	if len(extra) > 0 {
		add("warning", "The parent delegates to %s, which the zone's NS records don't list. Resolvers ask them too; remove them at the registrar if they no longer serve the zone, or add NS records for them.", extra)
	}
	if len(missing) > 0 {
		add("warning", "The zone's NS records list %s, which the parent doesn't delegate to. Resolvers that start from the parent won't use them; add them at the registrar.", missing)
	}

	// Glue for name servers inside the zone.
	for _, n := range rep.ParentNS {
		if !dns.IsSubDomain(zone, n) {
			continue
		}
		g := GlueCheck{Name: n, Parent: keys(glue[n]), Jarl: canonAddrs(jarlAddrs[n])}
		g.Match = len(g.Parent) > 0 && strings.Join(g.Parent, " ") == strings.Join(g.Jarl, " ")
		rep.Glue = append(rep.Glue, g)
		switch {
		case len(g.Parent) == 0:
			add("error", "The parent has no glue address for %s, which is inside the zone, so resolvers can't find it. Add its address at the registrar (%s).", n, strings.Join(g.Jarl, ", "))
		case len(g.Jarl) == 0:
			add("warning", "The parent's glue for %s is %s, but the zone has no address record for it.", n, strings.Join(g.Parent, ", "))
		case !g.Match:
			add("warning", "The parent's glue for %s is %s, but the zone says %s. Update the glue at the registrar.", n, strings.Join(g.Parent, ", "), strings.Join(g.Jarl, ", "))
		}
	}

	// Every delegated server, reached the way resolvers reach it: by its
	// glue, or by looking its name up.
	var targets []Target
	for _, n := range rep.ParentNS {
		t := Target{Name: n, Source: "delegation"}
		if g := keys(glue[n]); len(g) > 0 {
			t.Addresses = g
		}
		targets = append(targets, t)
	}
	// Here every address is asked, this server's too: what matters is
	// what it answers from the outside.
	pub := *c
	pub.IsLocal = func(netip.Addr) bool { return false }
	rep.Servers = pub.Serials(ctx, zone, ref, targets)
	// One finding per server and outcome, however many addresses it has.
	type outcome struct{ name, status, detail string }
	var order []outcome
	addrs := map[outcome][]string{}
	bad := 0
	for _, r := range rep.Servers {
		var o outcome
		switch r.Status {
		case "in_sync", "answered", "no_route":
			continue
		case "behind", "ahead":
			o = outcome{r.Name, r.Status, fmt.Sprintf("serves serial %d, not %d", *r.Serial, *ref)}
		case "not_authoritative":
			o = outcome{r.Name, r.Status, ""}
		default:
			o = outcome{r.Name, r.Status, describe(r)}
		}
		bad++
		if addrs[o] == nil {
			order = append(order, o)
		}
		addrs[o] = append(addrs[o], nonEmpty(r.Address, "no address"))
	}
	for _, o := range order {
		at := strings.Join(addrs[o], ", ")
		if len(addrs[o]) > 2 {
			at = fmt.Sprintf("all %d addresses", len(addrs[o]))
		}
		if o.status == "not_authoritative" {
			add("error", "%s (%s) is delegated the zone but doesn't serve it (a lame delegation).", o.name, at)
		} else {
			add("warning", "%s (%s): %s.", o.name, at, o.detail)
		}
	}
	if bad == 0 && len(rep.Servers) > 0 {
		add("ok", "Every delegated name server answers for the zone with the current serial.")
	}
	return rep
}

// askParent asks one of the parent's servers (by name) for zone's NS.
func (c *Checker) askParent(ctx context.Context, zone, server string) ParentAnswer {
	pa := ParentAnswer{Name: server}
	addrs, err := c.lookup(ctx, server)
	if err != nil || len(addrs) == 0 {
		pa.Status, pa.Error = "error", "looking up its address: "+errOr(err, "no address").Error()
		return pa
	}
	m := new(dns.Msg)
	m.SetQuestion(zone, dns.TypeNS)
	m.RecursionDesired = false
	// IPv4 first: try each address until one answers.
	sort.SliceStable(addrs, func(i, j int) bool { return !strings.Contains(addrs[i], ":") && strings.Contains(addrs[j], ":") })
	var resp *dns.Msg
	for _, a := range addrs {
		pa.Address = a
		resp, _, err = c.exchange(ctx, m.Copy(), a, nil)
		if err == nil {
			break
		}
	}
	if err != nil {
		pa.Status, pa.Error = "unreachable", err.Error()
		return pa
	}
	if resp.Rcode == dns.RcodeNameError {
		pa.Status = "not_delegated"
		return pa
	}
	if resp.Rcode != dns.RcodeSuccess {
		pa.Status, pa.Error = "error", "answered "+dns.RcodeToString[resp.Rcode]
		return pa
	}
	// A referral puts the NS set in the authority section; a server that
	// also serves the child zone answers it directly.
	ns := map[string]bool{}
	for _, sec := range [][]dns.RR{resp.Ns, resp.Answer} {
		for _, rr := range sec {
			if n, ok := rr.(*dns.NS); ok && dns.CanonicalName(n.Hdr.Name) == zone {
				ns[dns.CanonicalName(n.Ns)] = true
			}
		}
	}
	if len(ns) == 0 {
		pa.Status = "not_delegated"
		return pa
	}
	pa.Status, pa.NS = "delegated", keys(ns)
	for _, rr := range resp.Extra {
		var name, addr string
		switch v := rr.(type) {
		case *dns.A:
			name, addr = v.Hdr.Name, v.A.String()
		case *dns.AAAA:
			name, addr = v.Hdr.Name, v.AAAA.String()
		default:
			continue
		}
		name = dns.CanonicalName(name)
		if ns[name] {
			if pa.Glue == nil {
				pa.Glue = map[string][]string{}
			}
			pa.Glue[name] = append(pa.Glue[name], addr)
		}
	}
	for n := range pa.Glue {
		pa.Glue[n] = canonAddrs(pa.Glue[n])
	}
	return pa
}

// enclosingZone finds the apex of the zone that holds name, from the
// owner of the SOA record the resolver returns (in the answer for an apex,
// in the authority section otherwise).
func (c *Checker) enclosingZone(ctx context.Context, res, name string) (string, error) {
	m := new(dns.Msg)
	m.SetQuestion(name, dns.TypeSOA)
	resp, _, err := c.exchangeHostPort(ctx, m, res, nil)
	if err != nil {
		return "", err
	}
	for _, sec := range [][]dns.RR{resp.Answer, resp.Ns} {
		for _, rr := range sec {
			if s, ok := rr.(*dns.SOA); ok {
				return dns.CanonicalName(s.Hdr.Name), nil
			}
		}
	}
	return "", fmt.Errorf("no SOA record in the answer for %s (%s)", name, dns.RcodeToString[resp.Rcode])
}

// nsNames asks the resolver for a zone's NS names.
func (c *Checker) nsNames(ctx context.Context, res, zone string) ([]string, error) {
	m := new(dns.Msg)
	m.SetQuestion(zone, dns.TypeNS)
	resp, _, err := c.exchangeHostPort(ctx, m, res, nil)
	if err != nil {
		return nil, err
	}
	set := map[string]bool{}
	for _, rr := range resp.Answer {
		if n, ok := rr.(*dns.NS); ok {
			set[dns.CanonicalName(n.Ns)] = true
		}
	}
	return keys(set), nil
}

// show is a name as people write it, without the trailing dot.
func show(name string) string {
	if name == "." || !strings.HasSuffix(name, ".") {
		return name
	}
	return strings.TrimSuffix(name, ".")
}

func showAll(names []string) string {
	out := make([]string, len(names))
	for i, n := range names {
		out[i] = show(n)
	}
	return strings.Join(out, ", ")
}

func parentName(zone string) string {
	labels := dns.SplitDomainName(zone)
	if len(labels) <= 1 {
		return "."
	}
	return dns.Fqdn(strings.Join(labels[1:], "."))
}

func describe(r Result) string {
	switch r.Status {
	case "refused":
		return "refuses queries for the zone"
	case "unreachable":
		return "no answer (" + r.Error + ")"
	case "no_address":
		return r.Error
	}
	return nonEmpty(r.Error, r.Status)
}

func keys(m map[string]bool) []string {
	out := make([]string, 0, len(m))
	for k := range m {
		out = append(out, k)
	}
	sort.Strings(out)
	return out
}

func canonAll(names []string) []string {
	out := make([]string, 0, len(names))
	for _, n := range names {
		out = append(out, dns.CanonicalName(n))
	}
	return out
}

func canonAddrs(addrs []string) []string {
	set := map[string]bool{}
	for _, a := range addrs {
		if ip, err := netip.ParseAddr(a); err == nil {
			set[ip.String()] = true
		}
	}
	return keys(set)
}

func errOr(err error, msg string) error {
	if err != nil {
		return err
	}
	return fmt.Errorf("%s", msg)
}

func nonEmpty(s, def string) string {
	if s != "" {
		return s
	}
	return def
}