package symm

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

	"b612.me/starcrypto/paddingx"
)

const (
	MODEECB = "ECB"
	MODECBC = "CBC"
	MODECFB = "CFB"
	MODEOFB = "OFB"
	MODECTR = "CTR"
	MODEGCM = "GCM"
)

var ErrUnsupportedCipherMode = errors.New("cipher mode not supported")

func normalizeCipherMode(mode string) string {
	return strings.ToUpper(strings.TrimSpace(mode))
}

func encryptWithBlockMode(block cipher.Block, data, iv []byte, mode, paddingType, defaultPadding string) ([]byte, error) {
	mode = normalizeCipherMode(mode)
	if mode == "" {
		mode = MODECBC
	}

	switch mode {
	case MODEECB:
		if paddingType == "" {
			paddingType = defaultPadding
		}
		content, err := paddingx.Pad(data, block.BlockSize(), paddingType)
		if err != nil {
			return nil, err
		}
		out := make([]byte, len(content))
		ecbEncryptBlocks(block, out, content)
		return out, nil
	case MODECBC:
		if len(iv) != block.BlockSize() {
			return nil, errors.New("iv length must match block size")
		}
		if paddingType == "" {
			paddingType = defaultPadding
		}
		content, err := paddingx.Pad(data, block.BlockSize(), paddingType)
		if err != nil {
			return nil, err
		}
		out := make([]byte, len(content))
		cipher.NewCBCEncrypter(block, iv).CryptBlocks(out, content)
		return out, nil
	case MODECFB, MODEOFB, MODECTR:
		if len(iv) != block.BlockSize() {
			return nil, errors.New("iv length must match block size")
		}
		stream, err := newCipherStream(block, iv, mode, false)
		if err != nil {
			return nil, err
		}
		out := make([]byte, len(data))
		stream.XORKeyStream(out, data)
		return out, nil
	default:
		return nil, ErrUnsupportedCipherMode
	}
}

func decryptWithBlockMode(block cipher.Block, src, iv []byte, mode, paddingType, defaultPadding string) ([]byte, error) {
	mode = normalizeCipherMode(mode)
	if mode == "" {
		mode = MODECBC
	}

	switch mode {
	case MODEECB:
		if len(src) == 0 || len(src)%block.BlockSize() != 0 {
			return nil, errors.New("ciphertext is not a full block size")
		}
		if paddingType == "" {
			paddingType = defaultPadding
		}
		decrypted := make([]byte, len(src))
		ecbDecryptBlocks(block, decrypted, src)
		return paddingx.Unpad(decrypted, block.BlockSize(), paddingType)
	case MODECBC:
		if len(iv) != block.BlockSize() {
			return nil, errors.New("iv length must match block size")
		}
		if len(src) == 0 || len(src)%block.BlockSize() != 0 {
			return nil, errors.New("ciphertext is not a full block size")
		}
		if paddingType == "" {
			paddingType = defaultPadding
		}
		decrypted := make([]byte, len(src))
		cipher.NewCBCDecrypter(block, iv).CryptBlocks(decrypted, src)
		return paddingx.Unpad(decrypted, block.BlockSize(), paddingType)
	case MODECFB, MODEOFB, MODECTR:
		if len(iv) != block.BlockSize() {
			return nil, errors.New("iv length must match block size")
		}
		stream, err := newCipherStream(block, iv, mode, true)
		if err != nil {
			return nil, err
		}
		out := make([]byte, len(src))
		stream.XORKeyStream(out, src)
		return out, nil
	default:
		return nil, ErrUnsupportedCipherMode
	}
}

func encryptWithBlockModeStream(block cipher.Block, dst io.Writer, src io.Reader, iv []byte, mode, paddingType, defaultPadding string) error {
	mode = normalizeCipherMode(mode)
	if mode == "" {
		mode = MODECBC
	}

	switch mode {
	case MODEECB:
		if paddingType == "" {
			paddingType = defaultPadding
		}
		return encryptPaddedBlockStream(dst, src, block.BlockSize(), paddingType, func(out, in []byte) {
			ecbEncryptBlocks(block, out, in)
		})
	case MODECBC:
		if len(iv) != block.BlockSize() {
			return errors.New("iv length must match block size")
		}
		if paddingType == "" {
			paddingType = defaultPadding
		}
		modeEnc := cipher.NewCBCEncrypter(block, iv)
		return encryptPaddedBlockStream(dst, src, block.BlockSize(), paddingType, modeEnc.CryptBlocks)
	case MODECFB, MODEOFB, MODECTR:
		if len(iv) != block.BlockSize() {
			return errors.New("iv length must match block size")
		}
		stream, err := newCipherStream(block, iv, mode, false)
		if err != nil {
			return err
		}
		return xorStreamCopy(dst, src, stream)
	default:
		return ErrUnsupportedCipherMode
	}
}

