package symm

import (
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"errors"
	"io"

	"b612.me/starcrypto/ccm"
	"b612.me/starcrypto/paddingx"
)

const (
	PKCS5PADDING    = paddingx.PKCS5
	PKCS7PADDING    = paddingx.PKCS7
	ZEROPADDING     = paddingx.ZERO
	ANSIX923PADDING = paddingx.ANSIX923
)

var (
	ErrInvalidGCMNonceLength = errors.New("gcm nonce length must be 12 bytes")
	ErrInvalidCCMNonceLength = errors.New("ccm nonce length must be 12 bytes")
)

const (
	aeadCCMTagSize   = 16
	aeadCCMNonceSize = 12
)

func EncryptAes(data, key, iv []byte, mode, paddingType string) ([]byte, error) {
	normalizedMode := normalizeCipherMode(mode)
	if normalizedMode == "" {
		normalizedMode = MODEGCM
	}
	if normalizedMode == MODEGCM {
		return EncryptAesGCM(data, key, iv, nil)
	}
	if normalizedMode == MODECCM {
		return EncryptAesCCM(data, key, iv, nil)
	}

	block, err := aes.NewCipher(key)
	if err != nil {
		return nil, err
	}
	return encryptWithBlockMode(block, data, iv, normalizedMode, paddingType, PKCS7PADDING)
}

func DecryptAes(src, key, iv []byte, mode, paddingType string) ([]byte, error) {
	normalizedMode := normalizeCipherMode(mode)
	if normalizedMode == "" {
		normalizedMode = MODEGCM
	}
	if normalizedMode == MODEGCM {
		return DecryptAesGCM(src, key, iv, nil)
	}
	if normalizedMode == MODECCM {
		return DecryptAesCCM(src, key, iv, nil)
	}

	block, err := aes.NewCipher(key)
	if err != nil {
		return nil, err
	}
	return decryptWithBlockMode(block, src, iv, normalizedMode, paddingType, PKCS7PADDING)
}

func EncryptAesStream(dst io.Writer, src io.Reader, key, iv []byte, mode, paddingType string) error {
	normalizedMode := normalizeCipherMode(mode)
	if normalizedMode == "" {
		normalizedMode = MODEGCM
	}
	if normalizedMode == MODEGCM {
		return EncryptAesGCMStream(dst, src, key, iv, nil)
	}
	if normalizedMode == MODECCM {
		return EncryptAesCCMStream(dst, src, key, iv, nil)
	}

	block, err := aes.NewCipher(key)
	if err != nil {
		return err
	}
	return encryptWithBlockModeStream(block, dst, src, iv, normalizedMode, paddingType, PKCS7PADDING)
}

func DecryptAesStream(dst io.Writer, src io.Reader, key, iv []byte, mode, paddingType string) error {
	normalizedMode := normalizeCipherMode(mode)
	if normalizedMode == "" {
		normalizedMode = MODEGCM
	}
	if normalizedMode == MODEGCM {
		return DecryptAesGCMStream(dst, src, key, iv, nil)
	}
	if normalizedMode == MODECCM {
		return DecryptAesCCMStream(dst, src, key, iv, nil)
	}

	block, err := aes.NewCipher(key)
	if err != nil {
		return err
	}
	return decryptWithBlockModeStream(block, dst, src, iv, normalizedMode, paddingType, PKCS7PADDING)
}

func EncryptAesWithOptions(data, key []byte, opts *CipherOptions) ([]byte, error) {
	cfg := normalizeCipherOptions(opts)
	mode := normalizeCipherMode(cfg.Mode)
	if mode == "" {
		mode = MODEGCM
	}
	if mode == MODEGCM {
		return EncryptAesGCM(data, key, nonceFromOptions(cfg), cfg.AAD)
	}
	if mode == MODECCM {
		return EncryptAesCCM(data, key, nonceFromOptions(cfg), cfg.AAD)
	}
	return EncryptAes(data, key, cfg.IV, mode, cfg.Padding)
}

