// Package dkim implements DKIM signing (outbound) and verification (inbound)
// for RSA-2048 and Ed25519 keys per RFC 6376 and RFC 8463.
package dkim

import (
	gocrypto "crypto"
	"crypto/ed25519"
	"crypto/rand"
	"crypto/rsa"
	"crypto/sha256"
	"crypto/x509"
	"encoding/base64"
	"encoding/pem"
	"fmt"
	"net"
	"strings"
	"time"
)

// Signer holds a loaded private key and signing metadata.
type Signer struct {
	privateKey gocrypto.PrivateKey // *rsa.PrivateKey or ed25519.PrivateKey
	selector   string
	domain     string
	algo       string // "rsa2048" | "ed25519"
}

// GenerateKeyPair generates a DKIM key pair for the given algorithm.
// algo must be "rsa2048" or "ed25519".
// Returns PEM-encoded private key and PEM-encoded public key.
func GenerateKeyPair(algo string) (privateKeyPEM, publicKeyPEM string, err error) {
	switch algo {
	case "rsa2048":
		priv, genErr := rsa.GenerateKey(rand.Reader, 2048)
		if genErr != nil {
			return "", "", fmt.Errorf("dkim: generate RSA key: %w", genErr)
		}
		privDER := x509.MarshalPKCS1PrivateKey(priv)
		privBlock := &pem.Block{Type: "RSA PRIVATE KEY", Bytes: privDER}
		privateKeyPEM = string(pem.EncodeToMemory(privBlock))

		pubDER, marshalErr := x509.MarshalPKIXPublicKey(&priv.PublicKey)
		if marshalErr != nil {
			return "", "", fmt.Errorf("dkim: marshal RSA public key: %w", marshalErr)
		}
		pubBlock := &pem.Block{Type: "PUBLIC KEY", Bytes: pubDER}
		publicKeyPEM = string(pem.EncodeToMemory(pubBlock))
		return privateKeyPEM, publicKeyPEM, nil

	case "ed25519":
		pub, priv, genErr := ed25519.GenerateKey(rand.Reader)
		if genErr != nil {
			return "", "", fmt.Errorf("dkim: generate Ed25519 key: %w", genErr)
		}
		privDER, marshalErr := x509.MarshalPKCS8PrivateKey(priv)
		if marshalErr != nil {
			return "", "", fmt.Errorf("dkim: marshal Ed25519 private key: %w", marshalErr)
		}
		privBlock := &pem.Block{Type: "PRIVATE KEY", Bytes: privDER}
		privateKeyPEM = string(pem.EncodeToMemory(privBlock))

		pubDER, marshalErr := x509.MarshalPKIXPublicKey(pub)
		if marshalErr != nil {
			return "", "", fmt.Errorf("dkim: marshal Ed25519 public key: %w", marshalErr)
		}
		pubBlock := &pem.Block{Type: "PUBLIC KEY", Bytes: pubDER}
		publicKeyPEM = string(pem.EncodeToMemory(pubBlock))
		return privateKeyPEM, publicKeyPEM, nil

	default:
		return "", "", fmt.Errorf("dkim: unsupported algorithm %q (want rsa2048 or ed25519)", algo)
	}
}

// NewSigner parses a PEM private key and returns a Signer.
// Tries PKCS1 RSA first, then PKCS8 (Ed25519 or RSA).
func NewSigner(privateKeyPEM, domain, selector string) (*Signer, error) {
	if domain == "" {
		return nil, fmt.Errorf("dkim: domain must not be empty")
	}
	if selector == "" {
		return nil, fmt.Errorf("dkim: selector must not be empty")
	}

	block, _ := pem.Decode([]byte(privateKeyPEM))
	if block == nil {
		return nil, fmt.Errorf("dkim: failed to decode PEM block")
	}

	// Try PKCS1 RSA.
	if rsaKey, err := x509.ParsePKCS1PrivateKey(block.Bytes); err == nil {
		return &Signer{
			privateKey: rsaKey,
			selector:   selector,
			domain:     domain,
			algo:       "rsa2048",
		}, nil
	}

	// Try PKCS8 (covers Ed25519 and RSA PKCS8).
	key, err := x509.ParsePKCS8PrivateKey(block.Bytes)
	if err != nil {
		return nil, fmt.Errorf("dkim: parse private key: %w", err)
	}
	switch k := key.(type) {
	case ed25519.PrivateKey:
		return &Signer{
			privateKey: k,
			selector:   selector,
			domain:     domain,
			algo:       "ed25519",
		}, nil
	case *rsa.PrivateKey:
		return &Signer{
			privateKey: k,
			selector:   selector,
			domain:     domain,
			algo:       "rsa2048",
		}, nil
	default:
		return nil, fmt.Errorf("dkim: unsupported private key type %T", key)
	}
}

