package utils

import (
	"bytes"
	"crypto/aes"
	"crypto/cipher"
	"fmt"
	"os"
)

var keys = map[string]string{
	"TW": "mikumikulukaluka",
	"JP": "mikumikulukaluka",
	"GL": "RXzFAC87WR28eNcP",
	"CN": "wakakabiliwakaka",
}

var ivs = map[string]string{
	"TW": "lukalukamikumiku",
	"JP": "lukalukamikumiku",
	"GL": "hrk0k7knXlit3S_y",
	"CN": "biliwakakawohaha",
}

func getKeyIV() ([]byte, []byte, error) {
	server := os.Getenv("SERVER")
	key, ok := keys[server]
	if !ok {
		return nil, nil, fmt.Errorf("unknown server: %s", server)
	}
	return []byte(key), []byte(ivs[server]), nil
}

func Decrypt(buf []byte) ([]byte, error) {
	key, iv, err := getKeyIV()
	if err != nil {
		return nil, err
	}

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

	if len(buf)%aes.BlockSize != 0 {
		return nil, fmt.Errorf("invalid block size")
	}

	decrypted := make([]byte, len(buf))
	cipher.NewCBCDecrypter(block, iv).CryptBlocks(decrypted, buf)

	// 手動移除 padding（setAutoPadding(false) 對應）
	padLen := int(decrypted[len(decrypted)-1])
	if padLen > 0 && padLen <= 16 && padLen <= len(decrypted) {
		decrypted = decrypted[:len(decrypted)-padLen]
	}

	return decrypted, nil
}

func Encrypt(buf []byte) ([]byte, error) {
	key, iv, err := getKeyIV()
	if err != nil {
		return nil, err
	}

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

	// 手動加 PKCS7 padding（setAutoPadding(true) 對應）
	blockSize := block.BlockSize()
	padding := blockSize - len(buf)%blockSize
	padded := append(buf, bytes.Repeat([]byte{byte(padding)}, padding)...)

	encrypted := make([]byte, len(padded))
	cipher.NewCBCEncrypter(block, iv).CryptBlocks(encrypted, padded)

	return encrypted, nil
}