func DecryptAesWithOptions(src, key []byte, opts *CipherOptions) ([]byte, error) {
	cfg := normalizeCipherOptions(opts)
	mode := normalizeCipherMode(cfg.Mode)
	if mode == "" {
		mode = MODEGCM
	}
	if mode == MODEGCM {
		return DecryptAesGCM(src, key, nonceFromOptions(cfg), cfg.AAD)
	}
	if mode == MODECCM {
		return DecryptAesCCM(src, key, nonceFromOptions(cfg), cfg.AAD)
	}
	return DecryptAes(src, key, cfg.IV, mode, cfg.Padding)
}

func EncryptAesStreamWithOptions(dst io.Writer, src io.Reader, key []byte, opts *CipherOptions) error {
	cfg := normalizeCipherOptions(opts)
	mode := normalizeCipherMode(cfg.Mode)
	if mode == "" {
		mode = MODEGCM
	}
	if mode == MODEGCM {
		return EncryptAesGCMStream(dst, src, key, nonceFromOptions(cfg), cfg.AAD)
	}
	if mode == MODECCM {
		return EncryptAesCCMStream(dst, src, key, nonceFromOptions(cfg), cfg.AAD)
	}
	return EncryptAesStream(dst, src, key, cfg.IV, mode, cfg.Padding)
}

func DecryptAesStreamWithOptions(dst io.Writer, src io.Reader, key []byte, opts *CipherOptions) error {
	cfg := normalizeCipherOptions(opts)
	mode := normalizeCipherMode(cfg.Mode)
	if mode == "" {
		mode = MODEGCM
	}
	if mode == MODEGCM {
		return DecryptAesGCMStream(dst, src, key, nonceFromOptions(cfg), cfg.AAD)
	}
	if mode == MODECCM {
		return DecryptAesCCMStream(dst, src, key, nonceFromOptions(cfg), cfg.AAD)
	}
	return DecryptAesStream(dst, src, key, cfg.IV, mode, cfg.Padding)
}

func EncryptAesGCM(plain, key, nonce, aad []byte) ([]byte, error) {
	block, err := aes.NewCipher(key)
	if err != nil {
		return nil, err
	}
	gcm, err := cipher.NewGCM(block)
	if err != nil {
		return nil, err
	}
	if len(nonce) != gcm.NonceSize() {
		return nil, ErrInvalidGCMNonceLength
	}
	return gcm.Seal(nil, nonce, plain, aad), nil
}

func DecryptAesGCM(ciphertext, key, nonce, aad []byte) ([]byte, error) {
	block, err := aes.NewCipher(key)
	if err != nil {
		return nil, err
	}
	gcm, err := cipher.NewGCM(block)
	if err != nil {
		return nil, err
	}
	if len(nonce) != gcm.NonceSize() {
		return nil, ErrInvalidGCMNonceLength
	}
	return gcm.Open(nil, nonce, ciphertext, aad)
}

func newAesCCM(key []byte) (cipher.AEAD, error) {
	block, err := aes.NewCipher(key)
	if err != nil {
		return nil, err
	}
	return ccm.NewCCM(block, aeadCCMTagSize, aeadCCMNonceSize)
}

func EncryptAesCCM(plain, key, nonce, aad []byte) ([]byte, error) {
	aead, err := newAesCCM(key)
	if err != nil {
		return nil, err
	}
	if len(nonce) != aead.NonceSize() {
		return nil, ErrInvalidCCMNonceLength
	}
	return aead.Seal(nil, nonce, plain, aad), nil
}

func DecryptAesCCM(ciphertext, key, nonce, aad []byte) ([]byte, error) {
	aead, err := newAesCCM(key)
	if err != nil {
		return nil, err
	}
	if len(nonce) != aead.NonceSize() {
		return nil, ErrInvalidCCMNonceLength
	}
	return aead.Open(nil, nonce, ciphertext, aad)
}

