package symm

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

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

func EncryptSM4(data, key, iv []byte, mode, paddingType string) ([]byte, error) {
	normalizedMode := normalizeCipherMode(mode)
	if normalizedMode == MODEGCM {
		return EncryptSM4GCM(data, key, iv, nil)
	}

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

func DecryptSM4(src, key, iv []byte, mode, paddingType string) ([]byte, error) {
	normalizedMode := normalizeCipherMode(mode)
	if normalizedMode == MODEGCM {
		return DecryptSM4GCM(src, key, iv, nil)
	}

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

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

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

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

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

func EncryptSM4WithOptions(data, key []byte, opts *CipherOptions) ([]byte, error) {
	cfg := normalizeCipherOptions(opts)
	mode := normalizeCipherMode(cfg.Mode)
	if mode == "" {
		mode = MODEGCM
	}
	if mode == MODEGCM {
		return EncryptSM4GCM(data, key, nonceFromOptions(cfg), cfg.AAD)
	}
	return EncryptSM4(data, key, cfg.IV, mode, cfg.Padding)
}

func DecryptSM4WithOptions(src, key []byte, opts *CipherOptions) ([]byte, error) {
	cfg := normalizeCipherOptions(opts)
	mode := normalizeCipherMode(cfg.Mode)
	if mode == "" {
		mode = MODEGCM
	}
	if mode == MODEGCM {
		return DecryptSM4GCM(src, key, nonceFromOptions(cfg), cfg.AAD)
	}
	return DecryptSM4(src, key, cfg.IV, mode, cfg.Padding)
}

func EncryptSM4StreamWithOptions(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 EncryptSM4GCMStream(dst, src, key, nonceFromOptions(cfg), cfg.AAD)
	}
	return EncryptSM4Stream(dst, src, key, cfg.IV, mode, cfg.Padding)
}

func DecryptSM4StreamWithOptions(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 DecryptSM4GCMStream(dst, src, key, nonceFromOptions(cfg), cfg.AAD)
	}
	return DecryptSM4Stream(dst, src, key, cfg.IV, mode, cfg.Padding)
}

func EncryptSM4GCM(plain, key, nonce, aad []byte) ([]byte, error) {
	block, err := sm4.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 DecryptSM4GCM(ciphertext, key, nonce, aad []byte) ([]byte, error) {
	block, err := sm4.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 EncryptSM4GCMStream(dst io.Writer, src io.Reader, key, nonce, aad []byte) error {
	plain, err := io.ReadAll(src)
	if err != nil {
		return err
	}
	out, err := EncryptSM4GCM(plain, key, nonce, aad)
	if err != nil {
		return err
	}
	_, err = dst.Write(out)
	return err
}

func DecryptSM4GCMStream(dst io.Writer, src io.Reader, key, nonce, aad []byte) error {
	enc, err := io.ReadAll(src)
	if err != nil {
		return err
	}
	out, err := DecryptSM4GCM(enc, key, nonce, aad)
	if err != nil {
		return err
	}
	_, err = dst.Write(out)
	return err
}

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 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, "")
}