package cryptox

import (
	"crypto/aes"
	"crypto/cipher"
	"crypto/md5"
	"crypto/sha256"
	"encoding/base64"
	"encoding/hex"
	"fmt"
	"time"
)

// generateKey: SHA256(secret)[:32]
func generateKey(secret string) []byte {
	h := sha256.Sum256([]byte(secret))
	return h[:32]
}

// generateIv: SHA256(hex(md5(nonce)) + secret)[:16]
func generateIv(nonce, secret string) []byte {
	md5h := md5.Sum([]byte(nonce))
	md5hex := hex.EncodeToString(md5h[:])
	sha := sha256.Sum256([]byte(md5hex + secret))
	return sha[:16]
}

// Encrypt 加密明文，返回 base64密文 和 nonce（hex(UnixNano)）
func Encrypt(plaintext []byte, secret string) (dataB64 string, nonce string, err error) {
	nonce = fmt.Sprintf("%x", time.Now().UnixNano())
	key := generateKey(secret)
	iv := generateIv(nonce, secret)

	block, err := aes.NewCipher(key)
	if err != nil {
		return
	}
	padded := pkcs7Pad(plaintext, aes.BlockSize)
	mode := cipher.NewCBCEncrypter(block, iv)
	dst := make([]byte, len(padded))
	mode.CryptBlocks(dst, padded)
	dataB64 = base64.StdEncoding.EncodeToString(dst)
	return
}

// Decrypt 解密，nonce 用于派生 IV
func Decrypt(dataB64, secret, nonce string) ([]byte, error) {
	ciphertext, err := base64.StdEncoding.DecodeString(dataB64)
	if err != nil {
		return nil, err
	}
	key := generateKey(secret)
	iv := generateIv(nonce, secret)
	block, err := aes.NewCipher(key)
	if err != nil {
		return nil, err
	}
	if len(ciphertext)%aes.BlockSize != 0 {
		return nil, fmt.Errorf("ciphertext length not multiple of block size")
	}
	mode := cipher.NewCBCDecrypter(block, iv)
	dst := make([]byte, len(ciphertext))
	mode.CryptBlocks(dst, ciphertext)
	return pkcs7Unpad(dst)
}

func pkcs7Pad(data []byte, blockSize int) []byte {
	pad := blockSize - len(data)%blockSize
	padding := make([]byte, pad)
	for i := range padding {
		padding[i] = byte(pad)
	}
	return append(data, padding...)
}

func pkcs7Unpad(data []byte) ([]byte, error) {
	if len(data) == 0 {
		return nil, fmt.Errorf("empty data")
	}
	pad := int(data[len(data)-1])
	if pad == 0 || pad > aes.BlockSize {
		return nil, fmt.Errorf("invalid padding size: %d", pad)
	}
	return data[:len(data)-pad], nil
}