package checker

import (
	"context"
	"fmt"
	"sort"

	"github.com/miekg/dns"
	sdk "git.happydns.org/checker-sdk-go/checker"
)

// allDNSKEYs flattens every server's DNSKEY into a single deduplicated slice,
// keyed by (KeyTag, Algorithm). The first occurrence wins because all servers
// should agree (dnssec_dnskey_consistent enforces that separately).
func allDNSKEYs(d *DNSSECData) []DNSKEYRecord {
	seen := map[string]DNSKEYRecord{}
	for _, name := range sortedServers(d) {
		for _, k := range d.Servers[name].DNSKEYs {
			id := fmt.Sprintf("%d/%d", k.KeyTag, k.Algorithm)
			if _, ok := seen[id]; !ok {
				seen[id] = k
			}
		}
	}
	out := make([]DNSKEYRecord, 0, len(seen))
	for _, v := range seen {
		out = append(out, v)
	}
	sort.Slice(out, func(i, j int) bool { return out[i].KeyTag < out[j].KeyTag })
	return out
}

type algorithmAllowedRule struct{}

func (algorithmAllowedRule) Name() string { return "dnssec_algorithm_allowed" }
func (algorithmAllowedRule) Description() string {
	return "Rejects DNSKEYs that use a forbidden algorithm or are not in the allowed list."
}

func (algorithmAllowedRule) Evaluate(ctx context.Context, obs sdk.ObservationGetter, opts sdk.CheckerOptions) []sdk.CheckState {
	data, errState := loadDNSSEC(ctx, obs)
	if errState != nil {
		return errState
	}
	keys := allDNSKEYs(data)
	if len(keys) == 0 {
		return skipped("no DNSKEY observed")
	}
	allowed := defaultAllowedAlgorithms()
	if v, ok := sdk.GetOption[[]uint8](opts, "allowedAlgorithms"); ok && len(v) > 0 {
		allowed = v
	}
	forbidden := defaultForbiddenAlgorithms()
	if v, ok := sdk.GetOption[[]uint8](opts, "forbiddenAlgorithms"); ok && len(v) > 0 {
		forbidden = v
	}
	allowedSet := map[uint8]bool{}
	for _, a := range allowed {
		allowedSet[a] = true
	}
	forbiddenSet := map[uint8]bool{}
	for _, a := range forbidden {
		forbiddenSet[a] = true
	}

	var states []sdk.CheckState
	for _, k := range keys {
		switch {
		case forbiddenSet[k.Algorithm]:
			states = append(states, withMeta(sdk.CheckState{
				Status:  sdk.StatusCrit,
				Subject: fmt.Sprintf("KeyTag %d", k.KeyTag),
				Message: fmt.Sprintf("DNSKEY uses forbidden algorithm %d (%s)", k.Algorithm, dns.AlgorithmToString[k.Algorithm]),
			}, "Roll the key to a modern algorithm: 13 (ECDSAP256SHA256) or 15 (Ed25519).", "dnssec.algorithm_disallowed"))
		case !allowedSet[k.Algorithm]:
			states = append(states, withMeta(sdk.CheckState{
				Status:  sdk.StatusWarn,
				Subject: fmt.Sprintf("KeyTag %d", k.KeyTag),
				Message: fmt.Sprintf("DNSKEY uses algorithm %d (%s), not in the allowed list", k.Algorithm, dns.AlgorithmToString[k.Algorithm]),
			}, "Add the algorithm to allowedAlgorithms or roll the key to one of: 8, 13, 14, 15, 16.", "dnssec.algorithm_disallowed"))
		default:
			states = append(states, sdk.CheckState{
				Status:  sdk.StatusOK,
				Subject: fmt.Sprintf("KeyTag %d", k.KeyTag),
				Message: fmt.Sprintf("DNSKEY algorithm %d (%s) accepted", k.Algorithm, dns.AlgorithmToString[k.Algorithm]),
			})
		}
	}
	return states
}

type algorithmModernRule struct{}

func (algorithmModernRule) Name() string { return "dnssec_algorithm_modern" }
func (algorithmModernRule) Description() string {
	return "Recommends ECDSAP256SHA256 (13) or Ed25519 (15) over RSA."
}

func (algorithmModernRule) Evaluate(ctx context.Context, obs sdk.ObservationGetter, _ sdk.CheckerOptions) []sdk.CheckState {
	data, errState := loadDNSSEC(ctx, obs)
	if errState != nil {
		return errState
	}
	keys := allDNSKEYs(data)
	if len(keys) == 0 {
		return skipped("no DNSKEY observed")
	}
	hasModern := false
	hasLegacy := false
	for _, k := range keys {
		switch k.Algorithm {
		case dns.ECDSAP256SHA256, dns.ECDSAP384SHA384, dns.ED25519, dns.ED448:
			hasModern = true
		case dns.RSASHA256, dns.RSASHA512, dns.RSASHA1, dns.RSASHA1NSEC3SHA1:
			hasLegacy = true
		}
	}
	switch {
	case hasModern && !hasLegacy:
		return okState(data.Domain, "zone uses modern elliptic-curve algorithms (13/14/15/16)")
	case hasLegacy:
		return []sdk.CheckState{withMeta(sdk.CheckState{
			Status:  sdk.StatusWarn,
			Subject: data.Domain,
			Message: "zone still uses RSA-family DNSKEYs; modern operators prefer 13 (ECDSAP256SHA256) or 15 (Ed25519) for smaller responses and faster validation",
		}, "Plan an algorithm rollover. `dnssec-keygen -a ECDSAP256SHA256 -K /var/lib/bind <zone>` (BIND), then add the new key, wait for the parent's DS to update, then drop the old key.",
			"dnssec.algorithm_legacy")}
	}
	return []sdk.CheckState{{
		Status:  sdk.StatusInfo,
		Subject: data.Domain,
		Message: "no modern or legacy algorithms detected; review DNSKEY policy manually",
	}}
}