// Sign produces a DKIM-Signature header for the given RFC822 message.
// Returns the complete "DKIM-Signature: ..." header line (no trailing CRLF).
func (s *Signer) Sign(message []byte) (string, error) {
	// Split at first blank line (\r\n\r\n or \n\n).
	headerBytes, bodyBytes := splitMessage(message)

	// Canonicalize body (relaxed).
	canonBody := canonicalizeBodyRelaxed(bodyBytes)

	// Body hash.
	bodyHash := sha256.Sum256(canonBody)
	bh := base64.StdEncoding.EncodeToString(bodyHash[:])

	// Determine algorithm tag.
	var aTag string
	switch s.algo {
	case "ed25519":
		aTag = "ed25519-sha256"
	default:
		aTag = "rsa-sha256"
	}

	// Signed header fields (lower-case, in sign order).
	signedFields := "from:to:subject:date:message-id"

	// Build DKIM-Signature with b= empty.
	ts := fmt.Sprintf("%d", time.Now().Unix())
	sigHeader := fmt.Sprintf(
		"DKIM-Signature: v=1; a=%s; c=relaxed/relaxed; d=%s; s=%s; t=%s; bh=%s; h=%s; b=",
		aTag, s.domain, s.selector, ts, bh, signedFields,
	)

	// Canonicalize headers to sign + the sig header (b= empty).
	hdrMap := parseHeaders(headerBytes)
	var sb strings.Builder
	for _, field := range strings.Split(signedFields, ":") {
		field = strings.TrimSpace(field)
		val, ok := hdrMap[strings.ToLower(field)]
		if !ok {
			continue
		}
		sb.WriteString(canonicalizeHeaderRelaxed(field, val))
		sb.WriteString("\r\n")
	}
	// Append the DKIM-Signature line itself (with b= empty), canonicalized.
	sb.WriteString(canonicalizeHeaderRelaxed("dkim-signature", strings.TrimPrefix(sigHeader, "DKIM-Signature: ")))

	dataToSign := []byte(sb.String())
	hash := sha256.Sum256(dataToSign)

	var sigBytes []byte
	var signErr error

	switch s.algo {
	case "ed25519":
		privKey, ok := s.privateKey.(ed25519.PrivateKey)
		if !ok {
			return "", fmt.Errorf("dkim: private key type mismatch for ed25519")
		}
		// RFC 8463: ed25519-sha256 — sign the SHA-256 hash of the data.
		sigBytes = ed25519.Sign(privKey, hash[:])

	default:
		privKey, ok := s.privateKey.(*rsa.PrivateKey)
		if !ok {
			return "", fmt.Errorf("dkim: private key type mismatch for RSA")
		}
		sigBytes, signErr = rsa.SignPKCS1v15(rand.Reader, privKey, gocrypto.SHA256, hash[:])
		if signErr != nil {
			return "", fmt.Errorf("dkim: RSA sign: %w", signErr)
		}
	}

	b := base64.StdEncoding.EncodeToString(sigBytes)
	return sigHeader + b, nil
}

