package generator

import (
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"encoding/hex"
	"errors"
	"io"
)

const (
	maxEncodeLen = 32
)

// KeyGenerator key generator
type KeyGenerator struct {
	AESKey string
}

func encryptAESCFB(msg, key []byte) ([]byte, error) {
	block, err := aes.NewCipher(key)
	if err != nil {
		return nil, err
	}

	ciphertext := make([]byte, aes.BlockSize+len(msg))
	iv := ciphertext[:aes.BlockSize]
	if _, err := io.ReadFull(rand.Reader, iv); err != nil {
		return nil, err
	}

	stream := cipher.NewCFBEncrypter(block, iv)
	stream.XORKeyStream(ciphertext[aes.BlockSize:], msg)

	return ciphertext, nil
}

func decryptAESCFB(msg, key []byte) ([]byte, error) {
	block, err := aes.NewCipher(key)
	if err != nil {
		return nil, err
	}

	if len(msg) < aes.BlockSize {
		return nil, errors.New("decrypt message too short")
	}
	iv := msg[:aes.BlockSize]
	msg = msg[aes.BlockSize:]

	stream := cipher.NewCFBDecrypter(block, iv)

	stream.XORKeyStream(msg, msg)
	return msg, nil
}

// NewKeyGenerator create a key generator
func NewKeyGenerator(key string) (*KeyGenerator, error) {
	l := len(key)
	if l != 16 && l != 24 && l != 32 {
		return nil, errors.New("invalid aes key length, should be 16, 24 or 32 bytes")
	}
	return &KeyGenerator{
		AESKey: key,
	}, nil
}

// GenRandomKey get random key
func (g *KeyGenerator) GenRandomKey(id string) (string, error) {
	binkey, err := encryptAESCFB([]byte(id), []byte(g.AESKey))
	if err != nil {
		return "", err
	}
	return hex.EncodeToString(binkey), nil
}

// DecodeIDFromRandomKey get id from encrypt strings
func (g *KeyGenerator) DecodeIDFromRandomKey(encrypted string) (string, error) {
	buf, err := hex.DecodeString(encrypted)
	if err != nil {
		return "", err
	}

	raw, err := decryptAESCFB(buf, []byte(g.AESKey))
	if err != nil {
		return "", err
	}



	return string(raw), nil

}