package checker

import (
	"context"
	"encoding/json"
	"fmt"
	"sort"
	"strings"
	"time"

	"github.com/miekg/dns"

	sdk "git.happydns.org/checker-sdk-go/checker"
)

// Collect runs the delegation testsuite and returns a *DelegationData
// populated with findings.
//
// The collector resolves the parent zone's authoritative servers, asks each
// of them for the delegation of the target FQDN, then turns around and
// queries every delegated server using ONLY the NS names + glue learned
// from the parent. The child zone is never used as a source of truth.
func (p *delegationProvider) Collect(ctx context.Context, opts sdk.CheckerOptions) (any, error) {
	svc, err := loadService(opts)
	if err != nil {
		return nil, err
	}

	parentZone, subdomain := loadNames(opts)
	if subdomain == "" {
		return nil, fmt.Errorf("missing 'subdomain' option")
	}
	if parentZone == "" {
		return nil, fmt.Errorf("missing 'domain_name' option")
	}

	delegatedFQDN := dns.Fqdn(strings.TrimSuffix(subdomain, ".") + "." + strings.TrimSuffix(parentZone, ".") + ".")

	data := &DelegationData{
		DelegatedFQDN: delegatedFQDN,
		ParentZone:    dns.Fqdn(parentZone),
	}

	requireDS := sdk.GetBoolOption(opts, "requireDS", false)
	requireTCP := sdk.GetBoolOption(opts, "requireTCP", true)
	minNS := sdk.GetIntOption(opts, "minNameServers", 2)
	allowGlueMismatch := sdk.GetBoolOption(opts, "allowGlueMismatch", false)

	// Declared NS / DS from the service.
	declaredNS := normalizeNSList(svc.NameServers)
	if len(declaredNS) < minNS {
		data.Findings = append(data.Findings, DelegationFinding{
			Code:     "delegation_too_few_ns",
			Severity: SeverityWarn,
			Message:  fmt.Sprintf("only %d name server(s) declared, RFC 1034 recommends at least %d", len(declaredNS), minNS),
		})
	}

	// Resolve parent's authoritative servers.
	_, parentServers, err := findParentZone(ctx, delegatedFQDN, parentZone)
	if err != nil {
		data.Findings = append(data.Findings, DelegationFinding{
			Code:     "delegation_no_parent_ns",
			Severity: SeverityCrit,
			Message:  err.Error(),
		})
		return data, nil
	}
	data.ParentNS = parentServers

	// Phase A: query every parent server.
	type parentView struct {
		server string
		ns     []string
		glue   map[string][]string
		ds     []*dns.DS
	}
	var views []parentView

	for _, ps := range parentServers {
		ns, glue, _, qerr := queryDelegation(ctx, ps, delegatedFQDN)
		if qerr != nil {
			data.Findings = append(data.Findings, DelegationFinding{
				Code:     "delegation_parent_query_failed",
				Severity: SeverityCrit,
				Message:  fmt.Sprintf("parent NS query failed: %v", qerr),
				Server:   ps,
			})
			continue
		}
		if len(ns) == 0 {
			data.Findings = append(data.Findings, DelegationFinding{
				Code:     "delegation_no_parent_ns",
				Severity: SeverityCrit,
				Message:  "parent returned an empty NS RRset",
				Server:   ps,
			})
			continue
		}

		// TCP reachability of the parent for the same query.
		if _, _, _, terr := queryDelegationTCP(ctx, ps, delegatedFQDN); terr != nil {
			sev := SeverityCrit
			if !requireTCP {
				sev = SeverityWarn
			}
			data.Findings = append(data.Findings, DelegationFinding{
				Code:     "delegation_parent_tcp_failed",
				Severity: sev,
				Message:  fmt.Sprintf("parent NS query over TCP failed: %v", terr),
				Server:   ps,
			})
		}

		// Compare NS to the declared list.
		missing, extra := diffStringSets(declaredNS, ns)
		if len(missing) > 0 || len(extra) > 0 {
			data.Findings = append(data.Findings, DelegationFinding{
				Code:     "delegation_ns_mismatch",
				Severity: SeverityCrit,
				Message:  fmt.Sprintf("NS RRset at parent does not match declared service: missing=%v extra=%v", missing, extra),
				Server:   ps,
			})
		}

		// Glue sanity: in-bailiwick NS must have glue, out-of-bailiwick NS must not.
		for _, n := range ns {
			inBailiwick := strings.HasSuffix(n, "."+delegatedFQDN) || strings.HasSuffix(n, delegatedFQDN)
			if inBailiwick {
				if len(glue[n]) == 0 {
					data.Findings = append(data.Findings, DelegationFinding{
						Code:     "delegation_missing_glue",
						Severity: SeverityCrit,
						Message:  fmt.Sprintf("in-bailiwick NS %s has no glue", n),
						Server:   ps,
					})
				}
			} else {
				if len(glue[n]) > 0 {
					data.Findings = append(data.Findings, DelegationFinding{
						Code:     "delegation_unnecessary_glue",
						Severity: SeverityWarn,
						Message:  fmt.Sprintf("out-of-bailiwick NS %s has glue records, which the parent should not return", n),
						Server:   ps,
					})
				}
			}
		}

		// DS at parent.
		ds, sigs, dserr := queryDS(ctx, ps, delegatedFQDN)
		if dserr != nil {
			data.Findings = append(data.Findings, DelegationFinding{
				Code:     "delegation_ds_query_failed",
				Severity: SeverityWarn,
				Message:  fmt.Sprintf("DS query failed: %v", dserr),
				Server:   ps,
			})
		} else {
			// Compare DS with declared service DS.
			declaredDS := svc.DS
			if len(declaredDS) > 0 || len(ds) > 0 {
				dsMissing, dsExtra := diffDS(declaredDS, ds)
				if len(dsMissing) > 0 || len(dsExtra) > 0 {
					sev := SeverityCrit
					if len(declaredDS) == 0 {
						// Service does not declare any DS but parent has some — warn only.
						sev = SeverityWarn
					}
					data.Findings = append(data.Findings, DelegationFinding{
						Code:     "delegation_ds_mismatch",
						Severity: sev,
						Message:  fmt.Sprintf("DS RRset at parent does not match declared service: missing=%d extra=%d", len(dsMissing), len(dsExtra)),
						Server:   ps,
					})
				}
			}

			if len(declaredDS) > 0 && len(ds) == 0 {
				sev := SeverityInfo
				if requireDS {
					sev = SeverityCrit
				}
				data.Findings = append(data.Findings, DelegationFinding{
					Code:     "delegation_ds_missing",
					Severity: sev,
					Message:  "service declares DS records but parent serves none",
					Server:   ps,
				})
			}

			// Validate DS RRSIG validity period if a signature is present.
			for _, sig := range sigs {
				if !sig.ValidityPeriod(time.Now()) {
					data.Findings = append(data.Findings, DelegationFinding{
						Code:     "delegation_ds_rrsig_invalid",
						Severity: SeverityCrit,
						Message:  fmt.Sprintf("DS RRSIG: %s", validityWindow(sig)),
						Server:   ps,
					})
				}
			}

			if len(ds) > 0 {
				dsTexts := make([]string, len(ds))
				for i, d := range ds {
					dsTexts[i] = d.String()
				}
				data.ParentDS = dsTexts
			}
		}

		views = append(views, parentView{server: ps, ns: ns, glue: glue, ds: ds})
	}

	if len(views) == 0 {
		// All parent servers failed; no point in continuing.
		return data, nil
	}

	// Pick the first successful parent view as the source of truth for
	// Phase B. We rely on the per-parent NS_mismatch findings already
	// emitted above to flag inconsistencies between parents.
	parent := views[0]
	data.AdvertisedNS = parent.ns
	data.AdvertisedGlue = parent.glue

	// Phase B: query each child name server using only parent-supplied data.
	data.ChildSerials = map[string]uint32{}
	for _, nsName := range parent.ns {
		addrs := parent.glue[nsName]
		if len(addrs) == 0 {
			// Out-of-bailiwick: resolve via the system resolver.
			resolved, rerr := resolveHost(ctx, nsName)
			if rerr != nil {
				data.Findings = append(data.Findings, DelegationFinding{
					Code:     "delegation_ns_unresolvable",
					Severity: SeverityCrit,
					Message:  fmt.Sprintf("cannot resolve NS %s: %v", nsName, rerr),
					Server:   nsName,
				})
				continue
			}
			addrs = resolved
		}

		var lastSerial uint32
		var sawAA bool
		for _, addr := range addrs {
			srv := hostPort(addr, "53")

			// UDP reachability + AA check.
			soa, aa, qerr := querySOA(ctx, "", srv, delegatedFQDN)
			if qerr != nil {
				data.Findings = append(data.Findings, DelegationFinding{
					Code:     "delegation_unreachable",
					Severity: SeverityCrit,
					Message:  fmt.Sprintf("UDP SOA query failed at %s (%s): %v", nsName, addr, qerr),
					Server:   srv,
				})
				continue
			}
			if !aa {
				data.Findings = append(data.Findings, DelegationFinding{
					Code:     "delegation_lame",
					Severity: SeverityCrit,
					Message:  fmt.Sprintf("server %s (%s) is not authoritative for %s", nsName, addr, delegatedFQDN),
					Server:   srv,
				})
				continue
			}
			sawAA = true
			if soa != nil {
				if lastSerial != 0 && lastSerial != soa.Serial {
					data.Findings = append(data.Findings, DelegationFinding{
						Code:     "delegation_soa_serial_drift",
						Severity: SeverityWarn,
						Message:  fmt.Sprintf("SOA serial drift on %s: %d vs %d", nsName, lastSerial, soa.Serial),
						Server:   srv,
					})
				}
				lastSerial = soa.Serial
				data.ChildSerials[srv] = soa.Serial
			}

			// TCP reachability.
			if _, _, terr := querySOA(ctx, "tcp", srv, delegatedFQDN); terr != nil {
				sev := SeverityCrit
				if !requireTCP {
					sev = SeverityWarn
				}
				data.Findings = append(data.Findings, DelegationFinding{
					Code:     "delegation_tcp_failed",
					Severity: sev,
					Message:  fmt.Sprintf("TCP SOA query failed at %s (%s): %v", nsName, addr, terr),
					Server:   srv,
				})
			}

			// NS RRset agreement with parent.
			childNS, nerr := queryNSAt(ctx, srv, delegatedFQDN)
			if nerr == nil {
				missing, extra := diffStringSets(parent.ns, childNS)
				if len(missing) > 0 || len(extra) > 0 {
					data.Findings = append(data.Findings, DelegationFinding{
						Code:     "delegation_ns_drift",
						Severity: SeverityWarn,
						Message:  fmt.Sprintf("child NS RRset differs from parent: missing=%v extra=%v", missing, extra),
						Server:   srv,
					})
				}
			}

			// In-bailiwick glue agreement.
			if isInBailiwick(nsName, delegatedFQDN) {
				childAddrs, _ := queryAddrsAt(ctx, srv, nsName)
				missing, _ := diffStringSets(parent.glue[nsName], childAddrs)
				if len(missing) > 0 {
					sev := SeverityCrit
					if allowGlueMismatch {
						sev = SeverityWarn
					}
					data.Findings = append(data.Findings, DelegationFinding{
						Code:     "delegation_glue_mismatch",
						Severity: sev,
						Message:  fmt.Sprintf("addresses served by child for %s differ from parent glue: missing=%v", nsName, missing),
						Server:   srv,
					})
				}
			}

			// DNSKEY hand-off, only if the parent has DS records.
			if len(parent.ds) > 0 {
				keys, kerr := queryDNSKEY(ctx, srv, delegatedFQDN)
				if kerr != nil {
					data.Findings = append(data.Findings, DelegationFinding{
						Code:     "delegation_dnskey_query_failed",
						Severity: SeverityWarn,
						Message:  fmt.Sprintf("DNSKEY query failed at %s: %v", nsName, kerr),
						Server:   srv,
					})
				} else if !dsMatchesAnyKey(parent.ds, keys) {
					data.Findings = append(data.Findings, DelegationFinding{
						Code:     "delegation_dnskey_no_match",
						Severity: SeverityCrit,
						Message:  fmt.Sprintf("none of the DNSKEY records served by %s match the DS published by the parent", nsName),
						Server:   srv,
					})
				}
			}
		}
		if !sawAA && len(addrs) > 0 {
			// At least record we tried.
			data.Findings = append(data.Findings, DelegationFinding{
				Code:     "delegation_no_authoritative_answer",
				Severity: SeverityCrit,
				Message:  fmt.Sprintf("no authoritative answer obtained from any address of %s", nsName),
				Server:   nsName,
			})
		}
	}

	return data, nil
}

