package dkim

import (
	"context"
	"crypto"
	"crypto/ed25519"
	"crypto/rand"
	"crypto/rsa"
	"crypto/x509"
	"encoding/base64"
	"encoding/pem"
	"fmt"
	"io"
	"log"
	"strings"

	"postern/internal/model"
)

type privateKey interface {
	Public() crypto.PublicKey
}

func genRSAPrivateKey() (privateKey, model.DKIMAlgorithmType) {
	nBits := 2048
	log.Printf("Generating a %v-bit RSA key", nBits)
	privKey, err := rsa.GenerateKey(rand.Reader, nBits)

	if err != nil {
		log.Fatalf("Failed to generate key: %v", err)
	}
	return privKey, model.DKIMAlgorithmRSA2048
}

func genEd25519PrivateKey() (privateKey, model.DKIMAlgorithmType) {
	_, privKey, err := ed25519.GenerateKey(rand.Reader)
	if err != nil {
		log.Fatalf("Failed to generate key: %v", err)
	}
	return privKey, model.DKIMAlgorithmEd25519
}

func encodePrivateKey(privateKey privateKey) []byte {
	privBytes, err := x509.MarshalPKCS8PrivateKey(privateKey)
	if err != nil {
		log.Fatalf("Failed to marshal private key: %v", err)
	}

	privBlock := pem.Block{
		Type:  "PRIVATE KEY",
		Bytes: privBytes,
	}

	encoded := pem.EncodeToMemory(&privBlock)
	if encoded == nil {
		log.Fatalf("Failed to write key PEM block: %v", err)
	}
	return encoded
}

func decodePrivateKey(encoded string) (privateKey, error) {
	// Decode PEM block from string
	block, _ := pem.Decode([]byte(encoded))
	if block == nil {
		return nil, fmt.Errorf("failed to decode PEM block: no valid PEM data found")
	}

	// Validate PEM type
	if block.Type != "PRIVATE KEY" {
		return nil, fmt.Errorf("invalid PEM type: expected 'PRIVATE KEY', got '%s'", block.Type)
	}

	// Parse PKCS#8 private key
	key, err := x509.ParsePKCS8PrivateKey(block.Bytes)
	if err != nil {
		return nil, fmt.Errorf("failed to parse private key: %w", err)
	}

	switch k := key.(type) {
	case *rsa.PrivateKey:
		return k, nil
	case ed25519.PrivateKey:
		return k, nil
	default:
		return nil, fmt.Errorf("unsupported private key type: %T", k)
	}
}

func pubKeyRecord(pubKey crypto.PublicKey) string {
	var pubBytes []byte
	var keyType string
	switch p := pubKey.(type) {
	case *rsa.PublicKey:
		keyType = "rsa"
		// RFC 6376 is inconsistent about whether RSA public keys should
		// be formatted as RSAPublicKey or SubjectPublicKeyInfo.
		// Erratum 3017 (https://www.rfc-editor.org/errata/eid3017)
		// proposes allowing both.  We use SubjectPublicKeyInfo for
		// consistency with other implementations including opendkim,
		// Gmail, and Fastmail.
		var err error
		pubBytes, err = x509.MarshalPKIXPublicKey(p)
		if err != nil {
			log.Fatalf("Failed to marshal public key: %v", err)
		}
	case ed25519.PublicKey:
		keyType = "ed25519"
		pubBytes = p
	default:
		panic("unreachable")
	}

	params := []string{
		"v=DKIM1",
		"k=" + keyType,
		"p=" + base64.StdEncoding.EncodeToString(pubBytes),
	}
	return strings.Join(params, "; ")
}

func (d *DKIM) retrievePrivateKey(ctx context.Context, selector string, domain string) (privateKey, string, error) {
	var dbRecord model.DKIM
	var err error
	if selector == "" {
		dbRecord, err = d.db.GetDKIMForDomain(ctx, domain)
	} else {
		dbRecord, err = d.db.GetDKIMForSelectorDomain(ctx, selector, domain)
	}
	if err != nil {
		return nil, "", err
	}

	pk, err := d.persistence.BlobReader(dbRecord.BlobAddress)
	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
	}

	privKey, err := decodePrivateKey(string(privKeyBytes))
	return privKey, dbRecord.Selector, err
}