// Package dkim manages per-domain DKIM keys and signs outbound mail, mirroring
// email_server/dkim_manager.py.
package dkim

import (
	"crypto"
	"crypto/rand"
	"crypto/rsa"
	"crypto/x509"
	"encoding/base64"
	"encoding/pem"
	"fmt"
	"math/big"
	"strings"
	"time"

	msgdkim "github.com/emersion/go-msgauth/dkim"
	"mailgoserver/internal/db"
)

// FixedHeaders is the exact 8-header list DKIM signs over, in this fixed order,
// mirroring dkim_manager.sign_email's `headers` list.
var FixedHeaders = []string{
	"from", "to", "subject", "date", "message-id", "mime-version", "content-type", "content-transfer-encoding",
}

const selectorChars = "abcdefghijklmnopqrstuvwxyz0123456789"

// GenerateSelector mirrors DKIMManager._generate_random_selector(length=12).
func GenerateSelector() string {
	b := make([]byte, 12)
	max := big.NewInt(int64(len(selectorChars)))
	for i := range b {
		n, _ := rand.Int(rand.Reader, max)
		b[i] = selectorChars[n.Int64()]
	}
	return string(b)
}

// Manager mirrors DKIMManager, keyed to a DB handle.
type Manager struct {
	DB      *db.DB
	KeySize int
}

func New(database *db.DB, keySize int) *Manager {
	if keySize == 0 {
		keySize = 2048
	}
	return &Manager{DB: database, KeySize: keySize}
}

// GenerateDKIMKeypair mirrors DKIMManager.generate_dkim_keypair. Returns false if the
// domain doesn't exist (looked up by exact name, active or not — matching the Python
// query, which has no is_active filter here).
func (m *Manager) GenerateDKIMKeypair(domainName, selector string, forceNewKey bool) (bool, error) {
	dom, err := m.DB.GetDomainByNameExact(domainName)
	if err != nil {
		return false, err
	}
	if dom == nil {
		return false, nil
	}

	now := time.Now()
	if _, err := m.DB.Exec(`UPDATE esrv_dkim_keys SET is_active = 0, replaced_at = ? WHERE domain_id = ? AND is_active = 1`, now, dom.ID); err != nil {
		return false, err
	}

	if selector == "" {
		selector = GenerateSelector()
	}

	if !forceNewKey {
		existing, err := m.DB.GetDKIMKeyByDomainAndSelector(dom.ID, selector)
		if err != nil {
			return false, err
		}
		if existing != nil {
			if _, err := m.DB.Exec(`UPDATE esrv_dkim_keys SET is_active = 1, replaced_at = NULL WHERE id = ?`, existing.ID); err != nil {
				return false, err
			}
			return true, nil
		}
	}

	priv, err := rsa.GenerateKey(rand.Reader, m.KeySize)
	if err != nil {
		return false, err
	}
	privPEM, pubPEM, err := encodeKeyPair(priv)
	if err != nil {
		return false, err
	}

	if _, err := m.DB.Exec(`INSERT INTO esrv_dkim_keys (domain_id, selector, private_key, public_key, is_active, created_at)
		VALUES (?, ?, ?, ?, 1, ?)`, dom.ID, selector, privPEM, pubPEM, now); err != nil {
		return false, err
	}
	return true, nil
}

func encodeKeyPair(priv *rsa.PrivateKey) (privPEM, pubPEM string, err error) {
	privBytes, err := x509.MarshalPKCS8PrivateKey(priv)
	if err != nil {
		return "", "", err
	}
	privPEM = string(pem.EncodeToMemory(&pem.Block{Type: "PRIVATE KEY", Bytes: privBytes}))

	pubBytes, err := x509.MarshalPKIXPublicKey(&priv.PublicKey)
	if err != nil {
		return "", "", err
	}
	pubPEM = string(pem.EncodeToMemory(&pem.Block{Type: "PUBLIC KEY", Bytes: pubBytes}))
	return privPEM, pubPEM, nil
}

// GetActiveDKIMKey mirrors DKIMManager.get_active_dkim_key.
func (m *Manager) GetActiveDKIMKey(domainName string) (*db.DKIMKey, error) {
	dom, err := m.DB.GetDomainByName(domainName)
	if err != nil || dom == nil {
		return nil, err
	}
	return m.DB.GetActiveDKIMKeyByDomainID(dom.ID)
}

