package com.ruoyi.security.service.impl;

import java.nio.ByteBuffer;
import java.security.NoSuchAlgorithmException;
import java.security.SecureRandom;
import java.util.Base64;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicReference;

import javax.crypto.Cipher;
import javax.crypto.KeyGenerator;
import javax.crypto.SecretKey;
import javax.crypto.spec.GCMParameterSpec;
import javax.crypto.spec.SecretKeySpec;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Service;

import com.ruoyi.security.service.EncryptionService;

/**
 * AES-256-GCM加密服务实现
 * 
 * @author ruoyi
 */
@Service
public class EncryptionServiceImpl implements EncryptionService {
    
    private static final Logger log = LoggerFactory.getLogger(EncryptionServiceImpl.class);
    
    private static final String ALGORITHM = "AES";
    private static final String TRANSFORMATION = "AES/GCM/NoPadding";
    private static final int GCM_IV_LENGTH = 12;
    private static final int GCM_TAG_LENGTH = 16;
    private static final int KEY_LENGTH = 256;
    
    @Value("${security.encryption.master-key:}")
    private String masterKeyBase64;
    
    private final SecureRandom secureRandom = new SecureRandom();
    private final ConcurrentHashMap<String, SecretKey> keyStore = new ConcurrentHashMap<>();
    private final AtomicReference<String> currentKeyId = new AtomicReference<>("default");
    
    @Override
    public String encrypt(String plaintext) {
        if (plaintext == null || plaintext.isEmpty()) {
            return plaintext;
        }
        
        try {
            byte[] encrypted = encrypt(plaintext.getBytes("UTF-8"));
            return Base64.getEncoder().encodeToString(encrypted);
        } catch (Exception e) {
            log.error("Failed to encrypt string data", e);
            throw new RuntimeException("Encryption failed", e);
        }
    }
    
    @Override
    public String decrypt(String ciphertext) {
        if (ciphertext == null || ciphertext.isEmpty()) {
            return ciphertext;
        }
        
        try {
            byte[] encrypted = Base64.getDecoder().decode(ciphertext);
            byte[] decrypted = decrypt(encrypted);
            return new String(decrypted, "UTF-8");
        } catch (Exception e) {
            log.error("Failed to decrypt string data", e);
            throw new RuntimeException("Decryption failed", e);
        }
    }
    
    @Override
    public byte[] encrypt(byte[] plaintext) {
        if (plaintext == null || plaintext.length == 0) {
            return plaintext;
        }
        
        try {
            SecretKey key = getCurrentKey();
            
            // 生成随机IV
            byte[] iv = new byte[GCM_IV_LENGTH];
            secureRandom.nextBytes(iv);
            
            // 初始化加密器
            Cipher cipher = Cipher.getInstance(TRANSFORMATION);
            GCMParameterSpec gcmSpec = new GCMParameterSpec(GCM_TAG_LENGTH * 8, iv);
            cipher.init(Cipher.ENCRYPT_MODE, key, gcmSpec);
            
            // 执行加密
            byte[] ciphertext = cipher.doFinal(plaintext);
            
            // 组合IV和密文
            ByteBuffer buffer = ByteBuffer.allocate(GCM_IV_LENGTH + ciphertext.length);
            buffer.put(iv);
            buffer.put(ciphertext);
            
            return buffer.array();
        } catch (Exception e) {
            log.error("Failed to encrypt byte array", e);
            throw new RuntimeException("Encryption failed", e);
        }
    }
    
    @Override
    public byte[] decrypt(byte[] ciphertext) {
        if (ciphertext == null || ciphertext.length <= GCM_IV_LENGTH) {
            return ciphertext;
        }
        
        try {
            // 分离IV和密文
            ByteBuffer buffer = ByteBuffer.wrap(ciphertext);
            byte[] iv = new byte[GCM_IV_LENGTH];
            buffer.get(iv);
            
            byte[] encrypted = new byte[buffer.remaining()];
            buffer.get(encrypted);
            
            // 初始化解密器
            SecretKey key = getCurrentKey();
            Cipher cipher = Cipher.getInstance(TRANSFORMATION);
            GCMParameterSpec gcmSpec = new GCMParameterSpec(GCM_TAG_LENGTH * 8, iv);
            cipher.init(Cipher.DECRYPT_MODE, key, gcmSpec);
            
            // 执行解密
            return cipher.doFinal(encrypted);
        } catch (Exception e) {
            log.error("Failed to decrypt byte array", e);
            throw new RuntimeException("Decryption failed", e);
        }
    }
    
    @Override
    public SecretKey generateKey() {
        try {
            KeyGenerator keyGenerator = KeyGenerator.getInstance(ALGORITHM);
            keyGenerator.init(KEY_LENGTH);
            return keyGenerator.generateKey();
        } catch (NoSuchAlgorithmException e) {
            log.error("Failed to generate encryption key", e);
            throw new RuntimeException("Key generation failed", e);
        }
    }
    
    @Override
    public String rotateKey() {
        try {
            SecretKey newKey = generateKey();
            String newKeyId = "key_" + System.currentTimeMillis();
            
            keyStore.put(newKeyId, newKey);
            currentKeyId.set(newKeyId);
            
            log.info("Key rotated successfully, new key ID: {}", newKeyId);
            return newKeyId;
        } catch (Exception e) {
            log.error("Failed to rotate encryption key", e);
            throw new RuntimeException("Key rotation failed", e);
        }
    }
    
    @Override
    public String getCurrentKeyId() {
        return currentKeyId.get();
    }
    
    @Override
    public boolean verifyIntegrity(String data, String signature) {
        if (data == null || signature == null) {
            return false;
        }
        
        try {
            String computedSignature = generateSignature(data);
            return signature.equals(computedSignature);
        } catch (Exception e) {
            log.error("Failed to verify data integrity", e);
            return false;
        }
    }
    
    @Override
    public String generateSignature(String data) {
        if (data == null) {
            return null;
        }
        
        try {
            // 使用HMAC-SHA256生成签名
            javax.crypto.Mac mac = javax.crypto.Mac.getInstance("HmacSHA256");
            SecretKey key = getCurrentKey();
            mac.init(key);
            
            byte[] signature = mac.doFinal(data.getBytes("UTF-8"));
            return Base64.getEncoder().encodeToString(signature);
        } catch (Exception e) {
            log.error("Failed to generate data signature", e);
            throw new RuntimeException("Signature generation failed", e);
        }
    }
    
    /**
     * 获取当前使用的密钥
     */
    private SecretKey getCurrentKey() {
        String keyId = currentKeyId.get();
        SecretKey key = keyStore.get(keyId);
        
        if (key == null) {
            // 如果没有密钥，尝试从配置加载或生成新密钥
            key = loadOrGenerateDefaultKey();
            keyStore.put(keyId, key);
        }
        
        return key;
    }
    
    /**
     * 加载或生成默认密钥
     */
    private SecretKey loadOrGenerateDefaultKey() {
        if (masterKeyBase64 != null && !masterKeyBase64.isEmpty()) {
            try {
                byte[] keyBytes = Base64.getDecoder().decode(masterKeyBase64);
                return new SecretKeySpec(keyBytes, ALGORITHM);
            } catch (Exception e) {
                log.warn("Failed to load master key from configuration, generating new key", e);
            }
        }
        
        log.info("Generating new default encryption key");
        return generateKey();
    }
}