package crypto

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

type EncryptType string

const (
	EncryptTypeNone EncryptType = "none"
	EncryptTypeAES  EncryptType = "aes"
	EncryptTypeRC4  EncryptType = "rc4"
)

type CryptoManager struct {
	encryptType EncryptType
	secretKey   string
}

func NewCryptoManager(encryptType EncryptType, secretKey string) *CryptoManager {
	return &CryptoManager{
		encryptType: encryptType,
		secretKey:   secretKey,
	}
}

func (cm *CryptoManager) GetEncryptType() EncryptType {
	return cm.encryptType
}

func (cm *CryptoManager) GetSecretKey() string {
	return cm.secretKey
}

func (cm *CryptoManager) Encrypt(plaintext string) (string, error) {
	if cm.encryptType == EncryptTypeNone || cm.secretKey == "" {
		return plaintext, nil
	}

	switch cm.encryptType {
	case EncryptTypeAES:
		return cm.encryptAES(plaintext)
	case EncryptTypeRC4:
		return cm.encryptRC4(plaintext)
	default:
		return plaintext, nil
	}
}

func (cm *CryptoManager) Decrypt(ciphertext string) (string, error) {
	if cm.encryptType == EncryptTypeNone || cm.secretKey == "" {
		return ciphertext, nil
	}

	switch cm.encryptType {
	case EncryptTypeAES:
		return cm.decryptAES(ciphertext)
	case EncryptTypeRC4:
		return cm.decryptRC4(ciphertext)
	default:
		return ciphertext, nil
	}
}

func (cm *CryptoManager) encryptAES(plaintext string) (string, error) {
	key := []byte(cm.secretKey)
	if len(key) != 16 && len(key) != 24 && len(key) != 32 {
		key = cm.padKey(key, 32)
	}

	block, err := aes.NewCipher(key)
	if err != nil {
		return "", err
	}

	gcm, err := cipher.NewGCM(block)
	if err != nil {
		return "", err
	}

	nonce := make([]byte, gcm.NonceSize())
	if _, err = io.ReadFull(rand.Reader, nonce); err != nil {
		return "", err
	}

	ciphertext := gcm.Seal(nonce, nonce, []byte(plaintext), nil)
	return base64.StdEncoding.EncodeToString(ciphertext), nil
}

func (cm *CryptoManager) decryptAES(ciphertext string) (string, error) {
	key := []byte(cm.secretKey)
	if len(key) != 16 && len(key) != 24 && len(key) != 32 {
		key = cm.padKey(key, 32)
	}

	data, err := base64.StdEncoding.DecodeString(ciphertext)
	if err != nil {
		return "", err
	}

	block, err := aes.NewCipher(key)
	if err != nil {
		return "", err
	}

	gcm, err := cipher.NewGCM(block)
	if err != nil {
		return "", err
	}

	nonceSize := gcm.NonceSize()
	if len(data) < nonceSize {
		return "", errors.New("ciphertext too short")
	}

	nonce, ciphertextBytes := data[:nonceSize], data[nonceSize:]
	plaintext, err := gcm.Open(nil, nonce, ciphertextBytes, nil)
	if err != nil {
		return "", err
	}

	return string(plaintext), nil
}

func (cm *CryptoManager) encryptRC4(plaintext string) (string, error) {
	key := []byte(cm.secretKey)
	if len(key) == 0 {
		return "", errors.New("secret key is empty")
	}

	cipher, err := NewRC4(key)
	if err != nil {
		return "", err
	}

	ciphertext := make([]byte, len(plaintext))
	cipher.XORKeyStream(ciphertext, []byte(plaintext))

	return base64.StdEncoding.EncodeToString(ciphertext), nil
}

func (cm *CryptoManager) decryptRC4(ciphertext string) (string, error) {
	key := []byte(cm.secretKey)
	if len(key) == 0 {
		return "", errors.New("secret key is empty")
	}

	data, err := base64.StdEncoding.DecodeString(ciphertext)
	if err != nil {
		return "", err
	}

	cipher, err := NewRC4(key)
	if err != nil {
		return "", err
	}

	plaintext := make([]byte, len(data))
	cipher.XORKeyStream(plaintext, data)

	return string(plaintext), nil
}

func (cm *CryptoManager) padKey(key []byte, targetLen int) []byte {
	if len(key) >= targetLen {
		return key[:targetLen]
	}

	padded := make([]byte, targetLen)
	copy(padded, key)
	for i := len(key); i < targetLen; i++ {
		padded[i] = key[i%len(key)]
	}
	return padded
}

type RC4 struct {
	s    [256]byte
	i, j uint8
}

func NewRC4(key []byte) (*RC4, error) {
	if len(key) == 0 {
		return nil, errors.New("RC4 key cannot be empty")
	}

	rc4 := &RC4{}
	for i := 0; i < 256; i++ {
		rc4.s[i] = byte(i)
	}

	j := uint8(0)
	for i := 0; i < 256; i++ {
		j = j + rc4.s[i] + key[i%len(key)]
		rc4.s[i], rc4.s[j] = rc4.s[j], rc4.s[i]
	}

	return rc4, nil
}

func (rc4 *RC4) XORKeyStream(dst, src []byte) {
	for i := range src {
		rc4.i++
		rc4.j += rc4.s[rc4.i]
		rc4.s[rc4.i], rc4.s[rc4.j] = rc4.s[rc4.j], rc4.s[rc4.i]
		dst[i] = src[i] ^ rc4.s[(rc4.s[rc4.i]+rc4.s[rc4.j])]
	}
}

func GenerateRandomKey(length int) (string, error) {
	key := make([]byte, length)
	if _, err := io.ReadFull(rand.Reader, key); err != nil {
		return "", err
	}
	return base64.StdEncoding.EncodeToString(key), nil
}

func GenerateAESKey() (string, error) {
	return GenerateRandomKey(32)
}

func GenerateRC4Key() (string, error) {
	return GenerateRandomKey(16)
}