func decryptWithBlockModeStream(block cipher.Block, dst io.Writer, src io.Reader, iv []byte, mode, paddingType, defaultPadding string) error {
	mode = normalizeCipherMode(mode)
	if mode == "" {
		mode = MODECBC
	}

	switch mode {
	case MODEECB:
		if paddingType == "" {
			paddingType = defaultPadding
		}
		return decryptPaddedBlockStream(dst, src, block.BlockSize(), paddingType, func(out, in []byte) {
			ecbDecryptBlocks(block, out, in)
		})
	case MODECBC:
		if len(iv) != block.BlockSize() {
			return errors.New("iv length must match block size")
		}
		if paddingType == "" {
			paddingType = defaultPadding
		}
		modeDec := cipher.NewCBCDecrypter(block, iv)
		return decryptPaddedBlockStream(dst, src, block.BlockSize(), paddingType, modeDec.CryptBlocks)
	case MODECFB, MODEOFB, MODECTR:
		if len(iv) != block.BlockSize() {
			return errors.New("iv length must match block size")
		}
		stream, err := newCipherStream(block, iv, mode, true)
		if err != nil {
			return err
		}
		return xorStreamCopy(dst, src, stream)
	default:
		return ErrUnsupportedCipherMode
	}
}

func newCipherStream(block cipher.Block, iv []byte, mode string, decrypt bool) (cipher.Stream, error) {
	switch mode {
	case MODECFB:
		if decrypt {
			return cipher.NewCFBDecrypter(block, iv), nil
		}
		return cipher.NewCFBEncrypter(block, iv), nil
	case MODEOFB:
		return cipher.NewOFB(block, iv), nil
	case MODECTR:
		return cipher.NewCTR(block, iv), nil
	default:
		return nil, ErrUnsupportedCipherMode
	}
}

func xorStreamCopy(dst io.Writer, src io.Reader, stream cipher.Stream) error {
	buf := make([]byte, 32*1024)
	out := make([]byte, 32*1024)
	for {
		n, err := src.Read(buf)
		if n > 0 {
			stream.XORKeyStream(out[:n], buf[:n])
			if _, werr := dst.Write(out[:n]); werr != nil {
				return werr
			}
		}
		if err != nil {
			if err == io.EOF {
				return nil
			}
			return err
		}
	}
}

func encryptPaddedBlockStream(dst io.Writer, src io.Reader, blockSize int, paddingType string, cryptBlocks func(dst, src []byte)) error {
	pending := make([]byte, 0, blockSize*2)
	buf := make([]byte, 32*1024)

	for {
		n, err := src.Read(buf)
		if n > 0 {
			pending = append(pending, buf[:n]...)
			processLen := len(pending) - blockSize
			if processLen > 0 {
				processLen -= processLen % blockSize
				if processLen > 0 {
					out := make([]byte, processLen)
					cryptBlocks(out, pending[:processLen])
					if _, werr := dst.Write(out); werr != nil {
						return werr
					}
					pending = append([]byte(nil), pending[processLen:]...)
				}
			}
		}
		if err != nil {
			if err == io.EOF {
				break
			}
			return err
		}
	}

	content, err := paddingx.Pad(pending, blockSize, paddingType)
	if err != nil {
		return err
	}
	out := make([]byte, len(content))
	cryptBlocks(out, content)
	_, err = dst.Write(out)
	return err
}

func decryptPaddedBlockStream(dst io.Writer, src io.Reader, blockSize int, paddingType string, cryptBlocks func(dst, src []byte)) error {
	pending := make([]byte, 0, blockSize*2)
	buf := make([]byte, 32*1024)

	for {
		n, err := src.Read(buf)
		if n > 0 {
			pending = append(pending, buf[:n]...)
			processLen := len(pending) - blockSize
			if processLen > 0 {
				processLen -= processLen % blockSize
				if processLen > 0 {
					out := make([]byte, processLen)
					cryptBlocks(out, pending[:processLen])
					if _, werr := dst.Write(out); werr != nil {
						return werr
					}
					pending = append([]byte(nil), pending[processLen:]...)
				}
			}
		}
		if err != nil {
			if err == io.EOF {
				break
			}
			return err
		}
	}

	if len(pending) == 0 || len(pending)%blockSize != 0 {
		return errors.New("ciphertext is not a full block size")
	}

	decrypted := make([]byte, len(pending))
	cryptBlocks(decrypted, pending)
	out, err := paddingx.Unpad(decrypted, blockSize, paddingType)
	if err != nil {
		return err
	}
	_, err = dst.Write(out)
	return err
}

func ecbEncryptBlocks(block cipher.Block, dst, src []byte) {
	blockSize := block.BlockSize()
	for i := 0; i < len(src); i += blockSize {
		block.Encrypt(dst[i:i+blockSize], src[i:i+blockSize])
	}
}

func ecbDecryptBlocks(block cipher.Block, dst, src []byte) {
	blockSize := block.BlockSize()
	for i := 0; i < len(src); i += blockSize {
		block.Decrypt(dst[i:i+blockSize], src[i:i+blockSize])
	}
}