package checker

import (
	"bytes"
	"context"
	"crypto/dsa"
	"crypto/ecdsa"
	"crypto/ed25519"
	"crypto/rsa"
	"crypto/sha256"
	"crypto/x509"
	"encoding/base64"
	"encoding/hex"
	"encoding/json"
	"encoding/pem"
	"fmt"
	"strings"
	"time"

	"github.com/ProtonMail/go-crypto/openpgp"
	"github.com/ProtonMail/go-crypto/openpgp/packet"
	"github.com/miekg/dns"

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

// maxKeyMaterialBytes caps the decoded byte size of an OPENPGPKEY
// payload or an SMIMEA certificate before it is handed to the parser.
// Anything larger is rejected outright to keep parser costs bounded; a
// rule (e.g. RulePGPRecordTooLarge at 4 KiB) flags more conservative
// limits separately. 64 KiB is well above any legitimate OpenPGP key
// size while staying clear of pathological input.
const maxKeyMaterialBytes = 64 * 1024

// serviceBody is the common envelope for the two services.
type serviceBody struct {
	Username string          `json:"username,omitempty"`
	OpenPGP  *dns.OPENPGPKEY `json:"openpgpkey,omitempty"`
	SMIMEA   *dns.SMIMEA     `json:"smimea,omitempty"`
}

// Collect runs the DANE-email data gathering pipeline and returns an
// *EmailKeyData carrying raw facts (DNS outcome, parsed key / cert
// structure). Judgment, severity, fix hints, option-driven thresholds,
// is deferred to the rules. A non-nil error is returned only for
// unrecoverable input problems (missing options, unknown service type).
func (p *emailKeyProvider) Collect(ctx context.Context, opts sdk.CheckerOptions) (any, error) {
	svcMsg, err := serviceFromOptions(opts)
	if err != nil {
		return nil, err
	}

	kind := kindForServiceType(svcMsg.Type)
	if kind == "" {
		return nil, fmt.Errorf("service type %q is not supported by this checker", svcMsg.Type)
	}

	var body serviceBody
	if err := json.Unmarshal(svcMsg.Service, &body); err != nil {
		return nil, fmt.Errorf("decode service body: %w", err)
	}

	originOpt, _ := sdk.GetOption[string](opts, "domain_name")
	subdomainOpt, _ := sdk.GetOption[string](opts, "subdomain")
	resolverOpt, _ := sdk.GetOption[string](opts, OptionResolver)

	origin := strings.TrimSuffix(firstNonEmpty(originOpt, svcMsg.Domain), ".")
	if origin == "" {
		return nil, fmt.Errorf("missing 'domain_name' option")
	}
	parent := joinSubdomain(subdomainOpt, origin)

	data := &EmailKeyData{
		Kind:        kind,
		Domain:      dns.Fqdn(origin),
		Subdomain:   strings.TrimSuffix(subdomainOpt, "."),
		Username:    body.Username,
		CollectedAt: time.Now().UTC(),
	}

	prefix := OpenPGPKeyPrefix
	if kind == KindSMIMEA {
		prefix = SMIMEACertPrefix
	}
	expectedOwner, recordedOwner := computeOwner(body, prefix, parent)
	data.ExpectedOwner = expectedOwner
	data.QueriedOwner = firstNonEmpty(recordedOwner, expectedOwner)

	// Owner-name hash inputs: rules compare the two and decide.
	if data.Username != "" {
		data.ExpectedOwnerPrefix = ownerHashHex(data.Username)
		data.ObservedOwnerPrefix = extractOwnerPrefix(data.QueriedOwner, prefix, parent)
	}

	// DNS lookup + DNSSEC flag.
	if data.QueriedOwner != "" {
		servers := resolvers(resolverOpt)
		qtype := dns.TypeOPENPGPKEY
		if kind == KindSMIMEA {
			qtype = dns.TypeSMIMEA
		}
		ans, err := lookup(ctx, servers, data.QueriedOwner, qtype)
		if err != nil {
			data.DNSQueryError = fmt.Sprintf("DNS lookup for %s %s failed: %v", dns.TypeToString[qtype], data.QueriedOwner, err)
		} else {
			data.Resolver = ans.Server
			secure := ans.AD
			data.DNSSECSecure = &secure
			data.RecordCount = len(ans.Records)

			present := !(ans.Rcode == dns.RcodeNameError || len(ans.Records) == 0)
			data.DNSAnswerPresent = &present

			// Compare DNS-returned record bytes with the service-declared ones
			// only when we actually have records to compare and a reference.
			if present {
				var match bool
				switch {
				case kind == KindOpenPGPKey && body.OpenPGP != nil:
					match = anyOpenPGPMatches(ans.Records, body.OpenPGP)
					data.DNSRecordMatchesService = &match
				case kind == KindSMIMEA && body.SMIMEA != nil:
					match = anySMIMEAMatches(ans.Records, body.SMIMEA)
					data.DNSRecordMatchesService = &match
				}
			}
		}
	}

	// Parse the payload from the service body (so rules can evaluate even
	// when the DNS lookup failed to reach the authoritative servers).
	if kind == KindOpenPGPKey {
		data.OpenPGP = analyzeOpenPGP(body)
	} else {
		data.SMIMEA = analyzeSMIMEA(body)
	}

	return data, nil
}

// serviceFromOptions pulls the "service" option out of the options map,
// accepting both the in-process plugin path (native Go value) and the
// HTTP path (JSON-decoded map[string]any). Normalising via a JSON
// round-trip keeps both paths working without importing the upstream
// type.
func serviceFromOptions(opts sdk.CheckerOptions) (*serviceMessage, error) {
	v, ok := opts["service"]
	if !ok {
		return nil, fmt.Errorf("service option missing")
	}
	raw, err := json.Marshal(v)
	if err != nil {
		return nil, fmt.Errorf("marshal service option: %w", err)
	}
	var svc serviceMessage
	if err := json.Unmarshal(raw, &svc); err != nil {
		return nil, fmt.Errorf("decode service option: %w", err)
	}
	// Fall back to the service_type option when the envelope doesn't
	// carry _svctype (older hosts).
	if svc.Type == "" {
		if st, ok := sdk.GetOption[string](opts, "service_type"); ok {
			svc.Type = st
		}
	}
	return &svc, nil
}

func kindForServiceType(t string) string {
	switch t {
	case ServiceOpenPGP:
		return KindOpenPGPKey
	case ServiceSMimeCert:
		return KindSMIMEA
	default:
		return ""
	}
}

// ownerHashHex returns the RFC 7929 / 8162 label: hex(sha256(localpart)[:28]).
func ownerHashHex(username string) string {
	sum := sha256.Sum256([]byte(username))
	return hex.EncodeToString(sum[:DANEOwnerHashSize])
}

// computeOwner derives the expected FQDN from the service body. It
// returns the expected-by-specification owner and, when the service
// body carries its own Hdr.Name, the recorded owner, so we can detect
// discrepancies between the two.
func computeOwner(body serviceBody, prefix, parent string) (expected, recorded string) {
	if body.Username != "" {
		expected = dns.Fqdn(ownerHashHex(body.Username) + "." + strings.TrimPrefix(prefix, "") + "." + strings.TrimSuffix(parent, "."))
		// Normalise: no double dots.
		expected = strings.Replace(expected, "..", ".", -1)
	}
	// happyDomain encodes service-embedded record owners relative to the
	// parent zone, so we must join with parent before treating as FQDN.
	switch {
	case body.OpenPGP != nil && body.OpenPGP.Hdr.Name != "":
		recorded = dns.Fqdn(sdk.JoinRelative(body.OpenPGP.Hdr.Name, parent))
	case body.SMIMEA != nil && body.SMIMEA.Hdr.Name != "":
		recorded = dns.Fqdn(sdk.JoinRelative(body.SMIMEA.Hdr.Name, parent))
	}
	return
}

// extractOwnerPrefix pulls the leading label from an owner name of the
// form <hash>._openpgpkey.<...> (or _smimecert), returning the hash
// portion only. Returns "" when the owner does not follow that shape.
func extractOwnerPrefix(owner, prefix, parent string) string {
	owner = strings.TrimSuffix(strings.ToLower(owner), ".")
	// Look for ".<prefix>." just after the first label.
	marker := "." + prefix + "."
	if i := strings.Index(owner, marker); i > 0 {
		return owner[:i]
	}
	return ""
}

// anyOpenPGPMatches reports whether any of rrs carries the same public
// key bytes as ref.
func anyOpenPGPMatches(rrs []dns.RR, ref *dns.OPENPGPKEY) bool {
	want := strings.TrimSpace(ref.PublicKey)
	for _, rr := range rrs {
		if r, ok := rr.(*dns.OPENPGPKEY); ok && strings.TrimSpace(r.PublicKey) == want {
			return true
		}
	}
	return false
}

// anySMIMEAMatches reports whether any of rrs matches ref on (usage,
// selector, matching type, certificate bytes).
func anySMIMEAMatches(rrs []dns.RR, ref *dns.SMIMEA) bool {
	want := strings.ToLower(strings.TrimSpace(ref.Certificate))
	for _, rr := range rrs {
		r, ok := rr.(*dns.SMIMEA)
		if !ok {
			continue
		}
		if r.Usage == ref.Usage && r.Selector == ref.Selector && r.MatchingType == ref.MatchingType &&
			strings.ToLower(strings.TrimSpace(r.Certificate)) == want {
			return true
		}
	}
	return false
}

// ── OpenPGP analysis ─────────────────────────────────────────────────────────

// analyzeOpenPGP parses the OpenPGP key from the service record and
// returns a structured fact summary. When parsing fails, ParseError is
// populated and the rest of the fields hold whatever could be recovered.
func analyzeOpenPGP(body serviceBody) *OpenPGPInfo {
	if body.OpenPGP == nil {
		return &OpenPGPInfo{ParseError: "Service body has no OPENPGPKEY record."}
	}

	encoded := body.OpenPGP.PublicKey
	// Reject pathological payloads before allocating: the base64-decoded
	// size is at most ceil(len(encoded)*3/4).
	if len(encoded)/4*3 > maxKeyMaterialBytes {
		return &OpenPGPInfo{
			RawSize:    len(encoded) / 4 * 3,
			ParseError: fmt.Sprintf("OPENPGPKEY payload exceeds the %d-byte parse limit.", maxKeyMaterialBytes),
		}
	}
	raw, err := base64.StdEncoding.DecodeString(encoded)
	if err != nil {
		return &OpenPGPInfo{ParseError: fmt.Sprintf("OPENPGPKEY record carries invalid base64: %v", err)}
	}
	if len(raw) > maxKeyMaterialBytes {
		return &OpenPGPInfo{
			RawSize:    len(raw),
			ParseError: fmt.Sprintf("OPENPGPKEY payload exceeds the %d-byte parse limit.", maxKeyMaterialBytes),
		}
	}

	info := &OpenPGPInfo{RawSize: len(raw)}

	entities, err := openpgp.ReadKeyRing(bytes.NewReader(raw))
	if err != nil || len(entities) == 0 {
		if err == nil {
			err = fmt.Errorf("no OpenPGP entity found")
		}
		info.ParseError = fmt.Sprintf("Cannot parse OpenPGP key: %v", err)
		return info
	}

	info.EntityCount = len(entities)

	ent := entities[0]
	pub := ent.PrimaryKey
	info.CreatedAt = pub.CreationTime
	info.Fingerprint = strings.ToUpper(hex.EncodeToString(pub.Fingerprint))
	info.KeyID = fmt.Sprintf("%016X", pub.KeyId)
	info.PrimaryAlgorithm = algorithmName(pub)
	info.PrimaryBits = publicKeyBits(pub)

	for name := range ent.Identities {
		info.UIDs = append(info.UIDs, name)
	}

	if len(ent.Revocations) > 0 {
		info.Revoked = true
	}

	// Expiry on the primary key, derived from the self-signature.
	now := time.Now()
	if selfSig, _ := ent.PrimarySelfSignature(); selfSig != nil {
		if selfSig.KeyLifetimeSecs != nil && *selfSig.KeyLifetimeSecs > 0 {
			info.ExpiresAt = pub.CreationTime.Add(time.Duration(*selfSig.KeyLifetimeSecs) * time.Second)
		}
	}

	// UID vs username matching.
	if len(ent.Identities) > 0 && body.Username != "" {
		wantedLocal := strings.ToLower(body.Username)
		matched := false
		for name := range ent.Identities {
			if strings.Contains(strings.ToLower(name), "<"+wantedLocal+"@") ||
				strings.Contains(strings.ToLower(name), wantedLocal+"@") {
				matched = true
				break
			}
		}
		info.MatchesUsername = &matched
	}

	// Subkeys + encryption capability.
	for _, sk := range ent.Subkeys {
		si := SubkeyInfo{
			Algorithm: algorithmName(sk.PublicKey),
			Bits:      publicKeyBits(sk.PublicKey),
			CreatedAt: sk.PublicKey.CreationTime,
			Revoked:   len(sk.Revocations) > 0,
		}
		if sk.Sig != nil {
			if sk.Sig.FlagsValid {
				si.CanSign = sk.Sig.FlagSign
				si.CanEncrypt = sk.Sig.FlagEncryptCommunications || sk.Sig.FlagEncryptStorage
				si.CanAuth = sk.Sig.FlagAuthenticate
			}
			if sk.Sig.KeyLifetimeSecs != nil && *sk.Sig.KeyLifetimeSecs > 0 {
				si.ExpiresAt = sk.PublicKey.CreationTime.Add(time.Duration(*sk.Sig.KeyLifetimeSecs) * time.Second)
			}
		}
		info.Subkeys = append(info.Subkeys, si)

		if si.CanEncrypt && !si.Revoked && (si.ExpiresAt.IsZero() || si.ExpiresAt.After(now)) {
			info.HasEncryptionCapability = true
		}
	}
	// Primary can also be an encryption key if flagged so.
	if selfSig, _ := ent.PrimarySelfSignature(); selfSig != nil && selfSig.FlagsValid &&
		(selfSig.FlagEncryptCommunications || selfSig.FlagEncryptStorage) &&
		!info.Revoked && (info.ExpiresAt.IsZero() || info.ExpiresAt.After(now)) {
		info.HasEncryptionCapability = true
	}

	return info
}

func algorithmName(pub *packet.PublicKey) string {
	switch pub.PubKeyAlgo {
	case packet.PubKeyAlgoRSA, packet.PubKeyAlgoRSAEncryptOnly, packet.PubKeyAlgoRSASignOnly:
		return "RSA"
	case packet.PubKeyAlgoDSA:
		return "DSA"
	case packet.PubKeyAlgoElGamal:
		return "ElGamal"
	case packet.PubKeyAlgoECDH:
		return "ECDH"
	case packet.PubKeyAlgoECDSA:
		return "ECDSA"
	case packet.PubKeyAlgoEdDSA:
		return "EdDSA"
	case packet.PubKeyAlgoX25519:
		return "X25519"
	case packet.PubKeyAlgoX448:
		return "X448"
	case packet.PubKeyAlgoEd25519:
		return "Ed25519"
	case packet.PubKeyAlgoEd448:
		return "Ed448"
	default:
		return fmt.Sprintf("algo-%d", pub.PubKeyAlgo)
	}
}

func publicKeyBits(pub *packet.PublicKey) int {
	if pub == nil {
		return 0
	}
	switch k := pub.PublicKey.(type) {
	case *rsa.PublicKey:
		if k == nil || k.N == nil {
			return 0
		}
		return k.N.BitLen()
	case *dsa.PublicKey:
		if k == nil || k.P == nil {
			return 0
		}
		return k.P.BitLen()
	case *ecdsa.PublicKey:
		if k == nil || k.Params() == nil {
			return 0
		}
		return k.Params().BitSize
	case ed25519.PublicKey:
		return 256
	}
	// Fallback to the packet's advertised length.
	if n, err := pub.BitLength(); err == nil {
		return int(n)
	}
	return 0
}

// ── SMIMEA analysis ──────────────────────────────────────────────────────────

// analyzeSMIMEA parses the SMIMEA certificate and returns a structured
// fact summary. When parsing fails, ParseError is populated.
func analyzeSMIMEA(body serviceBody) *SMIMEAInfo {
	if body.SMIMEA == nil {
		return &SMIMEAInfo{ParseError: "Service body has no SMIMEA record."}
	}
	rec := body.SMIMEA

	info := &SMIMEAInfo{
		Usage:        rec.Usage,
		Selector:     rec.Selector,
		MatchingType: rec.MatchingType,
		HashHex:      strings.ToLower(rec.Certificate),
	}

	// Matching types 1 and 2 only carry a digest; no certificate or SPKI
	// to parse. Rules surface that; here we just stop.
	if rec.MatchingType != 0 {
		return info
	}

	if len(rec.Certificate)/2 > maxKeyMaterialBytes {
		info.ParseError = fmt.Sprintf("SMIMEA payload exceeds the %d-byte parse limit.", maxKeyMaterialBytes)
		return info
	}
	der, err := hex.DecodeString(rec.Certificate)
	if err != nil || len(der) == 0 {
		info.ParseError = fmt.Sprintf("Cannot decode certificate bytes: %v", err)
		return info
	}
	if len(der) > maxKeyMaterialBytes {
		info.ParseError = fmt.Sprintf("SMIMEA payload exceeds the %d-byte parse limit.", maxKeyMaterialBytes)
		return info
	}

	// Selector 1 carries only a SubjectPublicKeyInfo; parse it that way.
	if rec.Selector == 1 {
		info.PublicKey = analyzeSPKI(der, info)
		return info
	}

	cert, err := x509.ParseCertificate(der)
	if err != nil {
		// Try a PEM fallback for robustness.
		if block, _ := pem.Decode(der); block != nil && block.Type == "CERTIFICATE" {
			cert, err = x509.ParseCertificate(block.Bytes)
		}
	}
	if err != nil || cert == nil {
		if err == nil {
			err = fmt.Errorf("no certificate found")
		}
		info.ParseError = fmt.Sprintf("Cannot parse X.509 certificate: %v", err)
		return info
	}

	ci := &CertInfo{
		Subject:            cert.Subject.String(),
		Issuer:             cert.Issuer.String(),
		SerialHex:          strings.ToUpper(hex.EncodeToString(cert.SerialNumber.Bytes())),
		NotBefore:          cert.NotBefore,
		NotAfter:           cert.NotAfter,
		SignatureAlgorithm: cert.SignatureAlgorithm.String(),
		PublicKeyAlgorithm: cert.PublicKeyAlgorithm.String(),
		EmailAddresses:     cert.EmailAddresses,
		DNSNames:           cert.DNSNames,
		IsCA:               cert.IsCA,
	}
	ci.IsSelfSigned = cert.Subject.String() == cert.Issuer.String() && cert.CheckSignatureFrom(cert) == nil
	ci.PublicKeyBits = x509PublicKeyBits(cert.PublicKey)

	for _, eku := range cert.ExtKeyUsage {
		if eku == x509.ExtKeyUsageEmailProtection {
			ci.HasEmailProtectionEKU = true
		}
	}
	if cert.KeyUsage&x509.KeyUsageDigitalSignature != 0 {
		ci.HasDigitalSignature = true
	}
	if cert.KeyUsage&x509.KeyUsageKeyEncipherment != 0 {
		ci.HasKeyEncipherment = true
	}

	// Email-address / username pairing fact.
	if body.Username != "" && len(cert.EmailAddresses) > 0 {
		wantPrefix := strings.ToLower(body.Username) + "@"
		matched := false
		for _, e := range cert.EmailAddresses {
			if strings.HasPrefix(strings.ToLower(e), wantPrefix) {
				matched = true
				break
			}
		}
		ci.EmailMatchesUsername = &matched
	}

	info.Certificate = ci
	return info
}

func analyzeSPKI(der []byte, info *SMIMEAInfo) *PubKeyInfo {
	pub, err := x509.ParsePKIXPublicKey(der)
	if err != nil {
		info.ParseError = fmt.Sprintf("Cannot parse SubjectPublicKeyInfo: %v", err)
		return nil
	}
	pk := &PubKeyInfo{Bits: x509PublicKeyBits(pub)}
	switch pub.(type) {
	case *rsa.PublicKey:
		pk.Algorithm = "RSA"
	case *ecdsa.PublicKey:
		pk.Algorithm = "ECDSA"
	case ed25519.PublicKey:
		pk.Algorithm = "Ed25519"
	default:
		pk.Algorithm = fmt.Sprintf("%T", pub)
	}
	return pk
}

func x509PublicKeyBits(pub any) int {
	switch k := pub.(type) {
	case *rsa.PublicKey:
		if k == nil || k.N == nil {
			return 0
		}
		return k.N.BitLen()
	case *ecdsa.PublicKey:
		if k == nil || k.Params() == nil {
			return 0
		}
		return k.Params().BitSize
	case ed25519.PublicKey:
		return 256
	}
	return 0
}

func firstNonEmpty(vals ...string) string {
	for _, v := range vals {
		if strings.TrimSpace(v) != "" {
			return v
		}
	}
	return ""
}