package nscheck

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

	"github.com/miekg/dns"

	"jarl/internal/dnssec"
)

// 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"`
	// DS is the zone's DS records at this server ("" in DSError when it
	// couldn't be asked).
	DS      []dnssec.DS `json:"ds,omitempty"`
	DSError string      `json:"ds_error,omitempty"`
	Error   string      `json:"error,omitempty"`
}

// Signing is how Jarl signs a zone, for comparing with the parent's DS
// records.
type Signing struct {
	// Policy is the zone's dnssec-policy ("" when unsigned, "insecure"
	// while signing is being removed).
	Policy string
	// DS are the DS records of the zone's key-signing keys that named
	// still publishes.
	DS []dnssec.DS
}

// ParentDSReport is what the parent zone's servers say about the zone's
// DS records.
type ParentDSReport struct {
	Zone     string `json:"zone"`
	Parent   string `json:"parent,omitempty"`
	Resolver string `json:"resolver,omitempty"`
	// Delegated: at least one parent server delegates the zone. Reached
	// is how many answered.
	Delegated bool           `json:"delegated"`
	Reached   int            `json:"reached"`
	Servers   []ParentAnswer `json:"parent_servers"`
	// DS is every DS record any parent server has.
	DS []dnssec.DS `json:"ds"`
	// Consistent: every server that answered has the same DS records.
	Consistent bool   `json:"consistent"`
	Error      string `json:"error,omitempty"`
}

// AllHaveDS reports whether every parent server that answered has a DS
// record for key tag (with digest, unless "").
func (r *ParentDSReport) AllHaveDS(tag uint16, digest string) bool {
	n := 0
	for _, a := range r.Servers {
		if a.Status != "delegated" || a.DSError != "" {
			continue
		}
		n++
		found := false
		for _, d := range a.DS {
			if d.KeyTag == tag && (digest == "" || strings.EqualFold(d.Digest, digest)) {
				found = true
			}
		}
		if !found {
			return false
		}
	}
	return n > 0
}

// HasDS reports whether a parent server has a DS record for key tag.
func (r *ParentDSReport) HasDS(tag uint16, digest string) bool {
	for _, d := range r.DS {
		if d.KeyTag == tag && (digest == "" || strings.EqualFold(d.Digest, digest)) {
			return true
		}
	}
	return false
}

// 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"`
	// ParentDS are the zone's DS records at the parent (DNSSEC).
	ParentDS []dnssec.DS `json:"parent_ds"`
	// 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, glue and DS records, compares those with jarlNS and
// jarlAddrs (Jarl's NS records and its addresses for in-zone name
// servers) and with sig (how Jarl signs the zone; nil skips the DNSSEC
// part), 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, sig *Signing) *DelegationReport {
	zone = dns.CanonicalName(zone)
	rep := &DelegationReport{Zone: zone, JarlNS: canonAll(jarlNS), ParentNS: []string{}, Parents: []ParentAnswer{}, Glue: []GlueCheck{}, ParentDS: []dnssec.DS{}, 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...)})
	}
	pr := c.ParentDS(ctx, zone)
	rep.Resolver, rep.Parent, rep.Parents = pr.Resolver, pr.Parent, pr.Servers
	if pr.Error != "" {
		level := "error"
		if zone == "." {
			level = "info"
		}
		rep.Findings = append(rep.Findings, Finding{level, pr.Error})
		return rep
	}
	parent, answers := pr.Parent, pr.Servers

	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 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.ParentDS = pr.DS
	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, ", "))
		}
	}

	if sig != nil {
		for _, f := range dsFindings(zone, pr, sig) {
			rep.Findings = append(rep.Findings, f)
		}
	}

	// 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
}

