package dkim

import (
	"bytes"
	"crypto"
	"crypto/rand"
	"crypto/rsa"
	"crypto/sha256"
	"encoding/base64"
	"fmt"
	"regexp"
	"strings"
	"time"
)

// signedHeaders is the fixed set of headers we sign, in order, when present.
// Keeping this list small and stable avoids the classic DKIM pitfall of
// signing headers that get legitimately rewritten in transit (Received, etc).
var signedHeaders = []string{"from", "to", "subject", "date", "message-id"}

// Sign adds a DKIM-Signature header to raw using relaxed/relaxed
// canonicalization and RSA-SHA256, per RFC 6376. Returns the message with
// the DKIM-Signature header prepended.
func Sign(privateKeyPEM []byte, domain, selector string, raw []byte) ([]byte, error) {
	key, err := ParsePrivateKey(privateKeyPEM)
	if err != nil {
		return nil, err
	}

	headers, body := splitMessage(raw)
	bodyCanon := canonicalizeBodyRelaxed(body)
	bodyHash := sha256.Sum256(bodyCanon)
	bodyHashB64 := base64.StdEncoding.EncodeToString(bodyHash[:])

	headerMap := parseHeaders(headers)

	var presentSigned []string
	for _, h := range signedHeaders {
		if _, ok := headerMap[h]; ok {
			presentSigned = append(presentSigned, h)
		}
	}
	if len(presentSigned) == 0 {
		return nil, fmt.Errorf("no signable headers present in message")
	}

	// Build the DKIM-Signature header with an empty b= tag first — this
	// unsigned version is itself included (relaxed-canonicalized) in what we
	// sign, per RFC 6376 §3.7.
	dkimHeaderTemplate := buildDKIMHeader(domain, selector, presentSigned, bodyHashB64, "")

	signInput := canonicalizeHeadersRelaxed(headerMap, presentSigned)
	signInput = append(signInput, canonicalizeHeaderRelaxed("dkim-signature", dkimHeaderTemplate)...)
	// Per spec, the DKIM-Signature header itself is canonicalized WITHOUT a
	// trailing CRLF when it's the last (signed) header being hashed.
	signInput = bytes.TrimSuffix(signInput, []byte("\r\n"))

	hashed := sha256.Sum256(signInput)
	signature, err := rsa.SignPKCS1v15(rand.Reader, key, crypto.SHA256, hashed[:])
	if err != nil {
		return nil, fmt.Errorf("signing: %w", err)
	}
	sigB64 := base64.StdEncoding.EncodeToString(signature)

	finalHeader := buildDKIMHeader(domain, selector, presentSigned, bodyHashB64, sigB64)

	var out bytes.Buffer
	out.WriteString("DKIM-Signature: ")
	out.WriteString(finalHeader)
	out.WriteString("\r\n")
	out.Write(headers)
	out.Write(body)

	return out.Bytes(), nil
}

func buildDKIMHeader(domain, selector string, signedHdrs []string, bodyHashB64, sigB64 string) string {
	return fmt.Sprintf(
		"v=1; a=rsa-sha256; c=relaxed/relaxed; d=%s; s=%s; t=%d; h=%s; bh=%s; b=%s",
		domain, selector, time.Now().Unix(), strings.Join(signedHdrs, ":"), bodyHashB64, sigB64,
	)
}

// splitMessage separates the raw RFC 5322 message into its header block
// (including the trailing blank line's CRLF) and body.
func splitMessage(raw []byte) (headers, body []byte) {
	sep := []byte("\r\n\r\n")
	idx := bytes.Index(raw, sep)
	if idx == -1 {
		// Tolerate bare-LF input (shouldn't happen from our own DATA reader,
		// which always produces CRLF, but be defensive).
		sep = []byte("\n\n")
		idx = bytes.Index(raw, sep)
		if idx == -1 {
			return raw, nil
		}
	}
	return raw[:idx+len(sep)], raw[idx+len(sep):]
}

// parseHeaders builds a lowercase-name -> raw-value-with-original-case map,
// unfolding continuation lines (RFC 5322 §2.2.3).
func parseHeaders(headerBlock []byte) map[string]string {
	result := map[string]string{}
	lines := strings.Split(string(headerBlock), "\r\n")

	var currentName, currentValue string
	flush := func() {
		if currentName != "" {
			result[strings.ToLower(currentName)] = currentValue
		}
	}

	for _, line := range lines {
		if line == "" {
			continue
		}
		if (strings.HasPrefix(line, " ") || strings.HasPrefix(line, "\t")) && currentName != "" {
			currentValue += " " + strings.TrimSpace(line)
			continue
		}
		flush()
		name, value, found := strings.Cut(line, ":")
		if !found {
			currentName = ""
			continue
		}
		currentName = strings.TrimSpace(name)
		currentValue = strings.TrimSpace(value)
	}
	flush()
	return result
}

// canonicalizeHeadersRelaxed builds the signed-header block per RFC 6376
// §3.4.2: lowercase header name, unfold, collapse WSP runs to single space,
// trim trailing WSP on the value, each header terminated with CRLF, in the
// exact order listed by names.
func canonicalizeHeadersRelaxed(headerMap map[string]string, names []string) []byte {
	var buf bytes.Buffer
	for _, name := range names {
		value, ok := headerMap[name]
		if !ok {
			continue
		}
		buf.Write(canonicalizeHeaderRelaxed(name, value))
	}
	return buf.Bytes()
}

func canonicalizeHeaderRelaxed(name, value string) []byte {
	name = strings.ToLower(strings.TrimSpace(name))
	value = collapseWSP(strings.TrimSpace(value))
	return []byte(name + ":" + value + "\r\n")
}

var wspRunRE = regexp.MustCompile(`[ \t]+`)

func collapseWSP(s string) string {
	return wspRunRE.ReplaceAllString(s, " ")
}

// canonicalizeBodyRelaxed implements RFC 6376 §3.4.4: reduce WSP sequences
// within a line to a single space, remove trailing WSP from each line,
// remove trailing empty lines (but keep exactly one CRLF if the body is
// non-empty after trimming).
func canonicalizeBodyRelaxed(body []byte) []byte {
	if len(body) == 0 {
		return []byte("")
	}

	lines := bytes.Split(body, []byte("\r\n"))
	for i, line := range lines {
		line = wspRunRE.ReplaceAll(line, []byte(" "))
		lines[i] = bytes.TrimRight(line, " \t")
	}

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

	if len(lines) == 0 {
		return []byte("")
	}

	var buf bytes.Buffer
	for _, line := range lines {
		buf.Write(line)
		buf.WriteString("\r\n")
	}
	return buf.Bytes()
}