// Verify finds the DKIM-Signature in message, fetches the DNS public key,
// and verifies the signature. Returns the signing domain on success.
func Verify(message []byte) (domain string, err error) {
	headerBytes, bodyBytes := splitMessage(message)
	hdrMap := parseHeaders(headerBytes)

	sigVal, ok := hdrMap["dkim-signature"]
	if !ok {
		return "", fmt.Errorf("dkim: no DKIM-Signature header found")
	}

	params := parseDKIMParams(sigVal)

	sel, ok := params["s"]
	if !ok || sel == "" {
		return "", fmt.Errorf("dkim: missing selector (s=) in DKIM-Signature")
	}
	dom, ok := params["d"]
	if !ok || dom == "" {
		return "", fmt.Errorf("dkim: missing domain (d=) in DKIM-Signature")
	}
	bh64, _ := params["bh"]
	b64, ok := params["b"]
	if !ok {
		return "", fmt.Errorf("dkim: missing signature (b=) in DKIM-Signature")
	}
	signedFields, _ := params["h"]

	// Verify body hash.
	canonBody := canonicalizeBodyRelaxed(bodyBytes)
	bodyHash := sha256.Sum256(canonBody)
	expectedBH := base64.StdEncoding.EncodeToString(bodyHash[:])
	// Strip whitespace from DNS-retrieved bh for comparison.
	cleanBH := strings.Map(func(r rune) rune {
		if r == ' ' || r == '\t' || r == '\r' || r == '\n' {
			return -1
		}
		return r
	}, bh64)
	if cleanBH != expectedBH {
		return "", fmt.Errorf("dkim: body hash mismatch")
	}

	// DNS lookup.
	lookupName := sel + "._domainkey." + dom
	txts, lookupErr := net.LookupTXT(lookupName)
	if lookupErr != nil {
		return "", fmt.Errorf("dkim: DNS lookup %s: %w", lookupName, lookupErr)
	}
	if len(txts) == 0 {
		return "", fmt.Errorf("dkim: no TXT record at %s", lookupName)
	}

	// Join TXT record parts (DNS may split at 255 bytes).
	dnsRecord := strings.Join(txts, "")
	dnsParams := parseDKIMParams(dnsRecord)
	pVal, ok := dnsParams["p"]
	if !ok || pVal == "" {
		return "", fmt.Errorf("dkim: no p= (public key) in DNS TXT record")
	}

	pubDER, decErr := base64.StdEncoding.DecodeString(pVal)
	if decErr != nil {
		return "", fmt.Errorf("dkim: decode public key base64: %w", decErr)
	}

	// Rebuild the data-to-sign exactly as Signer.Sign did.
	// Canonicalize the signed headers.
	var sb strings.Builder
	for _, field := range strings.Split(signedFields, ":") {
		field = strings.TrimSpace(field)
		val, ok2 := hdrMap[strings.ToLower(field)]
		if !ok2 {
			continue
		}
		sb.WriteString(canonicalizeHeaderRelaxed(field, val))
		sb.WriteString("\r\n")
	}
	// Append DKIM-Signature with b= stripped (zeroed), canonicalized.
	cleanedSig := stripBValue(sigVal)
	sb.WriteString(canonicalizeHeaderRelaxed("dkim-signature", cleanedSig))

	dataToVerify := []byte(sb.String())
	hash := sha256.Sum256(dataToVerify)

	sigBytes, decErr := base64.StdEncoding.DecodeString(
		strings.Map(func(r rune) rune {
			if r == ' ' || r == '\t' || r == '\r' || r == '\n' {
				return -1
			}
			return r
		}, b64),
	)
	if decErr != nil {
		return "", fmt.Errorf("dkim: decode signature base64: %w", decErr)
	}

	// Try RSA PKIX public key first.
	if pubKey, rsaErr := x509.ParsePKIXPublicKey(pubDER); rsaErr == nil {
		switch k := pubKey.(type) {
		case *rsa.PublicKey:
			if err := rsa.VerifyPKCS1v15(k, gocrypto.SHA256, hash[:], sigBytes); err != nil {
				return "", fmt.Errorf("dkim: RSA signature invalid: %w", err)
			}
			return dom, nil
		case ed25519.PublicKey:
			if !ed25519.Verify(k, hash[:], sigBytes) {
				return "", fmt.Errorf("dkim: Ed25519 signature invalid")
			}
			return dom, nil
		default:
			return "", fmt.Errorf("dkim: unsupported public key type %T", pubKey)
		}
	}

	// Try raw Ed25519 public key (32 bytes).
	if len(pubDER) == ed25519.PublicKeySize {
		edPub := ed25519.PublicKey(pubDER)
		if !ed25519.Verify(edPub, hash[:], sigBytes) {
			return "", fmt.Errorf("dkim: Ed25519 signature invalid")
		}
		return dom, nil
	}

	return "", fmt.Errorf("dkim: unable to parse public key from DNS record")
}

// DNSRecord returns the DKIM TXT record string for publishing.
func DNSRecord(selector, domain, publicKeyPEM string) string {
	block, _ := pem.Decode([]byte(publicKeyPEM))
	var p string
	if block != nil {
		p = base64.StdEncoding.EncodeToString(block.Bytes)
	}
	return fmt.Sprintf("%s._domainkey.%s. IN TXT \"v=DKIM1; k=rsa; p=%s\"", selector, domain, p)
}

