package secret

import (
	"bytes"
	"crypto/aes"
	"crypto/cipher"
	"crypto/ecdsa"
	"crypto/rsa"
	"crypto/sha256"
	"crypto/x509"
	"encoding/base64"
	"encoding/hex"
	"encoding/pem"
	"errors"
	"fmt"
	"hash"
	"os"
	"strings"

	"golang.org/x/crypto/chacha20poly1305"
	"golang.org/x/crypto/nacl/secretbox"
)

type aeadProvider struct {
	Provider
	aead cipher.AEAD
}

// WithChaCha20Poly1305 wraps the returned value and decrypts it using ChaCha20-Poly1305 AEAD.
//
// The used nonce size is 12 bytes.
func WithChaCha20Poly1305(p Provider, key []byte) (Provider, error) {
	cipher, err := chacha20poly1305.New(key)
	if err != nil {
		return nil, err
	}
	return aeadProvider{
		Provider: p,
		aead:     cipher,
	}, nil
}

// WithChaCha20Poly1305X wraps the returned value and decrypts it using ChaCha20-Poly1305 AEAD.
//
// The used nonce size is 24 bytes.
func WithChaCha20Poly1305X(p Provider, key []byte) (Provider, error) {
	cipher, err := chacha20poly1305.NewX(key)
	if err != nil {
		return nil, err
	}
	return aeadProvider{
		Provider: p,
		aead:     cipher,
	}, nil
}

// WithAESGCM wraps the returned value and decrypts it using AES-GCM AEAD.
//
// The accepted key sizes are 16 bytes for AES-128 and 32 bytes for AES-256.
func WithAESGCM(p Provider, key []byte) (Provider, error) {
	block, err := aes.NewCipher(key)
	if err != nil {
		return nil, err
	}

	gcm, err := cipher.NewGCM(block)
	if err != nil {
		return nil, err
	}

	return aeadProvider{
		Provider: p,
		aead:     gcm,
	}, nil
}

func (p aeadProvider) GetSecret(key string) (value []byte, err error) {
	if value, err = p.Provider.GetSecret(key); err != nil {
		return
	}

	nonceSize := p.aead.NonceSize()
	if len(value) < nonceSize {
		return nil, errors.New("secret: ciphertext is too short")
	}

	nonce, ciphertext := value[:nonceSize], value[nonceSize:]
	return p.aead.Open(nil, nonce, ciphertext, nil)
}

type oaepProvider struct {
	Provider
	key  *rsa.PrivateKey
	hash func() hash.Hash
}

// WithOAEP transparently decrypts secrets returned from the provider using RSA-OAEP.
//
// Encryption and decryption of a given message must use the same hash function and sha256.New()
// is a reasonable choice. When hashFunc is nil, then sha256 will be used.
func WithOAEP(p Provider, key *rsa.PrivateKey, hashFunc func() hash.Hash) Provider {
	if hashFunc == nil {
		hashFunc = func() hash.Hash {
			return sha256.New()
		}
	}
	return oaepProvider{
		Provider: p,
		key:      key,
		hash:     hashFunc,
	}
}

func (p oaepProvider) GetSecret(key string) (value []byte, err error) {
	if value, err = p.Provider.GetSecret(key); err != nil {
		return
	}
	return rsa.DecryptOAEP(p.hash(), nil, p.key, value, nil)
}

type pkcs1v15Provider struct {
	Provider
	key *rsa.PrivateKey
}

// WithPKCS1v15 transparently decrypts secrets returned from the provider using RSA and the padding scheme from PKCS #1 v1.5.
func WithPKCS1v15(p Provider, key *rsa.PrivateKey) Provider {
	return pkcs1v15Provider{
		Provider: p,
		key:      key,
	}
}

func (p pkcs1v15Provider) GetSecret(key string) (value []byte, err error) {
	if value, err = p.Provider.GetSecret(key); err != nil {
		return
	}
	return rsa.DecryptPKCS1v15(nil, p.key, value)
}

type secretboxProvider struct {
	Provider
	key [32]byte
}

// WithSecretBox transparently decrypts secrets returned from the provider using NaCL secretbox.
func WithSecretBox(p Provider, key [32]byte) Provider {
	return &secretboxProvider{
		Provider: p,
		key:      key,
	}
}

