package symm

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

	"github.com/emmansun/gmsm/sm4"
)

var (
	sm4GCMFactory = newGCMFactory(sm4.NewCipher)
	sm4CCMFactory = newCCMFactory(sm4.NewCipher)
)

func EncryptSM4(data, key, iv []byte, mode, paddingType string) ([]byte, error) {
	return encryptBlockCipher(data, key, iv, mode, paddingType, PKCS7PADDING, sm4.NewCipher, EncryptSM4GCM, EncryptSM4CCM)
}

func DecryptSM4(src, key, iv []byte, mode, paddingType string) ([]byte, error) {
	return decryptBlockCipher(src, key, iv, mode, paddingType, PKCS7PADDING, sm4.NewCipher, DecryptSM4GCM, DecryptSM4CCM)
}

func EncryptSM4Stream(dst io.Writer, src io.Reader, key, iv []byte, mode, paddingType string) error {
	return encryptBlockCipherStream(dst, src, key, iv, mode, paddingType, PKCS7PADDING, sm4.NewCipher, EncryptSM4GCMStream, EncryptSM4CCMStream)
}

func DecryptSM4Stream(dst io.Writer, src io.Reader, key, iv []byte, mode, paddingType string) error {
	return decryptBlockCipherStream(dst, src, key, iv, mode, paddingType, PKCS7PADDING, sm4.NewCipher, DecryptSM4GCMStream, DecryptSM4CCMStream)
}

func EncryptSM4WithOptions(data, key []byte, opts *CipherOptions) ([]byte, error) {
	return encryptBlockWithOptions(data, key, opts, EncryptSM4GCM, EncryptSM4CCM, EncryptSM4)
}

func DecryptSM4WithOptions(src, key []byte, opts *CipherOptions) ([]byte, error) {
	return decryptBlockWithOptions(src, key, opts, DecryptSM4GCM, DecryptSM4CCM, DecryptSM4)
}

func EncryptSM4StreamWithOptions(dst io.Writer, src io.Reader, key []byte, opts *CipherOptions) error {
	return encryptBlockStreamWithOptions(dst, src, key, opts, EncryptSM4GCMStream, EncryptSM4CCMStream, EncryptSM4Stream)
}

func DecryptSM4StreamWithOptions(dst io.Writer, src io.Reader, key []byte, opts *CipherOptions) error {
	return decryptBlockStreamWithOptions(dst, src, key, opts, DecryptSM4GCMStream, DecryptSM4CCMStream, DecryptSM4Stream)
}

func EncryptSM4GCM(plain, key, nonce, aad []byte) ([]byte, error) {
	return encryptAEAD(sm4GCMFactory, plain, key, nonce, aad, ErrInvalidGCMNonceLength)
}

func DecryptSM4GCM(ciphertext, key, nonce, aad []byte) ([]byte, error) {
	return decryptAEAD(sm4GCMFactory, ciphertext, key, nonce, aad, ErrInvalidGCMNonceLength)
}

func EncryptSM4GCMChunk(plain, key, nonce, aad []byte, chunkIndex uint64) ([]byte, error) {
	return encryptAEADChunk(sm4GCMFactory, plain, key, nonce, aad, chunkIndex, ErrInvalidGCMNonceLength, encryptGCMChunk)
}

func DecryptSM4GCMChunk(ciphertext, key, nonce, aad []byte, chunkIndex uint64) ([]byte, error) {
	return decryptAEADChunk(sm4GCMFactory, ciphertext, key, nonce, aad, chunkIndex, ErrInvalidGCMNonceLength, decryptGCMChunk)
}

func EncryptSM4GCMStream(dst io.Writer, src io.Reader, key, nonce, aad []byte) error {
	return encryptAEADStream(sm4GCMFactory, dst, src, key, nonce, aad, ErrInvalidGCMNonceLength, encryptGCMChunkedStream)
}

func DecryptSM4GCMStream(dst io.Writer, src io.Reader, key, nonce, aad []byte) error {
	return decryptAEADStream(sm4GCMFactory, dst, src, key, nonce, aad, ErrInvalidGCMNonceLength, decryptGCMChunkedOrLegacyStream)
}

func EncryptSM4CCM(plain, key, nonce, aad []byte) ([]byte, error) {
	return encryptAEAD(sm4CCMFactory, plain, key, nonce, aad, ErrInvalidCCMNonceLength)
}

func DecryptSM4CCM(ciphertext, key, nonce, aad []byte) ([]byte, error) {
	return decryptAEAD(sm4CCMFactory, ciphertext, key, nonce, aad, ErrInvalidCCMNonceLength)
}

func EncryptSM4CCMChunk(plain, key, nonce, aad []byte, chunkIndex uint64) ([]byte, error) {
	return encryptAEADChunk(sm4CCMFactory, plain, key, nonce, aad, chunkIndex, ErrInvalidCCMNonceLength, encryptCCMChunk)
}

func DecryptSM4CCMChunk(ciphertext, key, nonce, aad []byte, chunkIndex uint64) ([]byte, error) {
	return decryptAEADChunk(sm4CCMFactory, ciphertext, key, nonce, aad, chunkIndex, ErrInvalidCCMNonceLength, decryptCCMChunk)
}

func EncryptSM4CCMStream(dst io.Writer, src io.Reader, key, nonce, aad []byte) error {
	return encryptAEADStream(sm4CCMFactory, dst, src, key, nonce, aad, ErrInvalidCCMNonceLength, encryptCCMChunkedStream)
}

func DecryptSM4CCMStream(dst io.Writer, src io.Reader, key, nonce, aad []byte) error {
	return decryptAEADStream(sm4CCMFactory, dst, src, key, nonce, aad, ErrInvalidCCMNonceLength, decryptCCMChunkedOrLegacyStream)
}
func EncryptSM4CFB(origData, key []byte) ([]byte, error) {
	block, err := sm4.NewCipher(key)
	if err != nil {
		return nil, err
	}
	out := make([]byte, block.BlockSize()+len(origData))
	iv := out[:block.BlockSize()]
	if _, err := io.ReadFull(rand.Reader, iv); err != nil {
		return nil, err
	}
	body, err := EncryptSM4CFBNoBlock(origData, key, iv)
	if err != nil {
		return nil, err
	}
	copy(out[block.BlockSize():], body)
	return out, nil
}

func DecryptSM4CFB(encrypted, key []byte) ([]byte, error) {
	block, err := sm4.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 DecryptSM4CFBNoBlock(encrypted[block.BlockSize():], key, iv)
}

func EncryptSM4CFBNoBlock(origData, key, iv []byte) ([]byte, error) {
	return EncryptSM4(origData, key, iv, MODECFB, "")
}

func DecryptSM4CFBNoBlock(encrypted, key, iv []byte) ([]byte, error) {
	return DecryptSM4(encrypted, key, iv, MODECFB, "")
}

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

// DecryptSM4CBCFromSecondBlock 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 DecryptSM4CBCFromSecondBlock(src, key, prevCipherBlock []byte) ([]byte, error) {
	block, err := sm4.NewCipher(key)
	if err != nil {
		return nil, err
	}
	return decryptCBCFromSecondBlock(block, src, prevCipherBlock)
}

// DecryptSM4CFBFromSecondBlock 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 DecryptSM4CFBFromSecondBlock(src, key, prevCipherBlock []byte) ([]byte, error) {
	block, err := sm4.NewCipher(key)
	if err != nil {
		return nil, err
	}
	return decryptCFBFromSecondBlock(block, src, prevCipherBlock)
}