package checker

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

	"github.com/miekg/dns"

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

// Collect runs the delegation probe and returns a *DelegationData populated
// with raw facts only. All judgment (severity, option-driven thresholds,
// pass/fail) is deferred to the rules in rule.go.
//
// 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),
		DeclaredNS:    normalizeNSList(svc.NameServers),
	}
	for _, d := range svc.DS {
		if d == nil {
			continue
		}
		data.DeclaredDS = append(data.DeclaredDS, NewDSRecord(d))
	}

	// Resolve parent's authoritative servers.
	_, parentServers, err := findParentZone(ctx, delegatedFQDN, parentZone)
	if err != nil {
		data.ParentDiscoveryError = err.Error()
		return data, nil
	}
	data.ParentNS = parentServers

	// Phase A: query every parent server. Record raw outcomes only.
	for _, ps := range parentServers {
		view := ParentView{Server: ps}

		ns, glue, _, qerr := queryDelegation(ctx, ps, delegatedFQDN)
		if qerr != nil {
			view.UDPNSError = qerr.Error()
		} else {
			view.NS = ns
			view.Glue = glue
		}

		if _, _, _, terr := queryDelegationTCP(ctx, ps, delegatedFQDN); terr != nil {
			view.TCPNSError = terr.Error()
		}

		dsRRs, sigs, dserr := queryDS(ctx, ps, delegatedFQDN)
		if dserr != nil {
			view.DSQueryError = dserr.Error()
		} else {
			for _, d := range dsRRs {
				view.DS = append(view.DS, NewDSRecord(d))
			}
			for _, sig := range sigs {
				view.DSRRSIGs = append(view.DSRRSIGs, DSRRSIGObservation{
					Inception:  sig.Inception,
					Expiration: sig.Expiration,
				})
			}
		}

		data.ParentViews = append(data.ParentViews, view)
	}

	// Pick the first view that actually returned an NS RRset as the
	// source of truth for Phase B. If none succeeded, skip Phase B; the
	// rules will flag the absence of child data.
	var primary *ParentView
	for i := range data.ParentViews {
		if data.ParentViews[i].UDPNSError == "" && len(data.ParentViews[i].NS) > 0 {
			primary = &data.ParentViews[i]
			break
		}
	}
	if primary == nil {
		return data, nil
	}

	// Phase B: query each child name server using only parent-supplied data.
	for _, nsName := range primary.NS {
		child := ChildNSView{NSName: nsName}
		addrs := primary.Glue[nsName]
		if len(addrs) == 0 {
			// Out-of-bailiwick: resolve via the system resolver.
			resolved, rerr := resolveHost(ctx, nsName)
			if rerr != nil {
				child.ResolveError = rerr.Error()
				data.Children = append(data.Children, child)
				continue
			}
			addrs = resolved
		}

		for _, addr := range addrs {
			srv := hostPort(addr, "53")
			av := ChildAddressView{Address: addr, Server: srv}

			soa, aa, qerr := querySOA(ctx, "", srv, delegatedFQDN)
			if qerr != nil {
				av.UDPError = qerr.Error()
				av.Authoritative = aa
				child.Addresses = append(child.Addresses, av)
				continue
			}
			av.Authoritative = aa
			if soa != nil {
				av.SOASerial = soa.Serial
				av.SOASerialKnown = true
			}

			if _, _, terr := querySOA(ctx, "tcp", srv, delegatedFQDN); terr != nil {
				av.TCPError = terr.Error()
			}

			childNS, nerr := queryNSAt(ctx, srv, delegatedFQDN)
			if nerr != nil {
				av.ChildNSError = nerr.Error()
			} else {
				av.ChildNS = childNS
			}

			if isInBailiwick(nsName, delegatedFQDN) {
				addrsAt, _ := queryAddrsAt(ctx, srv, nsName)
				av.ChildGlueAddrs = addrsAt
			}

			// Only bother probing DNSKEY when the parent has at least one
			// DS to match against. The rule confirms this precondition.
			parentHasDS := false
			for _, pv := range data.ParentViews {
				if len(pv.DS) > 0 {
					parentHasDS = true
					break
				}
			}
			if parentHasDS {
				keys, kerr := queryDNSKEY(ctx, srv, delegatedFQDN)
				if kerr != nil {
					av.DNSKEYError = kerr.Error()
				} else {
					for _, k := range keys {
						av.DNSKEYs = append(av.DNSKEYs, NewDNSKEYRecord(k))
					}
				}
			}

			child.Addresses = append(child.Addresses, av)
		}

		data.Children = append(data.Children, child)
	}

	return data, nil
}

// queryDelegationTCP is the TCP variant of queryDelegation. It is split out
// so the per-server observations 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
}