package dkim

import (
	"crypto"
	"crypto/rsa"
	"crypto/sha256"
	"crypto/x509"
	"encoding/base64"
	"fmt"
	"strings"
)

// Verify checks a signed message's DKIM-Signature header against the given
// public key (as would be fetched from DNS in Phase 4's inbound pipeline).
// This lean version only handles rsa-sha256/relaxed-relaxed — the exact
// profile Sign() produces — since its purpose here is to prove the signer is
// correct. Phase 4 will build a fuller verifier (multiple algorithms,
// simple/simple and mixed canonicalization) for arbitrary inbound mail.
func Verify(publicKeyDER []byte, raw []byte) error {
	headers, body := splitMessage(raw)
	headerMap := parseHeaders(headers)

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

	tags := parseDKIMTags(dkimHeaderValue)
	if tags["a"] != "rsa-sha256" {
		return fmt.Errorf("unsupported algorithm: %s", tags["a"])
	}
	if tags["c"] != "relaxed/relaxed" {
		return fmt.Errorf("unsupported canonicalization: %s", tags["c"])
	}

	// Verify body hash.
	bodyCanon := canonicalizeBodyRelaxed(body)
	bodyHash := sha256.Sum256(bodyCanon)
	expectedBH := base64.StdEncoding.EncodeToString(bodyHash[:])
	if tags["bh"] != expectedBH {
		return fmt.Errorf("body hash mismatch: signature claims %s, computed %s", tags["bh"], expectedBH)
	}

	signedHdrNames := strings.Split(tags["h"], ":")

	// Rebuild the exact signing input: canonicalized signed headers, then the
	// DKIM-Signature header itself with b= emptied, no trailing CRLF.
	signInput := canonicalizeHeadersRelaxed(headerMap, signedHdrNames)

	dkimHeaderNoB := replaceDKIMTag(dkimHeaderValue, "b", "")
	signInput = append(signInput, canonicalizeHeaderRelaxed("dkim-signature", dkimHeaderNoB)...)
	signInput = trimTrailingCRLF(signInput)

	sigBytes, err := base64.StdEncoding.DecodeString(tags["b"])
	if err != nil {
		return fmt.Errorf("decoding signature: %w", err)
	}

	pubAny, err := x509.ParsePKIXPublicKey(publicKeyDER)
	if err != nil {
		return fmt.Errorf("parsing public key: %w", err)
	}
	pubKey, ok := pubAny.(*rsa.PublicKey)
	if !ok {
		return fmt.Errorf("public key is not RSA")
	}

	hashed := sha256.Sum256(signInput)
	if err := rsa.VerifyPKCS1v15(pubKey, crypto.SHA256, hashed[:], sigBytes); err != nil {
		return fmt.Errorf("signature verification failed: %w", err)
	}

	return nil
}

func parseDKIMTags(header string) map[string]string {
	tags := map[string]string{}
	for _, part := range strings.Split(header, ";") {
		part = strings.TrimSpace(part)
		if part == "" {
			continue
		}
		name, value, found := strings.Cut(part, "=")
		if !found {
			continue
		}
		tags[strings.TrimSpace(name)] = strings.TrimSpace(value)
	}
	return tags
}

func replaceDKIMTag(header, tag, newValue string) string {
	parts := strings.Split(header, ";")
	for i, part := range parts {
		trimmed := strings.TrimSpace(part)
		if strings.HasPrefix(trimmed, tag+"=") {
			parts[i] = " " + tag + "=" + newValue
		}
	}
	return strings.Join(parts, ";")
}

func trimTrailingCRLF(b []byte) []byte {
	for len(b) >= 2 && b[len(b)-2] == '\r' && b[len(b)-1] == '\n' {
		return b[:len(b)-2]
	}
	return b
}