package checker

import (
	"encoding/json"
	"sort"
	"strings"

	"github.com/miekg/dns"

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

// rawZone is the minimal slice of happyDomain's Zone JSON we consume to
// derive the RR types actually present at each owner. It mirrors the
// shape used by sibling checkers (see checker-legacy-records).
type rawZone struct {
	DomainName string                  `json:"domain_name,omitempty"`
	Services   map[string][]rawService `json:"services"`
}

type rawService struct {
	Type    string          `json:"_svctype"`
	Domain  string          `json:"_domain"`
	Service json.RawMessage `json:"Service"`
}

// fallbackQTypes is the legacy default applied when no zone is available
// and the user did not set recordTypes explicitly.
var fallbackQTypes = []uint16{
	dns.TypeSOA, dns.TypeNS, dns.TypeA, dns.TypeAAAA,
	dns.TypeMX, dns.TypeTXT, dns.TypeCAA,
}

// resolveQTypes returns the RR types to probe at each owner name plus the
// union across all owners (for reporting/metrics).
//
// Precedence:
//  1. Explicit "recordTypes" option → apply that list to every owner.
//  2. Auto-filled "zone" option → derive per-owner types from the zone's
//     services. The apex always carries SOA+NS even if the zone payload
//     omits them. Owners with no derivable types fall back to A,AAAA so
//     the probe still surfaces NXDOMAIN drift for user-requested
//     subdomains that are not present in the zone.
//  3. Neither → use the legacy default at every owner.
func resolveQTypes(opts sdk.CheckerOptions, recordTypesOpt, apex string, names []string) (map[string][]uint16, []uint16, error) {
	if recordTypesOpt != "" {
		qts := parseQTypes(recordTypesOpt)
		if len(qts) == 0 {
			return nil, nil, &invalidTypesError{raw: recordTypesOpt}
		}
		return uniformOwnerQTypes(names, qts), qts, nil
	}

	zone, _ := readWorkingZone(opts)
	if zone == nil {
		return uniformOwnerQTypes(names, fallbackQTypes), append([]uint16(nil), fallbackQTypes...), nil
	}

	owner := map[string]map[uint16]bool{}
	for _, n := range names {
		owner[n] = map[uint16]bool{}
	}

	for sub, services := range zone.Services {
		full := joinSubdomain(sub, apex)
		set, ok := owner[full]
		if !ok {
			continue
		}
		for _, svc := range services {
			for _, qt := range typesFromService(svc) {
				set[qt] = true
			}
		}
	}

	// SOA + NS at apex are foundational; the rules depend on them.
	apexLower := strings.ToLower(dns.Fqdn(apex))
	if set, ok := owner[apexLower]; ok {
		set[dns.TypeSOA] = true
		set[dns.TypeNS] = true
	}

	out := make(map[string][]uint16, len(names))
	unionSet := map[uint16]bool{}
	for _, n := range names {
		set := owner[n]
		if len(set) == 0 {
			// Owner present in the probe list but unknown to the zone:
			// keep a minimal probe so a missing-record finding can fire.
			set = map[uint16]bool{dns.TypeA: true, dns.TypeAAAA: true}
		}
		qts := sortedTypes(set)
		out[n] = qts
		for _, qt := range qts {
			unionSet[qt] = true
		}
	}
	return out, sortedTypes(unionSet), nil
}

func uniformOwnerQTypes(names []string, qts []uint16) map[string][]uint16 {
	out := make(map[string][]uint16, len(names))
	for _, n := range names {
		out[n] = qts
	}
	return out
}

func sortedTypes(set map[uint16]bool) []uint16 {
	out := make([]uint16, 0, len(set))
	for q := range set {
		out = append(out, q)
	}
	sort.Slice(out, func(i, j int) bool { return out[i] < out[j] })
	return out
}

// readWorkingZone parses the "zone" auto-fill option. The host may pass
// the value either as a native struct (in-process plugin) or as a JSON
// object (HTTP path); we round-trip through JSON in both cases for a
// single decoding path. A missing zone is not an error — standalone /
// HTTP callers may simply not provide one.
func readWorkingZone(opts sdk.CheckerOptions) (*rawZone, error) {
	v, ok := opts["zone"]
	if !ok || v == nil {
		return nil, nil
	}
	raw, err := json.Marshal(v)
	if err != nil {
		return nil, err
	}
	z := &rawZone{}
	if err := json.Unmarshal(raw, z); err != nil {
		return nil, err
	}
	return z, nil
}

// typesFromService extracts every RR type referenced by a service body.
// happyDomain service envelopes are opaque to us (the registry is in the
// host), so we scan the JSON for any nested "Rrtype": <number> field —
// every dns.RR_Header instance carries one, which catches MX, CAA,
// orphan, CNAME, SRV, … without needing a per-service decoder.
func typesFromService(svc rawService) []uint16 {
	if len(svc.Service) == 0 {
		return nil
	}
	var v any
	if err := json.Unmarshal(svc.Service, &v); err != nil {
		return nil
	}
	seen := map[uint16]bool{}
	collectRrtypes(v, seen)
	if len(seen) == 0 {
		return nil
	}
	out := make([]uint16, 0, len(seen))
	for q := range seen {
		out = append(out, q)
	}
	return out
}

func collectRrtypes(v any, out map[uint16]bool) {
	switch x := v.(type) {
	case map[string]any:
		for k, vv := range x {
			if k == "Rrtype" {
				if n, ok := vv.(float64); ok && n > 0 && n < 65536 {
					out[uint16(n)] = true
				}
				continue
			}
			collectRrtypes(vv, out)
		}
	case []any:
		for _, vv := range x {
			collectRrtypes(vv, out)
		}
	}
}

type invalidTypesError struct{ raw string }

func (e *invalidTypesError) Error() string {
	return "no valid record types in \"" + e.raw + "\""
}