package parsemail

import (
	"bytes"
	"crypto"
	"crypto/rsa"
	"crypto/x509"
	"encoding/base64"
	"encoding/pem"
	"fmt"
	"git.pcux.cn/xuans/Smail/config"
	"git.pcux.cn/xuans/Smail/utils/context"
	securefile "git.pcux.cn/xuans/Smail/utils/file"
	"git.pcux.cn/xuans/Smail/utils/runtimepath"
	"github.com/emersion/go-msgauth/dkim"
	log "github.com/sirupsen/logrus"
	"io"
	"net"
	"os"
	"strings"
)

type Dkim struct {
	privateKey crypto.Signer
}

var instance *Dkim

func Init() error {
	cfg := config.CurrentConfig()
	if cfg == nil {
		return fmt.Errorf("runtime configuration is unavailable")
	}
	privateKey, err := loadPrivateKey(cfg.DkimPrivateKeyPath)
	if err != nil {
		return fmt.Errorf("load DKIM private key: %w", err)
	}

	instance = &Dkim{
		privateKey: privateKey,
	}
	return nil
}

func loadPrivateKey(path string) (crypto.Signer, error) {
	basePath := runtimepath.DKIMDirectory(config.ROOT_PATH)
	privateName, err := securefile.ControlledRelativePath(basePath, path)
	if err != nil {
		return nil, err
	}
	root, err := os.OpenRoot(basePath)
	if err != nil {
		return nil, err
	}
	defer root.Close()
	info, err := root.Lstat(privateName)
	if err != nil {
		return nil, err
	}
	if info.Mode()&os.ModeSymlink != 0 || !info.Mode().IsRegular() || info.Size() <= 0 || info.Size() > 64<<10 {
		return nil, fmt.Errorf("invalid DKIM private key file")
	}
	data, err := root.ReadFile(privateName)
	if err != nil {
		return nil, err
	}
	block, rest := pem.Decode(data)
	if block == nil || len(bytes.TrimSpace(rest)) != 0 {
		return nil, fmt.Errorf("invalid DKIM private key PEM")
	}

	var privateKey *rsa.PrivateKey
	switch strings.ToUpper(block.Type) {
	case "PRIVATE KEY":
		parsed, parseErr := x509.ParsePKCS8PrivateKey(block.Bytes)
		if parseErr != nil {
			return nil, parseErr
		}
		var ok bool
		privateKey, ok = parsed.(*rsa.PrivateKey)
		if !ok {
			return nil, fmt.Errorf("unsupported DKIM private key type")
		}
	case "RSA PRIVATE KEY":
		privateKey, err = x509.ParsePKCS1PrivateKey(block.Bytes)
		if err != nil {
			return nil, err
		}
	default:
		return nil, fmt.Errorf("unsupported DKIM private key type: %q", block.Type)
	}
	if privateKey.N == nil || privateKey.N.BitLen() < 2048 {
		return nil, fmt.Errorf("DKIM RSA private key is too weak")
	}
	if err := privateKey.Validate(); err != nil {
		return nil, err
	}
	return privateKey, nil
}
func (p *Dkim) Sign(msgData string) ([]byte, error) {
	if p == nil || p.privateKey == nil {
		return nil, fmt.Errorf("DKIM signer is not initialized")
	}
	domains := config.CurrentMailDomains()
	if strings.TrimSpace(domains.Primary) == "" {
		return nil, fmt.Errorf("DKIM domain is not configured")
	}

	var b bytes.Buffer
	r := strings.NewReader(msgData)

	options := &dkim.SignOptions{
		Domain:   domains.Primary,
		Selector: "default",
		Signer:   p.privateKey,
	}

	if err := dkim.Sign(&b, r, options); err != nil {
		return nil, err
	}
	return b.Bytes(), nil
}

func Check(ctx *context.Context, mail io.Reader) bool {
	return CheckForDomain(ctx, mail, "")
}

// CheckLocalDelivery 仅对回环投递使用当前主域公钥验签，不能仅凭来源地址放行。
func CheckLocalDelivery(ctx *context.Context, mail io.Reader, domain string, remote net.Addr) bool {
	if instance == nil {
		return false
	}
	return checkLocalDelivery(ctx, mail, domain, config.CurrentMailDomains().Primary, remote, instance.privateKey)
}

func checkLocalDelivery(ctx *context.Context, mail io.Reader, domain, primary string, remote net.Addr, signer crypto.Signer) bool {
	address, ok := remote.(*net.TCPAddr)
	if !ok || address == nil || !address.IP.IsLoopback() || signer == nil {
		return false
	}
	domain = strings.TrimSpace(domain)
	primary = strings.TrimSpace(primary)
	if primary == "" || !strings.EqualFold(domain, primary) {
		return false
	}
	publicKey, err := x509.MarshalPKIXPublicKey(signer.Public())
	if err != nil {
		return false
	}
	// 只提供本机实际使用的选择器和域名公钥，其他签名交回正常 DNS 验证流程。
	options := &dkim.VerifyOptions{LookupTXT: func(name string) ([]string, error) {
		if !strings.EqualFold(name, "default._domainkey."+primary) {
			return nil, fmt.Errorf("signature does not use the current local DKIM selector")
		}
		return []string{"v=DKIM1; k=rsa; p=" + base64.StdEncoding.EncodeToString(publicKey)}, nil
	}}
	verified := checkForDomain(ctx, mail, domain, options)
	if verified {
		log.WithContext(ctx).Info("DKIM local delivery signature verified")
	}
	return verified
}

// CheckForDomain 在指定域名时要求有效签名与发件域严格一致，避免其他域签名冒用本域。
func CheckForDomain(ctx *context.Context, mail io.Reader, domain string) bool {
	return checkForDomain(ctx, mail, domain, nil)
}

func checkForDomain(ctx *context.Context, mail io.Reader, domain string, options *dkim.VerifyOptions) bool {
	verifications, err := dkim.VerifyWithOptions(mail, options)
	if err != nil {
		log.WithContext(ctx).Warnf("DKIM Error:%v", err)
		return false
	}
	domain = strings.TrimSpace(domain)
	for _, v := range verifications {
		if v == nil {
			continue
		}
		if v.Err == nil {
			log.Println("Valid signature for:", v.Domain)
			if domain == "" || strings.EqualFold(v.Domain, domain) {
				return true
			}
		} else {
			log.Println("Invalid signature for:", v.Domain, v.Err)
		}
	}
	return false
}