package checker

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

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

// Rules returns the full list of CheckRules exposed by the DANE checker.
// Each rule covers exactly one concern so the UI can show per-concern
// status rather than a single monolithic rule that multiplexes many codes.
func Rules() []sdk.CheckRule {
	return []sdk.CheckRule{
		&hasRecordsRule{},
		&dnssecValidatedRule{},
		&probeAvailableRule{},
		&handshakeOKRule{},
		&recordsMatchChainRule{},
		&pkixChainValidRule{},
		&usageCoherentRule{},
	}
}

// ruleContext bundles the data rules typically need: the checker's own
// observation plus the map of related TLS probes keyed by endpoint Ref.
type ruleContext struct {
	data   DANEData
	probes map[string]*tls.TLSProbe
	// relatedErr is a non-fatal error encountered while loading related
	// probes (e.g. the cross-checker lineage was unreachable). Rules
	// surface it as an error state so operators can spot misconfiguration.
	relatedErr error
	// err is a fatal error loading the checker's own observation.
	err error
}

// loadRuleContext fetches the DANE observation and the related TLS probes.
// Rules call this once and then filter on the fields they care about.
func loadRuleContext(ctx context.Context, obs sdk.ObservationGetter) *ruleContext {
	rc := &ruleContext{}
	if err := obs.Get(ctx, ObservationKeyDANE, &rc.data); err != nil {
		rc.err = err
		return rc
	}
	rc.probes, rc.relatedErr = relatedTLSProbes(ctx, obs)
	return rc
}

// observationErrorState is the canonical short-circuit state emitted when a
// rule cannot load the DANE observation at all.
func observationErrorState(err error) sdk.CheckState {
	return sdk.CheckState{
		Status:  sdk.StatusError,
		Message: fmt.Sprintf("Failed to read %s: %v", ObservationKeyDANE, err),
		Code:    "dane_observation_error",
	}
}

// targetMeta builds the common Meta map for per-endpoint states.
func targetMeta(t TargetResult) map[string]any {
	return map[string]any{
		"host":     t.Host,
		"port":     t.Port,
		"proto":    t.Proto,
		"owner":    t.Owner,
		"starttls": t.STARTTLS,
		"records":  len(t.Records),
	}
}

// targetSubject is the human-readable subject tag used on per-endpoint states.
func targetSubject(t TargetResult) string {
	return fmt.Sprintf("%s:%d (%s)", t.Host, t.Port, t.Proto)
}

// probeUsable reports whether p carries a successfully-observed certificate
// chain. Rules that need to compare against the chain skip endpoints where
// this is false; the missing/failed cases are surfaced by probeAvailableRule
// and handshakeOKRule respectively, so other rules stay focused.
func probeUsable(p *tls.TLSProbe) bool {
	return p != nil && p.Error == "" && len(p.Chain) > 0
}

// matchSummary aggregates per-target match outcomes so callers don't redo the
// per-record loop. firstUnmatchedIdx is -1 when every record matched.
type matchSummary struct {
	matched, unmatched   int
	firstUnmatchedIdx    int
	firstUnmatchedReason string
}

// summarizeMatches walks t.Records once and reports how many matched p's
// chain, plus the first unmatched index and reason for messaging.
func summarizeMatches(t TargetResult, p *tls.TLSProbe) matchSummary {
	s := matchSummary{firstUnmatchedIdx: -1}
	if p == nil {
		return s
	}
	for i, rec := range t.Records {
		ok, reason := matchRecord(rec, p)
		if ok {
			s.matched++
			continue
		}
		s.unmatched++
		if s.firstUnmatchedIdx < 0 {
			s.firstUnmatchedIdx = i
			s.firstUnmatchedReason = reason
		}
	}
	return s
}

// matchRecord returns true when rec matches some certificate at the chain
// slot implied by rec.Usage. reason explains the miss on a false return.
//
// Slot selection:
//
//   - Usage 1 (PKIX-EE) and 3 (DANE-EE): leaf only.
//   - Usage 0 (PKIX-TA) and 2 (DANE-TA): intermediates + the root the
//     server presented (if any). We match against every non-leaf cert the
//     server sent, because some deployments publish the root and some the
//     intermediate; either is a valid TA reference for the connection's
//     path.
func matchRecord(rec TLSARecord, p *tls.TLSProbe) (bool, string) {
	if len(p.Chain) == 0 {
		return false, "no certificates observed on the endpoint"
	}
	var slots []tls.CertInfo
	switch rec.Usage {
	case UsagePKIXEE, UsageDANEEE:
		slots = p.Chain[:1]
	case UsagePKIXTA, UsageDANETA:
		if len(p.Chain) > 1 {
			slots = p.Chain[1:]
		} else {
			// Self-signed / bundle with only a leaf: allow matching against
			// the leaf as a degenerate TA so the user gets a hash comparison
			// rather than a silent "no slot".
			slots = p.Chain[:1]
		}
	default:
		return false, fmt.Sprintf("unsupported TLSA usage %d", rec.Usage)
	}

	var lastErr string
	for _, c := range slots {
		got, err := recordCandidate(rec, c)
		if err != nil {
			lastErr = err.Error()
			continue
		}
		if strings.EqualFold(got, rec.Certificate) {
			return true, ""
		}
	}
	if lastErr != "" {
		return false, lastErr
	}
	return false, fmt.Sprintf("expected %s, got none matching in chain", truncHex(rec.Certificate))
}

