package crypto

import (
	"agent/defines/derrs"
	"crypto/aes"
	"crypto/cipher"
)

func NewAesCtr(key any) (ICipher, error) {
	cipher := &aesctr{}
	if err := cipher.SetKey(key); err != nil {
		return nil, err
	}
	return cipher, nil
}

type aesctr struct {
	encryptIndex int64
	decryptIndex int64

	key   []byte
	block cipher.Block

	encryptIV []byte
	encrypt   cipher.Stream

	decryptIV []byte
	decrypt   cipher.Stream
}

func (n *aesctr) New(key any) ICipher {
	cipher := &aesctr{}
	if err := cipher.SetKey(key); err != nil {
		panic(err)
	}
	return cipher
}

func (n *aesctr) ResetState() (err error) {
	n.encryptIV = make([]byte, aes.BlockSize)
	n.decryptIV = make([]byte, aes.BlockSize)

	n.encrypt = cipher.NewCTR(n.block, n.encryptIV)
	n.decrypt = cipher.NewCTR(n.block, n.decryptIV)

	n.encryptIndex = 0
	n.decryptIndex = 0
	return nil
}

func (n *aesctr) SetKey(key any) (err error) {
	if n.key, err = any2key(key); err != nil {
		return err
	}
	if len(n.key) != 16 {
		return derrs.NewIncorrectParamsError(`len(key) != 16`)
	}
	if n.block, err = aes.NewCipher(n.key); err != nil {
		return err
	}
	n.ResetState()
	return nil
}

func (n *aesctr) SetKey2(key any, _ any) error {
	return n.SetKey(key)
}

func (n *aesctr) Encrypt(input []byte, output []byte) (int, error) {
	outputSize := n.EncryptOutputSize(len(input))
	if len(output) < outputSize {
		return outputSize, derrs.NewIncorrectParamsError("len(output) < outputSize")
	}

	n.encrypt.XORKeyStream(output, input)
	n.encryptIndex += int64(outputSize)

	return outputSize, nil
}

func (n *aesctr) Encrypt2(input []byte) ([]byte, error) {
	outputSize := n.EncryptOutputSize(len(input))
	output := make([]byte, outputSize)

	n.encrypt.XORKeyStream(output, input)
	n.encryptIndex += int64(outputSize)

	return output, nil
}

func (n *aesctr) Decrypt(input []byte, output []byte) (int, error) {
	outputSize := n.DecryptOutputSize(len(input))
	if len(output) < outputSize {
		return outputSize, derrs.NewIncorrectParamsError("len(output) < outputSize")
	}

	n.decrypt.XORKeyStream(output, input)
	n.decryptIndex += int64(outputSize)

	return outputSize, nil
}

func (n *aesctr) Decrypt2(input []byte) ([]byte, error) {
	outputSize := n.DecryptOutputSize(len(input))
	output := make([]byte, outputSize)

	n.decrypt.XORKeyStream(output, input)
	n.decryptIndex += int64(outputSize)

	return output, nil
}

func (n *aesctr) DecryptBlock(input []byte, output []byte, blocks int) (int, error) {
	outputSize := n.EncryptOutputSize(n.BlockSize() * blocks)
	if len(output) < outputSize {
		return outputSize, derrs.NewIncorrectParamsError("len(output) < outputSize")
	}

	n.decrypt.XORKeyStream(output, input)
	n.decryptIndex += int64(outputSize)

	return outputSize, nil
}

func (n *aesctr) EncryptOutputSize(inputSize int) int {
	return inputSize
}

func (n *aesctr) DecryptOutputSize(inputSize int) int {
	return inputSize
}

func (n *aesctr) BlockSize() int {
	return 1
}

func (n *aesctr) PlainTextBlockSize() int {
	return 1
}

func (n *aesctr) Encrypted() int {
	return int(n.encryptIndex)
}

func (n *aesctr) Decrypted() int {
	return int(n.decryptIndex)
}
