package symm

import (
	"crypto/cipher"
	"io"

	"b612.me/starcrypto/ccm"
)

type blockCipherFactory func(key []byte) (cipher.Block, error)
type aeadFactory func(key []byte) (cipher.AEAD, error)

type aeadBytesFunc func(data, key, nonce, aad []byte) ([]byte, error)
type aeadStreamFunc func(dst io.Writer, src io.Reader, key, nonce, aad []byte) error
type blockModeBytesFunc func(data, key, iv []byte, mode, paddingType string) ([]byte, error)
type blockModeStreamFunc func(dst io.Writer, src io.Reader, key, iv []byte, mode, paddingType string) error

type aeadChunkEncryptFunc func(aead cipher.AEAD, plain, nonce, aad []byte, chunkIndex uint64) []byte
type aeadChunkDecryptFunc func(aead cipher.AEAD, ciphertext, nonce, aad []byte, chunkIndex uint64) ([]byte, error)
type aeadStreamCodec func(dst io.Writer, src io.Reader, aead cipher.AEAD, nonce, aad []byte) error

func newGCMFactory(newBlock blockCipherFactory) aeadFactory {
	return func(key []byte) (cipher.AEAD, error) {
		block, err := newBlock(key)
		if err != nil {
			return nil, err
		}
		return cipher.NewGCM(block)
	}
}

func newCCMFactory(newBlock blockCipherFactory) aeadFactory {
	return func(key []byte) (cipher.AEAD, error) {
		block, err := newBlock(key)
		if err != nil {
			return nil, err
		}
		return ccm.NewCCM(block, aeadCCMTagSize, aeadCCMNonceSize)
	}
}

func normalizeModeOrDefaultAEAD(mode string) string {
	mode = normalizeCipherMode(mode)
	if mode == "" {
		mode = MODEGCM
	}
	return mode
}

func encryptBlockCipher(data, key, iv []byte, mode, paddingType, defaultPadding string, newBlock blockCipherFactory, encryptGCM, encryptCCM aeadBytesFunc) ([]byte, error) {
	mode = normalizeModeOrDefaultAEAD(mode)
	switch mode {
	case MODEGCM:
		return encryptGCM(data, key, iv, nil)
	case MODECCM:
		return encryptCCM(data, key, iv, nil)
	default:
		block, err := newBlock(key)
		if err != nil {
			return nil, err
		}
		return encryptWithBlockMode(block, data, iv, mode, paddingType, defaultPadding)
	}
}

func decryptBlockCipher(src, key, iv []byte, mode, paddingType, defaultPadding string, newBlock blockCipherFactory, decryptGCM, decryptCCM aeadBytesFunc) ([]byte, error) {
	mode = normalizeModeOrDefaultAEAD(mode)
	switch mode {
	case MODEGCM:
		return decryptGCM(src, key, iv, nil)
	case MODECCM:
		return decryptCCM(src, key, iv, nil)
	default:
		block, err := newBlock(key)
		if err != nil {
			return nil, err
		}
		return decryptWithBlockMode(block, src, iv, mode, paddingType, defaultPadding)
	}
}

func encryptBlockCipherStream(dst io.Writer, src io.Reader, key, iv []byte, mode, paddingType, defaultPadding string, newBlock blockCipherFactory, encryptGCMStream, encryptCCMStream aeadStreamFunc) error {
	mode = normalizeModeOrDefaultAEAD(mode)
	switch mode {
	case MODEGCM:
		return encryptGCMStream(dst, src, key, iv, nil)
	case MODECCM:
		return encryptCCMStream(dst, src, key, iv, nil)
	default:
		block, err := newBlock(key)
		if err != nil {
			return err
		}
		return encryptWithBlockModeStream(block, dst, src, iv, mode, paddingType, defaultPadding)
	}
}

func decryptBlockCipherStream(dst io.Writer, src io.Reader, key, iv []byte, mode, paddingType, defaultPadding string, newBlock blockCipherFactory, decryptGCMStream, decryptCCMStream aeadStreamFunc) error {
	mode = normalizeModeOrDefaultAEAD(mode)
	switch mode {
	case MODEGCM:
		return decryptGCMStream(dst, src, key, iv, nil)
	case MODECCM:
		return decryptCCMStream(dst, src, key, iv, nil)
	default:
		block, err := newBlock(key)
		if err != nil {
			return err
		}
		return decryptWithBlockModeStream(block, dst, src, iv, mode, paddingType, defaultPadding)
	}
}

func encryptBlockWithOptions(data, key []byte, opts *CipherOptions, encryptGCM, encryptCCM aeadBytesFunc, encryptBlock blockModeBytesFunc) ([]byte, error) {
	cfg := normalizeCipherOptions(opts)
	mode := normalizeModeOrDefaultAEAD(cfg.Mode)
	switch mode {
	case MODEGCM:
		return encryptGCM(data, key, nonceFromOptions(cfg), cfg.AAD)
	case MODECCM:
		return encryptCCM(data, key, nonceFromOptions(cfg), cfg.AAD)
	default:
		return encryptBlock(data, key, cfg.IV, mode, cfg.Padding)
	}
}

