// chaPolyShrink provides functions to encrypt or decrypt bytes using ChaPolyShrink
// ChaPolyShrink first compresses bytes using zlib, adds padding to obscure size of encrypted bytes, and ciphers data using chaPoly

package chaPolyShrink

import "encoding/binary"
import "golang.org/x/crypto/chacha20poly1305"
import "compress/zlib"
import "crypto/rand"
import "errors"
import "bytes"
import "io"
import "slices"

// A paddingSize is specified, which adds padding so that resulting encrypted result size is divisible by paddingSize

//Outputs:
//	-[]byte: ChaPoly Encrypted bytes
//	-error
func EncryptChaPolyShrink(input []byte, key [32]byte, nonce [24]byte, compressContent bool, paddingSize int, includeAdditionalData bool, additionalData [32]byte)([]byte, error){

	if (paddingSize > 10000000 || paddingSize < 0) {
		return nil, errors.New("EncryptChaPolyShrink called with invalid paddingSize.")
	}

	getCompressionStrength := func()int{
		if (compressContent == false){
			return 0
		}
		return 9
	}

	compressionStrength := getCompressionStrength()

	var compressedBuffer bytes.Buffer

	compressorWriter, err := zlib.NewWriterLevel(&compressedBuffer, compressionStrength)
	if (err != nil) { return nil, err }
	compressorWriter.Write(input)
	compressorWriter.Close()

	compressedBytes := compressedBuffer.Bytes()
	 
	compressedBytesLength := len(compressedBytes)

	getNeededPaddingLength := func()int{

		if (paddingSize == 0){
			return 0
		}

		if (compressedBytesLength < paddingSize){
			paddingLength := paddingSize - compressedBytesLength
			return paddingLength
		}

		remainder := compressedBytesLength % paddingSize

		if (remainder == 0){
			return 0
		}

		paddingLengthInBytes := paddingSize - remainder

		return paddingLengthInBytes	
	}

	neededPaddingLength := getNeededPaddingLength()

	paddingLengthUint32 := uint32(neededPaddingLength)

	paddingLengthHeader := make([]byte, 4)

	binary.LittleEndian.PutUint32(paddingLengthHeader, paddingLengthUint32)

	paddingBytes := make([]byte, neededPaddingLength)
	_, err = rand.Read(paddingBytes[:])
	if (err != nil){ return nil, err }

	compressedBytesWithPaddingAndHeader := slices.Concat(paddingLengthHeader, paddingBytes, compressedBytes)

	getAdditionalDataParameter := func()[]byte{
		if (includeAdditionalData == false){
			return nil
		}

		result := additionalData[:]

		return result
	}

	additionalDataParameter := getAdditionalDataParameter()

	cipherObject, err := chacha20poly1305.NewX(key[:])
	if (err != nil) { return nil, err }

	encryptedMessage := cipherObject.Seal(nil, nonce[:], compressedBytesWithPaddingAndHeader, additionalDataParameter)

	return encryptedMessage, nil
}


//Outputs:
//	-bool: Able to decrypt
//	-[]byte: Decrypted bytes
//	-error (will return error if inputs are invalid)
func DecryptChaPolyShrink(inputBytes []byte, key [32]byte, nonce [24]byte, additionalDataExists bool, additionalData [32]byte)(bool, []byte, error){

	getAdditionalDataParameter := func()[]byte{
		if (additionalDataExists == false){
			return nil
		}

		result := additionalData[:]

		return result
	}

	additionalDataParameter := getAdditionalDataParameter()

	cipherObject, err := chacha20poly1305.NewX(key[:])
	if (err != nil) { return false, nil, err }

	decryptedBytes, err := cipherObject.Open(nil, nonce[:], inputBytes, additionalDataParameter)
	if (err != nil) {
		return false, nil, nil
	}

	paddingHeader := decryptedBytes[:4]

	paddingBytesUint32 := binary.LittleEndian.Uint32(paddingHeader[:])

	endOfPaddingIndex := 4 + paddingBytesUint32

	compressedBytes := decryptedBytes[endOfPaddingIndex:]
	compressedBytesReader := bytes.NewReader(compressedBytes)

	decompressedReader, err := zlib.NewReader(compressedBytesReader)
	if (err != nil) {
		return false, nil, nil
	}

	decompressedBytes, err := io.ReadAll(decompressedReader)
	if (err != nil) {
		return false, nil, nil
	}

	return true, decompressedBytes, nil
}