package auth

import (
	"crypto/aes"
	"encoding/base64"
	"errors"
	"fmt"
)

func DecryptCredential(ciphertextBase64, key string) (string, error) {
	ciphertext, err := base64.StdEncoding.DecodeString(ciphertextBase64)
	if err != nil {
		return "", err
	}
	if len(ciphertext) == 0 || len(ciphertext)%aes.BlockSize != 0 {
		return "", errors.New("invalid AES ciphertext length")
	}

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

	plain := make([]byte, len(ciphertext))
	for start := 0; start < len(ciphertext); start += aes.BlockSize {
		block.Decrypt(plain[start:start+aes.BlockSize], ciphertext[start:start+aes.BlockSize])
	}

	plain, err = unpadPKCS7(plain, aes.BlockSize)
	if err != nil {
		return "", err
	}
	return string(plain), nil
}

func EncryptCredential(plaintext, key string) (string, error) {
	block, err := aes.NewCipher(repeatKey(key, aes.BlockSize))
	if err != nil {
		return "", err
	}
	plain := padPKCS7([]byte(plaintext), aes.BlockSize)
	out := make([]byte, len(plain))
	for start := 0; start < len(plain); start += aes.BlockSize {
		block.Encrypt(out[start:start+aes.BlockSize], plain[start:start+aes.BlockSize])
	}
	return base64.StdEncoding.EncodeToString(out), nil
}

func repeatKey(key string, size int) []byte {
	src := []byte(key)
	out := make([]byte, size)
	if len(src) == 0 {
		return out
	}
	for i := range out {
		out[i] = src[i%len(src)]
	}
	return out
}

func unpadPKCS7(in []byte, blockSize int) ([]byte, error) {
	if len(in) == 0 || len(in)%blockSize != 0 {
		return nil, errors.New("invalid padded data length")
	}
	pad := int(in[len(in)-1])
	if pad == 0 || pad > blockSize || pad > len(in) {
		return nil, fmt.Errorf("invalid padding size %d", pad)
	}
	for _, b := range in[len(in)-pad:] {
		if int(b) != pad {
			return nil, errors.New("invalid padding bytes")
		}
	}
	return in[:len(in)-pad], nil
}

func padPKCS7(in []byte, blockSize int) []byte {
	pad := blockSize - len(in)%blockSize
	out := make([]byte, len(in)+pad)
	copy(out, in)
	for i := len(in); i < len(out); i++ {
		out[i] = byte(pad)
	}
	return out
}