func decryptBlockWithOptions(data, key []byte, opts *CipherOptions, decryptGCM, decryptCCM aeadBytesFunc, decryptBlock blockModeBytesFunc) ([]byte, error) {
	cfg := normalizeCipherOptions(opts)
	mode := normalizeModeOrDefaultAEAD(cfg.Mode)
	switch mode {
	case MODEGCM:
		return decryptGCM(data, key, nonceFromOptions(cfg), cfg.AAD)
	case MODECCM:
		return decryptCCM(data, key, nonceFromOptions(cfg), cfg.AAD)
	default:
		return decryptBlock(data, key, cfg.IV, mode, cfg.Padding)
	}
}

func encryptBlockStreamWithOptions(dst io.Writer, src io.Reader, key []byte, opts *CipherOptions, encryptGCM, encryptCCM aeadStreamFunc, encryptBlockStream blockModeStreamFunc) error {
	cfg := normalizeCipherOptions(opts)
	mode := normalizeModeOrDefaultAEAD(cfg.Mode)
	switch mode {
	case MODEGCM:
		return encryptGCM(dst, src, key, nonceFromOptions(cfg), cfg.AAD)
	case MODECCM:
		return encryptCCM(dst, src, key, nonceFromOptions(cfg), cfg.AAD)
	default:
		return encryptBlockStream(dst, src, key, cfg.IV, mode, cfg.Padding)
	}
}

func decryptBlockStreamWithOptions(dst io.Writer, src io.Reader, key []byte, opts *CipherOptions, decryptGCM, decryptCCM aeadStreamFunc, decryptBlockStream blockModeStreamFunc) error {
	cfg := normalizeCipherOptions(opts)
	mode := normalizeModeOrDefaultAEAD(cfg.Mode)
	switch mode {
	case MODEGCM:
		return decryptGCM(dst, src, key, nonceFromOptions(cfg), cfg.AAD)
	case MODECCM:
		return decryptCCM(dst, src, key, nonceFromOptions(cfg), cfg.AAD)
	default:
		return decryptBlockStream(dst, src, key, cfg.IV, mode, cfg.Padding)
	}
}

func buildAEAD(factory aeadFactory, key, nonce []byte, errInvalidNonce error) (cipher.AEAD, error) {
	aead, err := factory(key)
	if err != nil {
		return nil, err
	}
	if len(nonce) != aead.NonceSize() {
		return nil, errInvalidNonce
	}
	return aead, nil
}

func encryptAEAD(factory aeadFactory, plain, key, nonce, aad []byte, errInvalidNonce error) ([]byte, error) {
	aead, err := buildAEAD(factory, key, nonce, errInvalidNonce)
	if err != nil {
		return nil, err
	}
	return aead.Seal(nil, nonce, plain, aad), nil
}

func decryptAEAD(factory aeadFactory, ciphertext, key, nonce, aad []byte, errInvalidNonce error) ([]byte, error) {
	aead, err := buildAEAD(factory, key, nonce, errInvalidNonce)
	if err != nil {
		return nil, err
	}
	return aead.Open(nil, nonce, ciphertext, aad)
}

func encryptAEADChunk(factory aeadFactory, plain, key, nonce, aad []byte, chunkIndex uint64, errInvalidNonce error, encryptChunk aeadChunkEncryptFunc) ([]byte, error) {
	aead, err := buildAEAD(factory, key, nonce, errInvalidNonce)
	if err != nil {
		return nil, err
	}
	return encryptChunk(aead, plain, nonce, aad, chunkIndex), nil
}

func decryptAEADChunk(factory aeadFactory, ciphertext, key, nonce, aad []byte, chunkIndex uint64, errInvalidNonce error, decryptChunk aeadChunkDecryptFunc) ([]byte, error) {
	aead, err := buildAEAD(factory, key, nonce, errInvalidNonce)
	if err != nil {
		return nil, err
	}
	return decryptChunk(aead, ciphertext, nonce, aad, chunkIndex)
}

func encryptAEADStream(factory aeadFactory, dst io.Writer, src io.Reader, key, nonce, aad []byte, errInvalidNonce error, encryptStream aeadStreamCodec) error {
	aead, err := buildAEAD(factory, key, nonce, errInvalidNonce)
	if err != nil {
		return err
	}
	return encryptStream(dst, src, aead, nonce, aad)
}

func decryptAEADStream(factory aeadFactory, dst io.Writer, src io.Reader, key, nonce, aad []byte, errInvalidNonce error, decryptStream aeadStreamCodec) error {
	aead, err := buildAEAD(factory, key, nonce, errInvalidNonce)
	if err != nil {
		return err
	}
	return decryptStream(dst, src, aead, nonce, aad)
}