func EncryptAesCCMChunk(plain, key, nonce, aad []byte, chunkIndex uint64) ([]byte, error) {
	aead, err := newAesCCM(key)
	if err != nil {
		return nil, err
	}
	if len(nonce) != aead.NonceSize() {
		return nil, ErrInvalidCCMNonceLength
	}
	return encryptCCMChunk(aead, plain, nonce, aad, chunkIndex), nil
}

func DecryptAesCCMChunk(ciphertext, key, nonce, aad []byte, chunkIndex uint64) ([]byte, error) {
	aead, err := newAesCCM(key)
	if err != nil {
		return nil, err
	}
	if len(nonce) != aead.NonceSize() {
		return nil, ErrInvalidCCMNonceLength
	}
	return decryptCCMChunk(aead, ciphertext, nonce, aad, chunkIndex)
}

func EncryptAesCCMStream(dst io.Writer, src io.Reader, key, nonce, aad []byte) error {
	aead, err := newAesCCM(key)
	if err != nil {
		return err
	}
	if len(nonce) != aead.NonceSize() {
		return ErrInvalidCCMNonceLength
	}
	return encryptCCMChunkedStream(dst, src, aead, nonce, aad)
}

func DecryptAesCCMStream(dst io.Writer, src io.Reader, key, nonce, aad []byte) error {
	aead, err := newAesCCM(key)
	if err != nil {
		return err
	}
	if len(nonce) != aead.NonceSize() {
		return ErrInvalidCCMNonceLength
	}
	return decryptCCMChunkedOrLegacyStream(dst, src, aead, nonce, aad)
}

func EncryptAesGCMChunk(plain, key, nonce, aad []byte, chunkIndex uint64) ([]byte, error) {
	block, err := aes.NewCipher(key)
	if err != nil {
		return nil, err
	}
	gcm, err := cipher.NewGCM(block)
	if err != nil {
		return nil, err
	}
	if len(nonce) != gcm.NonceSize() {
		return nil, ErrInvalidGCMNonceLength
	}
	return encryptGCMChunk(gcm, plain, nonce, aad, chunkIndex), nil
}

func DecryptAesGCMChunk(ciphertext, key, nonce, aad []byte, chunkIndex uint64) ([]byte, error) {
	block, err := aes.NewCipher(key)
	if err != nil {
		return nil, err
	}
	gcm, err := cipher.NewGCM(block)
	if err != nil {
		return nil, err
	}
	if len(nonce) != gcm.NonceSize() {
		return nil, ErrInvalidGCMNonceLength
	}
	return decryptGCMChunk(gcm, ciphertext, nonce, aad, chunkIndex)
}

func EncryptAesGCMStream(dst io.Writer, src io.Reader, key, nonce, aad []byte) error {
	block, err := aes.NewCipher(key)
	if err != nil {
		return err
	}
	gcm, err := cipher.NewGCM(block)
	if err != nil {
		return err
	}
	if len(nonce) != gcm.NonceSize() {
		return ErrInvalidGCMNonceLength
	}
	return encryptGCMChunkedStream(dst, src, gcm, nonce, aad)
}

func DecryptAesGCMStream(dst io.Writer, src io.Reader, key, nonce, aad []byte) error {
	block, err := aes.NewCipher(key)
	if err != nil {
		return err
	}
	gcm, err := cipher.NewGCM(block)
	if err != nil {
		return err
	}
	if len(nonce) != gcm.NonceSize() {
		return ErrInvalidGCMNonceLength
	}
	return decryptGCMChunkedOrLegacyStream(dst, src, gcm, nonce, aad)
}

func EncryptAesECB(data, key []byte, paddingType string) ([]byte, error) {
	return EncryptAes(data, key, nil, MODEECB, paddingType)
}

func DecryptAesECB(src, key []byte, paddingType string) ([]byte, error) {
	return DecryptAes(src, key, nil, MODEECB, paddingType)
}

func EncryptAesCBC(data, key, iv []byte, paddingType string) ([]byte, error) {
	return EncryptAes(data, key, iv, MODECBC, paddingType)
}