// queryDelegationTCP is the TCP variant of queryDelegation. It is split out
// so the per-server findings keep their UDP/TCP roles distinct.
func queryDelegationTCP(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, "tcp", 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])
	}
	return
}

// loadService extracts the abstract.Delegation payload from the auto-filled
// "service" option. We parse it into our local minimal type so this checker
// does not have to import the full happyDomain server module.
func loadService(opts sdk.CheckerOptions) (*delegationService, error) {
	svc, ok := sdk.GetOption[serviceMessage](opts, "service")
	if !ok {
		return nil, fmt.Errorf("missing 'service' option")
	}
	if svc.Type != "" && svc.Type != "abstract.Delegation" {
		return nil, fmt.Errorf("service is %s, expected abstract.Delegation", svc.Type)
	}
	var d delegationService
	if err := json.Unmarshal(svc.Service, &d); err != nil {
		return nil, fmt.Errorf("decoding delegation service: %w", err)
	}
	return &d, nil
}

func loadNames(opts sdk.CheckerOptions) (parentZone, subdomain string) {
	if v, ok := sdk.GetOption[string](opts, "domain_name"); ok {
		parentZone = v
	}
	if v, ok := sdk.GetOption[string](opts, "subdomain"); ok {
		subdomain = v
	}
	return
}