func (p secretboxProvider) GetSecret(key string) (value []byte, err error) {
	if value, err = p.Provider.GetSecret(key); err != nil {
		return
	}

	var nonce [24]byte
	if copy(nonce[:], value) < 24 {
		return nil, errors.New("secret: encrypted secretbox value is too short")
	}

	var ok bool
	if value, ok = secretbox.Open(nil, value[24:], &nonce, &p.key); !ok {
		return nil, errors.New("secret: encrypted secretbox open failed")
	}

	return
}

const (
	pemECPrivateKey  = "EC PRIVATE KEY"
	pemRSAPrivateKey = "RSA PRIVATE KEY"
	pemPrivateKey    = "PRIVATE KEY"
)

func loadPrivateKey(name string, blockType string) (any, error) {
	b, err := os.ReadFile(name)
	if err != nil {
		return nil, err
	}
	return ParsePrivateKey(b, blockType)
}

func LoadECPrivateKey(name string) (*ecdsa.PrivateKey, error) {
	k, err := loadPrivateKey(name, pemECPrivateKey)
	if err != nil {
		return nil, err
	}
	return k.(*ecdsa.PrivateKey), nil
}

func LoadRSAPrivateKey(name string) (*rsa.PrivateKey, error) {
	k, err := loadPrivateKey(name, pemRSAPrivateKey)
	if err != nil {
		return nil, err
	}
	return k.(*rsa.PrivateKey), nil
}

func ParseECPrivateKey(b []byte) (*ecdsa.PrivateKey, error) {
	k, err := ParsePrivateKey(b, pemECPrivateKey)
	if err != nil {
		return nil, err
	}
	return k.(*ecdsa.PrivateKey), nil
}

func ParseRSAPrivateKey(b []byte) (*rsa.PrivateKey, error) {
	k, err := ParsePrivateKey(b, pemRSAPrivateKey)
	if err != nil {
		return nil, err
	}
	return k.(*rsa.PrivateKey), nil
}

const (
	secretBoxHexEncodedLen       = 64 // hex.EncodedLen(32)
	secretBoxRawBase64EncodedLen = 43 // base64.RawStdEncoding.EncodedLen(32)
	secretBoxStdBase64EncodedLen = 44 // base64.StdEncoding.EncodedLen(32)
)

// ParseSecretBoxKey decodes a hex or base64 encoded SecretBox key.
func ParseSecretBoxKey(b []byte) ([32]byte, error) {
	var key [32]byte
	b = bytes.TrimRight(b, "\r\n")
	switch len(b) {
	case secretBoxHexEncodedLen:
		if _, err := hex.Decode(key[:], b); err != nil {
			return key, err
		}
	case secretBoxRawBase64EncodedLen:
		if _, err := base64.RawStdEncoding.Decode(key[:], b); err != nil {
			return key, err
		}
	case secretBoxStdBase64EncodedLen:
		if _, err := base64.StdEncoding.Decode(key[:], b); err != nil {
			return key, err
		}
	default:
		return key, fmt.Errorf("secret: secretbox expected %d hex bytes, or %d/%d base64 bytes, got %d",
			secretBoxHexEncodedLen, secretBoxRawBase64EncodedLen, secretBoxStdBase64EncodedLen, len(b))
	}
	return key, nil
}

func ParsePrivateKey(b []byte, blockType string) (any, error) {
	var (
		rest  = b
		block *pem.Block
	)
	for {
		if block, rest = pem.Decode(rest); block == nil {
			return nil, errors.New("secret: no PEM encoded data remains")
		}
		if block.Type == blockType || block.Type == pemPrivateKey || blockType == "" {
			break
		}
	}

	if strings.Contains(block.Headers["Proc-Type"], "ENCRYPTED") {
		return nil, fmt.Errorf("secret: can't decrypt encrypted %s PEM block", block.Type)
	}

	switch block.Type {
	case pemECPrivateKey:
		return x509.ParseECPrivateKey(block.Bytes)
	case pemRSAPrivateKey:
		return x509.ParsePKCS1PrivateKey(block.Bytes)
	case pemPrivateKey:
		return x509.ParsePKCS8PrivateKey(block.Bytes)
	default:
		return nil, fmt.Errorf("secret: don't know how to decode a %s PEM block", block.Type)
	}
}