// DNSRecord is the DNS TXT record for a domain's active DKIM key, mirroring
// DKIMManager.get_dkim_public_key_record's return shape.
type DNSRecord struct {
	Name  string
	Type  string
	Value string
}

// GetDKIMPublicKeyRecord mirrors DKIMManager.get_dkim_public_key_record.
func (m *Manager) GetDKIMPublicKeyRecord(domainName string) (*DNSRecord, error) {
	key, err := m.GetActiveDKIMKey(domainName)
	if err != nil || key == nil {
		return nil, err
	}
	raw := rawBase64FromPEM(key.PublicKey)
	return &DNSRecord{
		Name:  fmt.Sprintf("%s._domainkey.%s", key.Selector, domainName),
		Type:  "TXT",
		Value: fmt.Sprintf(`"v=DKIM1; k=rsa; p=%s"`, raw),
	}, nil
}

func rawBase64FromPEM(pemStr string) string {
	block, _ := pem.Decode([]byte(pemStr))
	if block == nil {
		return ""
	}
	return base64.StdEncoding.EncodeToString(block.Bytes)
}

// Sign mirrors DKIMManager.sign_email: strips any existing DKIM-Signature header,
// signs over the fixed 8-header list with relaxed/relaxed canonicalization, and
// returns the original content unmodified on any failure (including "no active key").
func (m *Manager) Sign(content, domainName string) string {
	key, err := m.GetActiveDKIMKey(domainName)
	if err != nil || key == nil {
		return content
	}
	block, _ := pem.Decode([]byte(key.PrivateKey))
	if block == nil {
		return content
	}
	privAny, err := x509.ParsePKCS8PrivateKey(block.Bytes)
	if err != nil {
		return content
	}
	priv, ok := privAny.(*rsa.PrivateKey)
	if !ok {
		return content
	}

	stripped := stripExistingSignature(content)

	var out strings.Builder
	err = msgdkim.Sign(&out, strings.NewReader(stripped), &msgdkim.SignOptions{
		Domain:                 domainName,
		Selector:               key.Selector,
		Signer:                 priv,
		Hash:                   crypto.SHA256,
		HeaderCanonicalization: msgdkim.CanonicalizationRelaxed,
		BodyCanonicalization:   msgdkim.CanonicalizationRelaxed,
		HeaderKeys:             FixedHeaders,
	})
	if err != nil {
		return content
	}
	return out.String()
}

// stripExistingSignature removes a pre-existing DKIM-Signature header (including any
// folded continuation lines), mirroring the regex in dkim_manager.sign_email.
func stripExistingSignature(content string) string {
	lines := strings.Split(content, "\n")
	var out []string
	skipping := false
	for _, line := range lines {
		lower := strings.ToLower(line)
		if !skipping && strings.HasPrefix(lower, "dkim-signature:") {
			skipping = true
			continue
		}
		if skipping {
			if len(line) > 0 && (line[0] == ' ' || line[0] == '\t') {
				continue // folded continuation line
			}
			skipping = false
		}
		out = append(out, line)
	}
	return strings.Join(out, "\n")
}

// VerifyInbound checks whether content carries at least one valid DKIM signature
// aligned to domainName (typically the From-header domain), mirroring Sign's
// error-swallowing-to-false style. Used only as one signal feeding the inbound spam
// heuristic (internal/mailstore) — never blocks delivery by itself.
func VerifyInbound(content, domainName string) bool {
	verifications, err := msgdkim.Verify(strings.NewReader(content))
	if err != nil {
		return false
	}
	for _, v := range verifications {
		if v.Err == nil && strings.EqualFold(v.Domain, domainName) {
			return true
		}
	}
	return false
}

// GetActiveCustomHeaders mirrors DKIMManager.get_active_custom_headers.
func (m *Manager) GetActiveCustomHeaders(domainName string) ([][2]string, error) {
	dom, err := m.DB.GetDomainByName(domainName)
	if err != nil || dom == nil {
		return nil, err
	}
	rows, err := m.DB.Query(`SELECT header_name, header_value FROM esrv_custom_headers WHERE domain_id = ? AND is_active = 1`, dom.ID)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	var out [][2]string
	for rows.Next() {
		var name, value string
		if err := rows.Scan(&name, &value); err != nil {
			return nil, err
		}
		out = append(out, [2]string{name, value})
	}
	return out, rows.Err()
}