// normalizeNSList lowercases and FQDN-normalizes a list of NS records.
func normalizeNSList(ns []*dns.NS) []string {
	out := make([]string, 0, len(ns))
	for _, n := range ns {
		if n == nil {
			continue
		}
		out = append(out, strings.ToLower(dns.Fqdn(n.Ns)))
	}
	sort.Strings(out)
	return out
}

// diffStringSets returns the elements of "want" missing from "got" and the
// elements of "got" not present in "want".
func diffStringSets(want, got []string) (missing, extra []string) {
	w := map[string]bool{}
	for _, v := range want {
		w[strings.ToLower(strings.TrimSuffix(v, "."))] = true
	}
	g := map[string]bool{}
	for _, v := range got {
		g[strings.ToLower(strings.TrimSuffix(v, "."))] = true
	}
	for k := range w {
		if !g[k] {
			missing = append(missing, k)
		}
	}
	for k := range g {
		if !w[k] {
			extra = append(extra, k)
		}
	}
	sort.Strings(missing)
	sort.Strings(extra)
	return
}

// diffDS returns the DS records present in "want" but missing from "got"
// and vice-versa.
func diffDS(want, got []*dns.DS) (missing, extra []*dns.DS) {
	for _, w := range want {
		found := false
		for _, g := range got {
			if dsEqual(w, g) {
				found = true
				break
			}
		}
		if !found {
			missing = append(missing, w)
		}
	}
	for _, g := range got {
		found := false
		for _, w := range want {
			if dsEqual(w, g) {
				found = true
				break
			}
		}
		if !found {
			extra = append(extra, g)
		}
	}
	return
}

// isInBailiwick reports whether host sits inside zone.
func isInBailiwick(host, zone string) bool {
	host = strings.ToLower(dns.Fqdn(host))
	zone = strings.ToLower(dns.Fqdn(zone))
	return host == zone || strings.HasSuffix(host, "."+zone)
}

// dsMatchesAnyKey reports whether at least one of the DNSKEY records hashes
// to one of the DS records.
func dsMatchesAnyKey(ds []*dns.DS, keys []*dns.DNSKEY) bool {
	for _, k := range keys {
		for _, d := range ds {
			expected := k.ToDS(d.DigestType)
			if expected != nil && dsEqual(expected, d) {
				return true
			}
		}
	}
	return false
}