// maxFullDERBytes caps the size of a "Full" (MatchingType 0) DER payload
// that this checker is willing to base64-decode and hex-encode. Real X.509
// certificates rarely exceed 8 KiB; 64 KiB leaves comfortable headroom for
// pathological-but-legitimate chains while preventing a hostile probe
// payload from forcing arbitrary heap allocations during evaluation.
const maxFullDERBytes = 64 * 1024

// decodeFullDER base64-decodes b after rejecting payloads whose decoded size
// would exceed maxFullDERBytes, so an attacker-controlled probe cannot make
// the rule allocate unbounded memory before the hex comparison.
func decodeFullDER(b string, what string) ([]byte, error) {
	// base64 decoded length is at most ceil(len(b)/4)*3; bail out cheaply
	// before allocating the destination buffer.
	if len(b)/4*3 > maxFullDERBytes {
		return nil, fmt.Errorf("%s exceeds %d bytes", what, maxFullDERBytes)
	}
	der, err := base64.StdEncoding.DecodeString(b)
	if err != nil {
		return nil, fmt.Errorf("decode %s: %w", what, err)
	}
	if len(der) > maxFullDERBytes {
		return nil, fmt.Errorf("%s exceeds %d bytes", what, maxFullDERBytes)
	}
	return der, nil
}

// recordCandidate returns the hex value the TLSA record should match for
// the (selector, matching_type) pair against this certificate slot. For
// matching_type 0 (Full), both sides are compared as hex-encoded DER.
func recordCandidate(rec TLSARecord, c tls.CertInfo) (string, error) {
	var source string
	switch rec.Selector {
	case SelectorCert:
		switch rec.MatchingType {
		case MatchingFull:
			der, err := decodeFullDER(c.DERBase64, "cert DER")
			if err != nil {
				return "", err
			}
			source = hex.EncodeToString(der)
		case MatchingSHA256:
			source = c.CertSHA256
		case MatchingSHA512:
			source = c.CertSHA512
		default:
			return "", fmt.Errorf("unsupported matching type %d", rec.MatchingType)
		}
	case SelectorSPKI:
		switch rec.MatchingType {
		case MatchingFull:
			spki, err := decodeFullDER(c.SPKIDERBase64, "SPKI DER")
			if err != nil {
				return "", err
			}
			source = hex.EncodeToString(spki)
		case MatchingSHA256:
			source = c.SPKISHA256
		case MatchingSHA512:
			source = c.SPKISHA512
		default:
			return "", fmt.Errorf("unsupported matching type %d", rec.MatchingType)
		}
	default:
		return "", fmt.Errorf("unsupported selector %d", rec.Selector)
	}
	return source, nil
}

// parseTLSProbeMap decodes one related-observation payload into its constituent
// probes, keyed by endpoint Ref. Returns nil on decode error (caller skips).
func parseTLSProbeMap(data []byte) map[string]tls.TLSProbe {
	var payload struct {
		Probes map[string]tls.TLSProbe `json:"probes"`
	}
	if err := json.Unmarshal(data, &payload); err != nil {
		return nil
	}
	return payload.Probes
}

// relatedTLSProbes indexes TLS probes fetched via GetRelated by endpoint Ref.
func relatedTLSProbes(ctx context.Context, obs sdk.ObservationGetter) (map[string]*tls.TLSProbe, error) {
	related, err := obs.GetRelated(ctx, tls.ObservationKeyTLSProbes)
	if err != nil {
		return nil, fmt.Errorf("related TLS probes unavailable: %w", err)
	}
	return indexProbes(related), nil
}

// indexProbes flattens a slice of related TLS-probe observations into a probe
// map keyed by endpoint Ref. Shared by the rule path (relatedTLSProbes) and
// the report path (GetHTMLReport), which receive the same RelatedObservation
// type from different SDK entry points.
func indexProbes(related []sdk.RelatedObservation) map[string]*tls.TLSProbe {
	out := map[string]*tls.TLSProbe{}
	for _, ro := range related {
		for k, v := range parseTLSProbeMap(ro.Data) {
			out[k] = &v
		}
	}
	return out
}

func truncHex(s string) string {
	if len(s) > 12 {
		return s[:12] + "…"
	}
	return s
}