package sm1

import (
	"crypto/cipher"

	"xdx.jelly/xgcl/gerrors"
	"xdx.jelly/xgcl/internal/modes"
	"xdx.jelly/xgcl/utils"
)

// EncryptECB len(encData) must exactly equal to len(data)
// or return error
// we not realloc memory for encData!!!
// data和encData要么完全重叠，要么完全不重叠！！！
//
// @param dst[out]	密文预分配空间，可以为nil
// @param key[in]	16字节密钥
// @param src[in]	明文,长度必须是16的倍数，函数不负责padding
// @return []byte	输入密文。
// @note 如果输入dst的空间（cap）比src小，则可能引起内存地址得变化。
// @usage dst = ECBEncrypt(dst, key, src)
func EncryptECB(dst, key, src []byte) ([]byte, error) {

	block, err := NewCipher(key)
	if err != nil {
		return dst, err
	}

	if len(src)%BlockSize != 0 {
		return dst, gerrors.WithAnnotatingf(ErrInvalidInputLength,
			"the length of input must be a multiple of %d", BlockSize)
	}

	if dst == nil || cap(dst) < len(src) {
		dst = make([]byte, len(src))
	}
	dst = dst[:len(src)]

	blockMode := modes.NewECBEncrypter(block)
	blockMode.CryptBlocks(dst, src)
	block.(utils.Clearable).Clear()
	return dst, nil
}

// DecryptECB 解密
// params:
//
//	dst[in]		明文
//	src[in]		密文，长度必须是16的倍数，函数不负责解padding
//	key[in]		16字节密钥
func DecryptECB(dst, key, src []byte) ([]byte, error) {
	block, err := NewCipher(key)
	if err != nil {
		return dst, err
	}
	if len(src)%BlockSize != 0 {
		return dst, gerrors.WithAnnotatingf(ErrInvalidInputLength,
			"the length of input must be a multiple of %d", BlockSize)
	}

	if dst == nil || cap(dst) < len(src) {
		dst = make([]byte, len(src))
	}
	dst = dst[:len(src)]

	blockMode := modes.NewECBDecrypter(block)
	blockMode.CryptBlocks(dst, src)
	block.(utils.Clearable).Clear()
	return dst, nil
}

// EncryptCBC cbc加密
// 函数返回时iv保存最后一个分组密文，因此对大包数据可以分组多次调用进行加密。
// params:
//
//	dst[out]	密文，输入时分配好长度，必须和data的长度相等
//	iv[in/out]	初始向量, 16字节，返回时是最后一个分组密文
//	key[in]		16字节密钥
//	src[in]		明文,长度必须是16的倍数，函数不负责padding
func EncryptCBC(dst, iv, key, src []byte) ([]byte, error) {
	block, err := NewCipher(key)
	if err != nil {
		return dst, err
	}

	if len(iv) != BlockSize {
		return dst, gerrors.WithAnnotatingf(ErrInvalidInputLength,
			"the length of input iv must be %d", BlockSize)
	}

	if len(src)%BlockSize != 0 {
		return dst, gerrors.WithAnnotatingf(ErrInvalidInputLength,
			"the length of input must be %d", BlockSize)
	}

	if dst == nil || cap(dst) < len(src) {
		dst = make([]byte, len(src))
	}
	dst = dst[:len(src)]

	blockMode := cipher.NewCBCEncrypter(block, iv)
	blockMode.CryptBlocks(dst, src)

	// The last dst block is the return iv
	if len(dst) > 0 {
		copy(iv, dst[len(dst)-BlockSize:])
	}
	block.(utils.Clearable).Clear()
	return dst, nil

}

// DecryptCBC cbc解密
// 函数返回时iv保存最后一个分组密文，因此对大包数据可以分组多次调用进行加密。
// params:
//
//	dst[out]	密文，输入时分配好长度，必须和data的长度相等
//	iv[in/out]	初始向量, 16字节，返回时是最后一个分组
//	key[in]		16字节密钥
//	src[in]		明文,长度必须是16的倍数，函数不负责padding
func DecryptCBC(dst, iv, key, src []byte) ([]byte, error) {
	block, err := NewCipher(key)
	if err != nil {
		return dst, err
	}

	if len(iv) != BlockSize {
		return dst, gerrors.WithAnnotatingf(ErrInvalidInputLength,
			"the length of input iv must be %d", BlockSize)
	}

	if len(src)%BlockSize != 0 {
		return dst, gerrors.WithAnnotatingf(ErrInvalidInputLength,
			"the length of input must be %d", BlockSize)
	}

	if dst == nil || cap(dst) < len(src) {
		dst = make([]byte, len(src))
	}
	dst = dst[:len(src)]

	blockMode := cipher.NewCBCDecrypter(block, iv)
	blockMode.CryptBlocks(dst, src)

	// The last src block is the return iv
	if len(src) > 0 {
		copy(iv, src[len(src)-block.BlockSize():])
	}
	if block, ok := block.(utils.Clearable); ok {
		block.Clear()
	}
	return dst, nil
}

