package crypto

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

var (
	ErrInvalidKey        = errors.New("invalid encryption key: must be 32 bytes")
	ErrInvalidCiphertext = errors.New("invalid ciphertext: too short")
	ErrDecryptFailed     = errors.New("decryption failed")
)

// MessageEncryptor 消息加密器
// 使用 AES-256-GCM 算法进行加密
type MessageEncryptor struct {
	key        []byte
	gcm        cipher.AEAD
	keyVersion int
	mu         sync.RWMutex
}

// globalEncryptor 全局加密器实例
var globalEncryptor *MessageEncryptor
var encryptorOnce sync.Once

// InitMessageEncryptor 初始化全局消息加密器
// key 必须是32字节（256位）的密钥
func InitMessageEncryptor(key string, keyVersion int) error {
	keyBytes := []byte(key)
	if len(keyBytes) != 32 {
		return ErrInvalidKey
	}

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

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

	encryptorOnce.Do(func() {
		globalEncryptor = &MessageEncryptor{
			key:        keyBytes,
			gcm:        gcm,
			keyVersion: keyVersion,
		}
	})

	return nil
}

// GetMessageEncryptor 获取全局加密器实例
func GetMessageEncryptor() *MessageEncryptor {
	return globalEncryptor
}

// Encrypt 加密数据
// 返回 base64 编码的密文（格式：nonce + ciphertext + tag）
func (e *MessageEncryptor) Encrypt(plaintext []byte) (string, error) {
	if e == nil {
		return "", errors.New("encryptor not initialized")
	}

	if len(plaintext) == 0 {
		return "", nil
	}

	e.mu.RLock()
	defer e.mu.RUnlock()

	// 生成随机nonce（12字节）
	nonce := make([]byte, e.gcm.NonceSize())
	if _, err := io.ReadFull(rand.Reader, nonce); err != nil {
		return "", err
	}

	// 加密数据，nonce附加在密文前面
	ciphertext := e.gcm.Seal(nonce, nonce, plaintext, nil)

	// 返回base64编码的密文
	return base64.StdEncoding.EncodeToString(ciphertext), nil
}

// Decrypt 解密数据
// 输入 base64 编码的密文
func (e *MessageEncryptor) Decrypt(ciphertextBase64 string) ([]byte, error) {
	if e == nil {
		return nil, errors.New("encryptor not initialized")
	}

	if ciphertextBase64 == "" {
		return nil, nil
	}

	e.mu.RLock()
	defer e.mu.RUnlock()

	// base64解码
	ciphertext, err := base64.StdEncoding.DecodeString(ciphertextBase64)
	if err != nil {
		return nil, err
	}

	// 检查密文长度
	nonceSize := e.gcm.NonceSize()
	if len(ciphertext) < nonceSize {
		return nil, ErrInvalidCiphertext
	}

	// 提取nonce和实际密文
	nonce := ciphertext[:nonceSize]
	actualCiphertext := ciphertext[nonceSize:]

	// 解密
	plaintext, err := e.gcm.Open(nil, nonce, actualCiphertext, nil)
	if err != nil {
		return nil, ErrDecryptFailed
	}

	return plaintext, nil
}

// GetKeyVersion 获取当前密钥版本
func (e *MessageEncryptor) GetKeyVersion() int {
	if e == nil {
		return 0
	}
	e.mu.RLock()
	defer e.mu.RUnlock()
	return e.keyVersion
}

// RotateKey 密钥轮换（用于密钥升级）
// 新密钥必须也是32字节
func (e *MessageEncryptor) RotateKey(newKey string, newVersion int) error {
	keyBytes := []byte(newKey)
	if len(keyBytes) != 32 {
		return ErrInvalidKey
	}

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

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

	e.mu.Lock()
	defer e.mu.Unlock()

	e.key = keyBytes
	e.gcm = gcm
	e.keyVersion = newVersion

	return nil
}