// SPFRecord returns the recommended SPF TXT record for a domain.
func SPFRecord(domain string) string {
	return fmt.Sprintf("%s IN TXT \"v=spf1 mx a ~all\"", domain)
}

// ---- Internals ----

// splitMessage splits at the first blank line (CRLF or LF variants).
func splitMessage(msg []byte) (headers, body []byte) {
	s := string(msg)
	for _, sep := range []string{"\r\n\r\n", "\n\n"} {
		if idx := strings.Index(s, sep); idx >= 0 {
			return []byte(s[:idx]), []byte(s[idx+len(sep):])
		}
	}
	return msg, nil
}

// canonicalizeBodyRelaxed applies RFC 6376 relaxed body canonicalization.
func canonicalizeBodyRelaxed(body []byte) []byte {
	if len(body) == 0 {
		return []byte("\r\n")
	}
	// Normalize line endings.
	s := strings.ReplaceAll(string(body), "\r\n", "\n")
	s = strings.ReplaceAll(s, "\r", "\n")

	lines := strings.Split(s, "\n")
	var out []string
	for _, line := range lines {
		// Collapse whitespace runs within each line, trim trailing whitespace.
		fields := strings.Fields(line)
		out = append(out, strings.Join(fields, " "))
	}

	// Strip trailing empty lines.
	for len(out) > 0 && out[len(out)-1] == "" {
		out = out[:len(out)-1]
	}

	result := strings.Join(out, "\r\n") + "\r\n"
	return []byte(result)
}

// canonicalizeHeaderRelaxed applies RFC 6376 relaxed header canonicalization
// to a single header name + value. Returns "lowername:value" (no trailing CRLF).
func canonicalizeHeaderRelaxed(name, value string) string {
	name = strings.ToLower(strings.TrimSpace(name))
	// Collapse all whitespace (including CRLF folding) to a single space.
	value = strings.ReplaceAll(value, "\r\n", " ")
	value = strings.ReplaceAll(value, "\r", " ")
	value = strings.ReplaceAll(value, "\n", " ")
	// Collapse multiple spaces.
	parts := strings.Fields(value)
	value = strings.Join(parts, " ")
	value = strings.TrimSpace(value)
	return name + ":" + value
}

// parseHeaders builds a map of lower-cased header name → last value.
// Handles multi-line (folded) headers.
func parseHeaders(headerBytes []byte) map[string]string {
	m := make(map[string]string)
	lines := strings.Split(strings.ReplaceAll(string(headerBytes), "\r\n", "\n"), "\n")

	var curName, curVal string
	flush := func() {
		if curName != "" {
			m[strings.ToLower(curName)] = curVal
		}
	}

	for _, line := range lines {
		if line == "" {
			continue
		}
		if len(line) > 0 && (line[0] == ' ' || line[0] == '\t') {
			// Folded continuation.
			curVal += " " + strings.TrimSpace(line)
			continue
		}
		flush()
		idx := strings.IndexByte(line, ':')
		if idx < 0 {
			curName = ""
			curVal = ""
			continue
		}
		curName = line[:idx]
		curVal = strings.TrimSpace(line[idx+1:])
	}
	flush()
	return m
}

// parseDKIMParams parses semicolon-separated tag=value pairs (DKIM and DNS TXT).
func parseDKIMParams(s string) map[string]string {
	m := make(map[string]string)
	for _, part := range strings.Split(s, ";") {
		part = strings.TrimSpace(part)
		if part == "" {
			continue
		}
		idx := strings.IndexByte(part, '=')
		if idx < 0 {
			continue
		}
		key := strings.TrimSpace(part[:idx])
		val := strings.TrimSpace(part[idx+1:])
		m[key] = val
	}
	return m
}

// stripBValue removes the value of the b= tag (sets it to empty) so the
// header can be re-canonicalized for verification.
func stripBValue(sigVal string) string {
	// Find b= and zero everything after it up to the next ;.
	parts := strings.Split(sigVal, ";")
	for i, p := range parts {
		trimmed := strings.TrimSpace(p)
		if strings.HasPrefix(trimmed, "b=") {
			parts[i] = " b="
		}
	}
	return strings.Join(parts, ";")
}