// EncryptCFB cfb加密
// params:
//
//	dst[out]	密文，输入时分配好长度，必须和data的长度相等
//	iv[in/out]	初始向量, 16字节
//	key[in]		16字节密钥
//	src[in]		明文,长度必须是16的倍数，函数不负责padding
func EncryptCFB(dst, iv, key, src []byte) ([]byte, error) {
	block, err := NewCipher(key)
	if err != nil {
		return dst, gerrors.WithStack(err)
	}

	if len(iv) != BlockSize {
		return dst, gerrors.WithAnnotatingf(ErrInvalidInputLength,
			"the length of input iv must be %d", BlockSize)
	}

	if len(src)%BlockSize != 0 {
		return dst, gerrors.WithAnnotatingf(ErrInvalidInputLength,
			"the length of input must be %d", BlockSize)
	}

	if dst == nil || cap(dst) < len(src) {
		dst = make([]byte, len(src))
	}
	dst = dst[:len(src)]

	blockMode := cipher.NewCFBEncrypter(block, iv)
	blockMode.XORKeyStream(dst, src)
	if block, ok := block.(utils.Clearable); ok {
		block.Clear()
	}
	return dst, nil
}

// DecryptCFB ...
func DecryptCFB(dst, iv, key, src []byte) ([]byte, error) {
	block, err := NewCipher(key)
	if err != nil {
		return dst, err
	}

	if len(iv) != BlockSize {
		return dst, gerrors.WithAnnotatingf(ErrInvalidInputLength,
			"the length of input iv must be %d", BlockSize)
	}

	if len(src)%BlockSize != 0 {
		return dst, gerrors.WithAnnotatingf(ErrInvalidInputLength,
			"the length of input must be %d", BlockSize)
	}

	if dst == nil || cap(dst) < len(src) {
		dst = make([]byte, len(src))
	}
	dst = dst[:len(src)]

	blockMode := cipher.NewCFBDecrypter(block, iv)
	blockMode.XORKeyStream(dst, src)
	if block, ok := block.(utils.Clearable); ok {
		block.Clear()
	}
	return dst, nil
}

// EncryptOFB ofb加密
func EncryptOFB(dst, iv, key, src []byte) ([]byte, error) {
	block, err := NewCipher(key)
	if err != nil {
		return dst, err
	}

	if len(iv) != BlockSize {
		return dst, gerrors.WithAnnotatingf(ErrInvalidInputLength,
			"the length of input iv must be %d", BlockSize)
	}

	if len(src)%BlockSize != 0 {
		return dst, gerrors.WithAnnotatingf(ErrInvalidInputLength,
			"the length of input must be %d", BlockSize)
	}

	if dst == nil || cap(dst) < len(src) {
		dst = make([]byte, len(src))
	}
	dst = dst[:len(src)]

	blockMode := cipher.NewOFB(block, iv)
	blockMode.XORKeyStream(dst, src)
	if block, ok := block.(utils.Clearable); ok {
		block.Clear()
	}
	return dst, nil
}

// DecryptOFB ofb解密
func DecryptOFB(dst, iv, key, src []byte) ([]byte, error) {
	block, err := NewCipher(key)
	if err != nil {
		return dst, err
	}

	if len(iv) != BlockSize {
		return dst, gerrors.WithAnnotatingf(ErrInvalidInputLength,
			"the length of input iv must be %d", BlockSize)
	}

	if len(src)%BlockSize != 0 {
		return dst, gerrors.WithAnnotatingf(ErrInvalidInputLength,
			"the length of input must be %d", BlockSize)
	}

	if dst == nil || cap(dst) < len(src) {
		dst = make([]byte, len(src))
	}
	dst = dst[:len(src)]

	blockMode := cipher.NewOFB(block, iv)
	blockMode.XORKeyStream(dst, src)
	if block, ok := block.(utils.Clearable); ok {
		block.Clear()
	}
	return dst, nil
}