package utils

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

// 加密密钥（32字节，用于AES-256）
var key = []byte("change-me-MUST-32-BIT-1234567890")

// 自定义Base64编码表，这里排除了'+'和'/'
const customBase64Table = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789-_"

func CustomBase64Encode(src []byte) string {
	// 创建一个使用自定义编码表的Base64编码器
	encoder := base64.NewEncoding(customBase64Table)
	return encoder.EncodeToString(src)
}

func CustomBase64Decode(encoded string) ([]byte, error) {
	// 创建一个使用自定义编码表的Base64解码器
	decoder := base64.NewEncoding(customBase64Table)
	return decoder.DecodeString(encoded)
}

// SetKey updates the current encryption key to the provided byte slice.
// The newKey should be 32 bytes long for AES-256 encryption.
// It modifies the global 'key' variable which is used as the encryption key.
func SetKey(newKey string) error {
	if len(newKey) != 32 {
		return fmt.Errorf("key must be 32 bytes long")
	}
	key = []byte(newKey)
	return nil
}

func SetKeyInBytes(newKey []byte) error {
	if len(newKey) != 32 {
		return fmt.Errorf("key must be 32 bytes long")
	}
	key = newKey
	return nil
}

// Encrypt2 encrypts plaintext using AES encryption.
func Encrypt2(plaintext string) (string, error) {
	block, err := aes.NewCipher(key)
	if err != nil {
		return "", err
	}

	ciphertext := make([]byte, aes.BlockSize+len(plaintext))
	iv := ciphertext[:aes.BlockSize]

	if _, err := io.ReadFull(rand.Reader, iv); err != nil {
		return "", err
	}

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

	return CustomBase64Encode(ciphertext), nil
}

// Decrypt2 decrypts ciphertext back to plaintext.
func Decrypt2(ciphertext string) (string, error) {
	ciphertextBytes, err := CustomBase64Decode(ciphertext)
	if err != nil {
		return "", err
	}

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

	if len(ciphertextBytes) < aes.BlockSize {
		return "", fmt.Errorf("ciphertext too short")
	}
	iv := ciphertextBytes[:aes.BlockSize]
	ciphertextBytes = ciphertextBytes[aes.BlockSize:]

	stream := cipher.NewCFBDecrypter(block, iv)
	stream.XORKeyStream(ciphertextBytes, ciphertextBytes)

	return string(ciphertextBytes), nil
}

// Encode function combines the elements of strings slice into a single string and returns it.
func Encode(data []string) (string, error) {
	for _, element := range data {
		if strings.Contains(element, "||") {
			return "", fmt.Errorf("elements cannot contain '||' character")
		}
	}
	combined := strings.Join(data, "||")
	return Encrypt2(combined)
}

func Decode(encoded string) ([]string, error) {
	decrypted, err := Decrypt2(encoded)
	if err != nil {
		return nil, err
	}
	parts := strings.Split(decrypted, "||")
	return parts, nil
}