func DecryptAesCBC(src, key, iv []byte, paddingType string) ([]byte, error) {
	return DecryptAes(src, key, iv, MODECBC, paddingType)
}

func EncryptAesCFB(data, key, iv []byte) ([]byte, error) {
	return EncryptAes(data, key, iv, MODECFB, "")
}

func DecryptAesCFB(src, key, iv []byte) ([]byte, error) {
	return DecryptAes(src, key, iv, MODECFB, "")
}

func EncryptAesOFB(data, key, iv []byte) ([]byte, error) {
	return EncryptAes(data, key, iv, MODEOFB, "")
}

func DecryptAesOFB(src, key, iv []byte) ([]byte, error) {
	return DecryptAes(src, key, iv, MODEOFB, "")
}

func EncryptAesCTR(data, key, iv []byte) ([]byte, error) {
	return EncryptAes(data, key, iv, MODECTR, "")
}

func DecryptAesCTR(src, key, iv []byte) ([]byte, error) {
	return DecryptAes(src, key, iv, MODECTR, "")
}

func EncryptAesCTRAt(data, key, iv []byte, offset int64) ([]byte, error) {
	block, err := aes.NewCipher(key)
	if err != nil {
		return nil, err
	}
	return xorCTRAtOffset(block, data, iv, offset)
}

func DecryptAesCTRAt(src, key, iv []byte, offset int64) ([]byte, error) {
	return EncryptAesCTRAt(src, key, iv, offset)
}

func EncryptAesECBStream(dst io.Writer, src io.Reader, key []byte, paddingType string) error {
	return EncryptAesStream(dst, src, key, nil, MODEECB, paddingType)
}

func DecryptAesECBStream(dst io.Writer, src io.Reader, key []byte, paddingType string) error {
	return DecryptAesStream(dst, src, key, nil, MODEECB, paddingType)
}

func EncryptAesCBCStream(dst io.Writer, src io.Reader, key, iv []byte, paddingType string) error {
	return EncryptAesStream(dst, src, key, iv, MODECBC, paddingType)
}

func DecryptAesCBCStream(dst io.Writer, src io.Reader, key, iv []byte, paddingType string) error {
	return DecryptAesStream(dst, src, key, iv, MODECBC, paddingType)
}

func EncryptAesCFBStream(dst io.Writer, src io.Reader, key, iv []byte) error {
	return EncryptAesStream(dst, src, key, iv, MODECFB, "")
}

func DecryptAesCFBStream(dst io.Writer, src io.Reader, key, iv []byte) error {
	return DecryptAesStream(dst, src, key, iv, MODECFB, "")
}

func EncryptAesOFBStream(dst io.Writer, src io.Reader, key, iv []byte) error {
	return EncryptAesStream(dst, src, key, iv, MODEOFB, "")
}

func DecryptAesOFBStream(dst io.Writer, src io.Reader, key, iv []byte) error {
	return DecryptAesStream(dst, src, key, iv, MODEOFB, "")
}

func EncryptAesCTRStream(dst io.Writer, src io.Reader, key, iv []byte) error {
	return EncryptAesStream(dst, src, key, iv, MODECTR, "")
}

func DecryptAesCTRStream(dst io.Writer, src io.Reader, key, iv []byte) error {
	return DecryptAesStream(dst, src, key, iv, MODECTR, "")
}

func CustomEncryptAesCFB(origData, key []byte) ([]byte, error) {
	block, err := aes.NewCipher(key)
	if err != nil {
		return nil, err
	}

	encrypted := make([]byte, aes.BlockSize+len(origData))
	iv := encrypted[:block.BlockSize()]
	if _, err := io.ReadFull(rand.Reader, iv); err != nil {
		return nil, err
	}

	body, err := EncryptAesCFB(origData, key, iv)
	if err != nil {
		return nil, err
	}
	copy(encrypted[block.BlockSize():], body)
	return encrypted, nil
}