// ParentDS finds the parent of zone through the resolver and asks each of
// its servers for the zone's delegation and DS records.
func (c *Checker) ParentDS(ctx context.Context, zone string) *ParentDSReport {
	zone = dns.CanonicalName(zone)
	rep := &ParentDSReport{Zone: zone, Servers: []ParentAnswer{}, DS: []dnssec.DS{}}
	res, err := c.resolver()
	if err != nil {
		rep.Error = "No resolver to find the parent zone with: " + err.Error()
		return rep
	}
	rep.Resolver = res
	if zone == "." {
		rep.Error = "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 {
		rep.Error = fmt.Sprintf("Couldn't find the parent zone of %s through %s: %v", show(zone), res, err)
		return rep
	}
	rep.Parent = parent
	pns, err := c.nsNames(ctx, res, parent)
	if err != nil || len(pns) == 0 {
		rep.Error = fmt.Sprintf("Couldn't find the name servers of the parent zone %s through %s: %v", show(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.Servers = answers
	all := map[string]dnssec.DS{}
	sets := map[string]bool{}
	for _, a := range answers {
		if a.Status != "delegated" && a.Status != "not_delegated" {
			continue
		}
		rep.Reached++
		if a.Status == "delegated" {
			rep.Delegated = true
		}
		if a.DSError != "" {
			continue
		}
		var set []string
		for _, d := range a.DS {
			k := dsKey(d)
			all[k] = d
			set = append(set, k)
		}
		sort.Strings(set)
		sets[strings.Join(set, " ")] = true
	}
	if rep.Reached == 0 {
		names := make([]string, len(pns))
		for i, n := range pns {
			names[i] = show(n)
		}
		rep.Error = fmt.Sprintf("None of the parent zone's servers (%s) answered; this server may not be able to reach them.", strings.Join(names, ", "))
		return rep
	}
	for _, k := range keys(boolSet(all)) {
		rep.DS = append(rep.DS, all[k])
	}
	rep.Consistent = len(sets) <= 1
	return rep
}

func dsKey(d dnssec.DS) string {
	return fmt.Sprintf("%05d %d %d %s", d.KeyTag, d.Algorithm, d.DigestType, strings.ToUpper(d.Digest))
}

func boolSet(m map[string]dnssec.DS) map[string]bool {
	out := map[string]bool{}
	for k := range m {
		out[k] = true
	}
	return out
}

// dsFindings compares the parent's DS records with how Jarl signs the
// zone.
func dsFindings(zone string, pr *ParentDSReport, sig *Signing) []Finding {
	var out []Finding
	add := func(level, format string, args ...any) {
		out = append(out, Finding{level, fmt.Sprintf(format, args...)})
	}
	z := show(zone)
	var tags []string
	for _, d := range pr.DS {
		tags = append(tags, fmt.Sprint(d.KeyTag))
	}
	if !pr.Consistent {
		add("warning", "The parent's servers have different DS records for %s (a change at the registrar may still be on its way to them all).", z)
	}
	matches := func(d dnssec.DS) bool {
		for _, k := range sig.DS {
			if k.KeyTag == d.KeyTag && k.Algorithm == d.Algorithm && k.DigestType == d.DigestType && strings.EqualFold(k.Digest, d.Digest) {
				return true
			}
		}
		return false
	}
	var good, stale []string
	for _, d := range pr.DS {
		if matches(d) {
			good = append(good, fmt.Sprint(d.KeyTag))
		} else {
			stale = append(stale, fmt.Sprint(d.KeyTag))
		}
	}
	switch {
	case len(pr.DS) == 0 && sig.Policy == "":
		// Unsigned and no DS: nothing to say about DNSSEC.
	case len(pr.DS) == 0 && sig.Policy == "insecure":
		add("ok", "The parent has no DS record for %s any more, so turning DNSSEC off can go ahead.", z)
	case len(pr.DS) == 0:
		add("info", "%s is signed, but the parent has no DS record for it, so resolvers treat it as unsigned. Add the DS record at your registrar when the zone's DNSSEC page says it's ready.", z)
	case sig.Policy == "":
		add("error", "The parent has a DS record for %s (key %s), but the zone isn't signed: validating resolvers reject all of its answers. Remove the DS record at your registrar now.", z, strings.Join(tags, ", "))
	case len(good) == 0:
		add("error", "The parent's DS record for %s (key %s) matches none of the zone's keys: validating resolvers reject all of its answers. Replace it at your registrar with the zone's DS record.", z, strings.Join(tags, ", "))
	default:
		if sig.Policy == "insecure" {
			add("warning", "The parent still has the DS record for %s (key %s). DNSSEC is being turned off: remove it at your registrar, then run the DS check.", z, strings.Join(good, ", "))
		} else {
			add("ok", "The parent's DS record matches the zone's key %s: validating resolvers check its answers.", strings.Join(good, ", "))
		}
		if len(stale) > 0 {
			add("warning", "The parent also has a DS record for key %s, which the zone no longer has. Remove it at your registrar.", strings.Join(stale, ", "))
		}
	}
	return out
}

// askParent asks one of the parent's servers (by name) for zone's NS and
// DS records.
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])
	}
	// The parent itself answers DS queries for the zone.
	m = new(dns.Msg)
	m.SetQuestion(zone, dns.TypeDS)
	m.RecursionDesired = false
	resp, _, err = c.exchange(ctx, m, pa.Address, nil)
	switch {
	case err != nil:
		pa.DSError = err.Error()
	case resp.Rcode != dns.RcodeSuccess:
		pa.DSError = "answered " + dns.RcodeToString[resp.Rcode]
	default:
		for _, rr := range resp.Answer {
			if d, ok := rr.(*dns.DS); ok && dns.CanonicalName(d.Hdr.Name) == zone {
				pa.DS = append(pa.DS, dnssec.DS{KeyTag: d.KeyTag, Algorithm: d.Algorithm, DigestType: d.DigestType, Digest: strings.ToUpper(d.Digest), Record: d.String()})
			}
		}
	}
	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
}