type rsaKeySizeRule struct{}

func (rsaKeySizeRule) Name() string { return "dnssec_rsa_keysize" }
func (rsaKeySizeRule) Description() string {
	return "Verifies RSA DNSKEYs reach a minimum modulus size (default 2048 bits)."
}

func (rsaKeySizeRule) Evaluate(ctx context.Context, obs sdk.ObservationGetter, opts sdk.CheckerOptions) []sdk.CheckState {
	data, errState := loadDNSSEC(ctx, obs)
	if errState != nil {
		return errState
	}
	keys := allDNSKEYs(data)
	if len(keys) == 0 {
		return skipped("no DNSKEY observed")
	}
	minSize := optionUint(opts, "minRSAKeySize", defaultMinRSAKeySize)

	var states []sdk.CheckState
	for _, k := range keys {
		switch k.Algorithm {
		case dns.RSASHA1, dns.RSASHA1NSEC3SHA1, dns.RSASHA256, dns.RSASHA512:
		default:
			continue
		}
		if k.KeySize == 0 {
			continue // could not estimate; rule_keys parser limitation
		}
		switch {
		case k.KeySize < 1024:
			states = append(states, withMeta(sdk.CheckState{
				Status:  sdk.StatusCrit,
				Subject: fmt.Sprintf("KeyTag %d", k.KeyTag),
				Message: fmt.Sprintf("RSA DNSKEY %d uses a %d-bit modulus: practically broken", k.KeyTag, k.KeySize),
			}, "Roll the key to at least 2048-bit RSA, or better, ECDSAP256SHA256 (algo 13).", "dnssec.rsa_keysize_small"))
		case uint(k.KeySize) < minSize:
			states = append(states, withMeta(sdk.CheckState{
				Status:  sdk.StatusWarn,
				Subject: fmt.Sprintf("KeyTag %d", k.KeyTag),
				Message: fmt.Sprintf("RSA DNSKEY %d uses a %d-bit modulus (recommended ≥ %d)", k.KeyTag, k.KeySize, minSize),
			}, "Roll to a 2048-bit RSA key, or migrate to ECDSAP256SHA256 (algo 13).", "dnssec.rsa_keysize_small"))
		default:
			states = append(states, sdk.CheckState{
				Status:  sdk.StatusOK,
				Subject: fmt.Sprintf("KeyTag %d", k.KeyTag),
				Message: fmt.Sprintf("RSA DNSKEY %d uses a %d-bit modulus", k.KeyTag, k.KeySize),
			})
		}
	}
	if len(states) == 0 {
		return okState(data.Domain, "no RSA DNSKEY in use")
	}
	return states
}

type kskPresentRule struct{}

func (kskPresentRule) Name() string { return "dnssec_ksk_present" }
func (kskPresentRule) Description() string {
	return "Verifies at least one DNSKEY has the SEP bit (KSK)."
}

func (kskPresentRule) Evaluate(ctx context.Context, obs sdk.ObservationGetter, opts sdk.CheckerOptions) []sdk.CheckState {
	data, errState := loadDNSSEC(ctx, obs)
	if errState != nil {
		return errState
	}
	keys := allDNSKEYs(data)
	if len(keys) == 0 {
		return skipped("no DNSKEY observed")
	}
	required := sdk.GetBoolOption(opts, "requireSEP", defaultRequireSEP)
	if !required {
		return okState(data.Domain, "requireSEP=false")
	}
	for _, k := range keys {
		if k.IsKSK {
			return okState(data.Domain, fmt.Sprintf("KSK present (KeyTag %d, algorithm %d)", k.KeyTag, k.Algorithm))
		}
	}
	return []sdk.CheckState{withMeta(sdk.CheckState{
		Status:  sdk.StatusCrit,
		Subject: data.Domain,
		Message: "no DNSKEY carries the SEP (KSK) flag",
	}, "Re-publish the apex DNSKEY RRset with at least one key flagged as SEP (flags 257). Most signers do this automatically; check that the KSK was not accidentally removed during a rollover.",
		"dnssec.no_ksk")}
}

type dnskeyCountRule struct{}

func (dnskeyCountRule) Name() string { return "dnssec_dnskey_count" }
func (dnskeyCountRule) Description() string {
	return "Warns when too many DNSKEYs are published, inflating responses and amplification potential."
}

func (dnskeyCountRule) Evaluate(ctx context.Context, obs sdk.ObservationGetter, _ sdk.CheckerOptions) []sdk.CheckState {
	data, errState := loadDNSSEC(ctx, obs)
	if errState != nil {
		return errState
	}
	keys := allDNSKEYs(data)
	n := len(keys)
	switch {
	case n == 0:
		return skipped("no DNSKEY observed")
	case n >= 8:
		return []sdk.CheckState{withMeta(sdk.CheckState{
			Status:  sdk.StatusWarn,
			Subject: data.Domain,
			Message: fmt.Sprintf("%d DNSKEYs published; large RRsets bloat responses and increase amplification factor", n),
		}, "Drop retired keys after their successor has fully rolled.", "dnssec.dnskey_count")}
	default:
		return okState(data.Domain, fmt.Sprintf("%d DNSKEY(s) published", n))
	}
}