func CustomDecryptAesCFB(encrypted, key []byte) ([]byte, error) {
	block, err := aes.NewCipher(key)
	if err != nil {
		return nil, err
	}
	if len(encrypted) < block.BlockSize() {
		return nil, errors.New("ciphertext too short")
	}

	iv := encrypted[:block.BlockSize()]
	return DecryptAesCFB(encrypted[block.BlockSize():], key, iv)
}

func CustomEncryptAesCFBNoBlock(origData, key, iv []byte) ([]byte, error) {
	return EncryptAesCFB(origData, key, iv)
}

func CustomDecryptAesCFBNoBlock(encrypted, key, iv []byte) ([]byte, error) {
	return DecryptAesCFB(encrypted, key, iv)
}

func PKCS5Padding(cipherText []byte, blockSize int) []byte {
	out, _ := paddingx.Pad(cipherText, blockSize, PKCS5PADDING)
	return out
}

func PKCS5Trimming(encrypted []byte) []byte {
	if len(encrypted) == 0 {
		return nil
	}
	padding := int(encrypted[len(encrypted)-1])
	if padding <= 0 || padding > len(encrypted) {
		return nil
	}
	for i := len(encrypted) - padding; i < len(encrypted); i++ {
		if int(encrypted[i]) != padding {
			return nil
		}
	}
	return encrypted[:len(encrypted)-padding]
}

func PKCS7Padding(cipherText []byte, blockSize int) []byte {
	out, _ := paddingx.Pad(cipherText, blockSize, PKCS7PADDING)
	return out
}

func PKCS7Trimming(encrypted []byte, blockSize int) []byte {
	out, err := paddingx.Unpad(encrypted, blockSize, PKCS7PADDING)
	if err != nil {
		return nil
	}
	return out
}

func EncryptAesCFB8(data, key, iv []byte) ([]byte, error) {
	block, err := aes.NewCipher(key)
	if err != nil {
		return nil, err
	}
	return encryptCFB8(block, data, iv)
}

func DecryptAesCFB8(src, key, iv []byte) ([]byte, error) {
	block, err := aes.NewCipher(key)
	if err != nil {
		return nil, err
	}
	return decryptCFB8(block, src, iv)
}

func EncryptAesCFB8Stream(dst io.Writer, src io.Reader, key, iv []byte) error {
	block, err := aes.NewCipher(key)
	if err != nil {
		return err
	}
	return encryptCFB8Stream(block, dst, src, iv, false)
}

func DecryptAesCFB8Stream(dst io.Writer, src io.Reader, key, iv []byte) error {
	block, err := aes.NewCipher(key)
	if err != nil {
		return err
	}
	return encryptCFB8Stream(block, dst, src, iv, true)
}

func DecryptAesECBBlocks(src, key []byte) ([]byte, error) {
	block, err := aes.NewCipher(key)
	if err != nil {
		return nil, err
	}
	return decryptECBBlocks(block, src)
}

// DecryptAesCBCFromSecondBlock decrypts a CBC ciphertext segment that starts from block 2 or later.
// prevCipherBlock must be the previous ciphertext block. For data from block 2, pass block 1 as prevCipherBlock.
func DecryptAesCBCFromSecondBlock(src, key, prevCipherBlock []byte) ([]byte, error) {
	block, err := aes.NewCipher(key)
	if err != nil {
		return nil, err
	}
	return decryptCBCFromSecondBlock(block, src, prevCipherBlock)
}

// DecryptAesCFBFromSecondBlock decrypts a CFB ciphertext segment that starts from block 2 or later.
// prevCipherBlock must be the previous ciphertext block. For data from block 2, pass block 1 as prevCipherBlock.
func DecryptAesCFBFromSecondBlock(src, key, prevCipherBlock []byte) ([]byte, error) {
	block, err := aes.NewCipher(key)
	if err != nil {
		return nil, err
	}
	return decryptCFBFromSecondBlock(block, src, prevCipherBlock)
}