package dkim

import (
	"bytes"
	"context"
	"crypto"
	"fmt"
	"io"
	"log/slog"
	"net/mail"
	"strings"
	"time"

	"github.com/emersion/go-msgauth/dkim"
)

func rfc6376Section541() []string {
	return []string{
		"From",
		"Reply-To",
		"Subject",
		"Date",
		"To",
		"Cc",
		"Resent-Date",
		"Resent-From",
		"Resent-To",
		"Resent-Cc",
		"In-Reply-To",
		"References",
		"List-Id",
		"List-Help",
		"List-Unsubscribe",
		"List-Subscribe",
		"List-Post",
		"List-Owner",
		"List-Archive",
	}
}

// Sign signs the message with the first enabled selector for the From: domain and returns the signed message.
func (d *DKIM) Sign(ctx context.Context, message []byte) ([]byte, error) {
	// TODO check if we can make this a zero-copy operation with io.Reader instead of message []byte
	// for DMARC alignment, we first need to find the From: domain
	msg, err := mail.ReadMessage(bytes.NewReader(message))
	if err != nil {
		return nil, fmt.Errorf("parsing message: %w", err)
	}

	fromHeader := msg.Header.Get("From")
	if fromHeader == "" {
		return nil, fmt.Errorf("from header is empty")
	}
	parsedFromHeader, err := mail.ParseAddress(fromHeader)
	if err != nil {
		return nil, fmt.Errorf("parsing from header: %w", err)
	}
	i := strings.LastIndex(parsedFromHeader.Address, "@")
	if i == -1 {
		return nil, fmt.Errorf("from header does not contain @")
	}
	domain := parsedFromHeader.Address[i+1:]

	slog.Info("[DKIM] signing based on From: header", "domain", domain, "from", fromHeader)

	dRec, err := d.db.GetDKIMForDomain(ctx, domain)
	if err != nil {
		return nil, fmt.Errorf("get DKIM for domain: %w", err)
	}

	signer, err := d.readPrivateKey(dRec.BlobAddress)
	if err != nil {
		return nil, fmt.Errorf("reading signing key: %w", err)
	}

	var buffer bytes.Buffer
	_, err = io.Copy(&buffer, bytes.NewReader(message))
	if err != nil {
		return nil, fmt.Errorf("copy buffer: %w", err)
	}

	err = dkim.Sign(&buffer, &buffer, &dkim.SignOptions{
		Domain:                 domain,
		Selector:               dRec.Selector,
		Identifier:             "",
		Signer:                 signer.(crypto.Signer),
		Hash:                   crypto.SHA256,
		HeaderCanonicalization: dkim.CanonicalizationSimple,
		BodyCanonicalization:   dkim.CanonicalizationSimple,
		HeaderKeys:             rfc6376Section541(),
		Expiration:             time.Time{},
		QueryMethods:           nil,
	})
	if err != nil {
		return nil, fmt.Errorf("creating signer: %w", err)
	}

	return buffer.Bytes(), nil
}

// readPrivateKey loads and decodes the private key blob referenced by nonceIdentifier.
func (d *DKIM) readPrivateKey(nonceIdentifier string) (privateKey, error) {
	pk, err := d.persistence.BlobReader(nonceIdentifier)
	if err != nil {
		return nil, err
	}
	defer func(privKey io.ReadCloser) {
		_ = privKey.Close()
	}(pk)

	privKeyBytes, err := io.ReadAll(pk)
	if err != nil {
		return nil, err
	}

	return decodePrivateKey(string(privKeyBytes))
}