package jwe

import (
	"crypto/ecdh"
	"crypto/ecdsa"
	"crypto/rsa"
	"fmt"

	"github.com/lestrrat-go/jwx/v3/internal/keyconv"
	"github.com/lestrrat-go/jwx/v3/jwa"
	"github.com/lestrrat-go/jwx/v3/jwe/internal/keygen"
	"github.com/lestrrat-go/jwx/v3/jwe/jwebb"
)

// encrypter is responsible for taking various components to encrypt a key.
// its operation is not concurrency safe. You must provide locking yourself
//
//nolint:govet
type encrypter struct {
	apu        []byte
	apv        []byte
	ctalg      jwa.ContentEncryptionAlgorithm
	keyalg     jwa.KeyEncryptionAlgorithm
	pubkey     any
	rawKey     any
	pbes2Count int
}

// newEncrypter creates a new Encrypter instance with all required parameters.
// The content cipher is built internally during construction.
//
// pubkey must be a public key in its "raw" format (i.e. something like
// *rsa.PublicKey, instead of jwk.Key)
//
// You should consider this object immutable once created.
func newEncrypter(keyalg jwa.KeyEncryptionAlgorithm, ctalg jwa.ContentEncryptionAlgorithm, pubkey any, rawKey any, apu, apv []byte, pbes2Count int) *encrypter {
	return &encrypter{
		apu:        apu,
		apv:        apv,
		ctalg:      ctalg,
		keyalg:     keyalg,
		pubkey:     pubkey,
		rawKey:     rawKey,
		pbes2Count: pbes2Count,
	}
}

func (e *encrypter) EncryptKey(cek []byte) (keygen.ByteSource, error) {
	keyalgStr := e.keyalg.String()
	ctalgStr := e.ctalg.String()

	if ke, ok := e.pubkey.(KeyEncrypter); ok {
		encrypted, err := ke.EncryptKey(cek)
		if err != nil {
			return nil, err
		}
		return keygen.ByteKey(encrypted), nil
	}

	if jwebb.IsDirect(keyalgStr) {
		sharedkey, ok := e.rawKey.([]byte)
		if !ok {
			return nil, fmt.Errorf("encrypt key: []byte is required as the key for %s (got %T)", keyalgStr, e.rawKey)
		}
		return jwebb.KeyEncryptDirect(cek, keyalgStr, sharedkey)
	}

	if jwebb.IsPBES2(keyalgStr) {
		password, ok := e.rawKey.([]byte)
		if !ok {
			return nil, fmt.Errorf("encrypt key: []byte is required as the password for %s (got %T)", keyalgStr, e.rawKey)
		}
		return jwebb.KeyEncryptPBES2(cek, keyalgStr, password, e.pbes2Count)
	}

	if jwebb.IsAESGCMKW(keyalgStr) {
		sharedkey, ok := e.rawKey.([]byte)
		if !ok {
			return nil, fmt.Errorf("encrypt key: []byte is required as the key for %s (got %T)", keyalgStr, e.rawKey)
		}
		return jwebb.KeyEncryptAESGCMKW(cek, keyalgStr, sharedkey)
	}

	if jwebb.IsECDHES(keyalgStr) {
		_, keysize, keywrap, err := jwebb.KeyEncryptionECDHESKeySize(keyalgStr, ctalgStr)
		if err != nil {
			return nil, fmt.Errorf(`failed to determine ECDH-ES key size: %w`, err)
		}

		// Use rawKey for ECDH-ES operations - it should contain the actual key material
		keyToUse := e.rawKey
		if keyToUse == nil {
			keyToUse = e.pubkey
		}

		switch key := keyToUse.(type) {
		case *ecdsa.PublicKey:
			// no op
		case ecdsa.PublicKey:
			keyToUse = &key
		case *ecdsa.PrivateKey:
			keyToUse = &key.PublicKey
		case ecdsa.PrivateKey:
			keyToUse = &key.PublicKey
		case *ecdh.PublicKey:
			// no op
		case ecdh.PublicKey:
			keyToUse = &key
		case ecdh.PrivateKey:
			keyToUse = key.PublicKey()
		case *ecdh.PrivateKey:
			keyToUse = key.PublicKey()
		}

		// Determine key type and call appropriate function
		switch key := keyToUse.(type) {
		case *ecdh.PublicKey:
			if key.Curve() == ecdh.X25519() {
				if !keywrap {
					return jwebb.KeyEncryptECDHESX25519(cek, keyalgStr, e.apu, e.apv, key, keysize, ctalgStr)
				}
				return jwebb.KeyEncryptECDHESKeyWrapX25519(cek, keyalgStr, e.apu, e.apv, key, keysize, ctalgStr)
			}

			var ecdsaKey *ecdsa.PublicKey
			if err := keyconv.ECDHToECDSA(&ecdsaKey, key); err != nil {
				return nil, fmt.Errorf(`encrypt: failed to convert ECDH public key to ECDSA: %w`, err)
			}
			keyToUse = ecdsaKey
		}

		switch key := keyToUse.(type) {
		case *ecdsa.PublicKey:
			if !keywrap {
				return jwebb.KeyEncryptECDHESECDSA(cek, keyalgStr, e.apu, e.apv, key, keysize, ctalgStr)
			}
			return jwebb.KeyEncryptECDHESKeyWrapECDSA(cek, keyalgStr, e.apu, e.apv, key, keysize, ctalgStr)
		default:
			return nil, fmt.Errorf(`encrypt: unsupported key type for ECDH-ES: %T`, keyToUse)
		}
	}

	if jwebb.IsRSA15(keyalgStr) {
		keyToUse := e.rawKey
		if keyToUse == nil {
			keyToUse = e.pubkey
		}

		// Handle rsa.PublicKey by value - convert to pointer
		if pk, ok := keyToUse.(rsa.PublicKey); ok {
			keyToUse = &pk
		}

		var pubkey *rsa.PublicKey
		if err := keyconv.RSAPublicKey(&pubkey, keyToUse); err != nil {
			return nil, fmt.Errorf(`encrypt: failed to convert to RSA public key: %w`, err)
		}

		return jwebb.KeyEncryptRSA15(cek, keyalgStr, pubkey)
	}

	if jwebb.IsRSAOAEP(keyalgStr) {
		keyToUse := e.rawKey
		if keyToUse == nil {
			keyToUse = e.pubkey
		}

		// Handle rsa.PublicKey by value - convert to pointer
		if pk, ok := keyToUse.(rsa.PublicKey); ok {
			keyToUse = &pk
		}

		var pubkey *rsa.PublicKey
		if err := keyconv.RSAPublicKey(&pubkey, keyToUse); err != nil {
			return nil, fmt.Errorf(`encrypt: failed to convert to RSA public key: %w`, err)
		}

		return jwebb.KeyEncryptRSAOAEP(cek, keyalgStr, pubkey)
	}

	if jwebb.IsAESKW(keyalgStr) {
		sharedkey, ok := e.rawKey.([]byte)
		if !ok {
			return nil, fmt.Errorf("[]byte is required as the key to encrypt %s", keyalgStr)
		}
		return jwebb.KeyEncryptAESKW(cek, keyalgStr, sharedkey)
	}

	return nil, fmt.Errorf(`unsupported algorithm for key encryption (%s)`, keyalgStr)
}
