package cn.novalon.manage.gateway.service.impl;

import cn.novalon.manage.gateway.service.JwtKeyService;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Service;

import javax.crypto.Cipher;
import javax.crypto.SecretKey;
import javax.crypto.SecretKeyFactory;
import javax.crypto.spec.GCMParameterSpec;
import javax.crypto.spec.PBEKeySpec;
import javax.crypto.spec.SecretKeySpec;
import java.nio.ByteBuffer;
import java.nio.charset.StandardCharsets;
import java.security.SecureRandom;
import java.security.spec.KeySpec;
import java.util.Base64;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicReference;

@Service
public class JwtKeyServiceImpl implements JwtKeyService {

    private static final Logger logger = LoggerFactory.getLogger(JwtKeyServiceImpl.class);
    
    private static final String KEY_ALGORITHM = "AES";
    private static final String KEY_ENCRYPTION_ALGORITHM = "AES/GCM/NoPadding";
    private static final int GCM_TAG_LENGTH = 128;
    private static final int GCM_IV_LENGTH = 12;
    private static final int KEY_SIZE_BITS = 256;
    private static final int MIN_KEY_LENGTH = 32;
    private static final int KEY_ROTATION_INTERVAL_MS = 30 * 24 * 60 * 60 * 1000; // 30 days
    
    @Value("${jwt.secret:}")
    private String configuredSecret;
    
    @Value("${jwt.key.encryption.password:}")
    private String encryptionPassword;
    
    @Value("${jwt.key.rotation.enabled:true}")
    private boolean rotationEnabled;
    
    private final AtomicReference<String> currentKeyVersion = new AtomicReference<>("v1");
    private final Map<String, SecretKey> keyVersionMap = new ConcurrentHashMap<>();
    private final Map<String, Long> keyCreationTimeMap = new ConcurrentHashMap<>();
    private final SecureRandom secureRandom = new SecureRandom();

    @Override
    public SecretKey getCurrentSigningKey() {
        String version = getCurrentKeyVersion();
        return getSigningKeyByVersion(version);
    }

    @Override
    public SecretKey getSigningKeyByVersion(String version) {
        return keyVersionMap.get(version);
    }

    @Override
    public String getCurrentKeyVersion() {
        return currentKeyVersion.get();
    }

    @Override
    public void rotateKey() {
        if (!rotationEnabled) {
            logger.info("Key rotation is disabled");
            return;
        }

        logger.info("Starting JWT key rotation");
        
        try {
            String newVersion = generateNextVersion();
            String newKey = generateSecureKey();
            
            SecretKey signingKey = new SecretKeySpec(
                    newKey.getBytes(StandardCharsets.UTF_8),
                    KEY_ALGORITHM
            );
            
            keyVersionMap.put(newVersion, signingKey);
            keyCreationTimeMap.put(newVersion, System.currentTimeMillis());
            currentKeyVersion.set(newVersion);
            
            logger.info("JWT key rotated successfully. New version: {}", newVersion);
            
            cleanupOldKeys();
            
        } catch (Exception e) {
            logger.error("Failed to rotate JWT key", e);
            throw new RuntimeException("Key rotation failed", e);
        }
    }

    @Override
    public boolean validateKeyStrength(String key) {
        if (key == null || key.length() < MIN_KEY_LENGTH) {
            logger.warn("Key validation failed: key is null or too short");
            return false;
        }
        
        boolean hasUpperCase = !key.equals(key.toLowerCase());
        boolean hasLowerCase = !key.equals(key.toUpperCase());
        boolean hasDigit = key.matches(".*\\d.*");
        boolean hasSpecialChar = !key.matches("[a-zA-Z0-9]*");
        
        int strengthScore = (hasUpperCase ? 1 : 0) + 
                         (hasLowerCase ? 1 : 0) + 
                         (hasDigit ? 1 : 0) + 
                         (hasSpecialChar ? 1 : 0);
        
        boolean isValid = strengthScore >= 3 && key.length() >= MIN_KEY_LENGTH;
        
        if (!isValid) {
            logger.warn("Key validation failed: strength score = {}, length = {}", strengthScore, key.length());
        }
        
        return isValid;
    }

    @Override
    public String generateSecureKey() {
        byte[] keyBytes = new byte[KEY_SIZE_BITS / 8];
        secureRandom.nextBytes(keyBytes);
        return Base64.getEncoder().encodeToString(keyBytes);
    }

    @Override
    public String encryptKey(String key) {
        if (encryptionPassword == null || encryptionPassword.isEmpty()) {
            logger.warn("Encryption password not configured, returning plain key");
            return key;
        }

        try {
            byte[] iv = new byte[GCM_IV_LENGTH];
            secureRandom.nextBytes(iv);
            
            SecretKey encryptionKey = deriveEncryptionKey(encryptionPassword);
            Cipher cipher = Cipher.getInstance(KEY_ENCRYPTION_ALGORITHM);
            GCMParameterSpec spec = new GCMParameterSpec(GCM_TAG_LENGTH, iv);
            cipher.init(Cipher.ENCRYPT_MODE, encryptionKey, spec);
            
            byte[] encryptedBytes = cipher.doFinal(key.getBytes(StandardCharsets.UTF_8));
            
            ByteBuffer byteBuffer = ByteBuffer.allocate(iv.length + encryptedBytes.length);
            byteBuffer.put(iv);
            byteBuffer.put(encryptedBytes);
            
            String result = Base64.getEncoder().encodeToString(byteBuffer.array());
            logger.debug("Key encrypted successfully");
            return result;
            
        } catch (Exception e) {
            logger.error("Failed to encrypt key", e);
            throw new RuntimeException("Key encryption failed", e);
        }
    }

    @Override
    public String decryptKey(String encryptedKey) {
        if (encryptionPassword == null || encryptionPassword.isEmpty()) {
            logger.warn("Encryption password not configured, returning key as is");
            return encryptedKey;
        }

        try {
            byte[] decodedBytes = Base64.getDecoder().decode(encryptedKey);
            ByteBuffer byteBuffer = ByteBuffer.wrap(decodedBytes);
            
            byte[] iv = new byte[GCM_IV_LENGTH];
            byteBuffer.get(iv);
            
            byte[] encryptedBytes = new byte[byteBuffer.remaining()];
            byteBuffer.get(encryptedBytes);
            
            SecretKey encryptionKey = deriveEncryptionKey(encryptionPassword);
            Cipher cipher = Cipher.getInstance(KEY_ENCRYPTION_ALGORITHM);
            GCMParameterSpec spec = new GCMParameterSpec(GCM_TAG_LENGTH, iv);
            cipher.init(Cipher.DECRYPT_MODE, encryptionKey, spec);
            
            byte[] decryptedBytes = cipher.doFinal(encryptedBytes);
            String result = new String(decryptedBytes, StandardCharsets.UTF_8);
            logger.debug("Key decrypted successfully");
            return result;
            
        } catch (Exception e) {
            logger.error("Failed to decrypt key", e);
            throw new RuntimeException("Key decryption failed", e);
        }
    }

    public void initializeKeys() {
        try {
            String initialKey;
            
            if (configuredSecret != null && !configuredSecret.isEmpty()) {
                if (configuredSecret.startsWith("enc:")) {
                    initialKey = decryptKey(configuredSecret.substring(4));
                    logger.info("Decrypted JWT key from configuration");
                } else {
                    initialKey = configuredSecret;
                    logger.warn("Using plain JWT key from configuration (not recommended)");
                    
                    if (!validateKeyStrength(initialKey)) {
                        logger.error("Configured JWT key does not meet strength requirements");
                        throw new IllegalArgumentException("Weak JWT key configuration");
                    }
                }
            } else {
                initialKey = generateSecureKey();
                logger.info("Generated new secure JWT key");
            }
            
            SecretKey signingKey = new SecretKeySpec(
                    initialKey.getBytes(StandardCharsets.UTF_8),
                    KEY_ALGORITHM
            );
            
            keyVersionMap.put("v1", signingKey);
            keyCreationTimeMap.put("v1", System.currentTimeMillis());
            currentKeyVersion.set("v1");
            
            logger.info("JWT key service initialized with version v1");
            
        } catch (Exception e) {
            logger.error("Failed to initialize JWT keys", e);
            throw new RuntimeException("JWT key initialization failed", e);
        }
    }

    private String generateNextVersion() {
        String currentVersion = getCurrentKeyVersion();
        int versionNumber = Integer.parseInt(currentVersion.substring(1));
        return "v" + (versionNumber + 1);
    }

    private SecretKey deriveEncryptionKey(String password) throws Exception {
        byte[] salt = "NovalonManageSystemSalt".getBytes(StandardCharsets.UTF_8);
        
        KeySpec spec = new PBEKeySpec(
                password.toCharArray(),
                salt,
                65536,
                256
        );
        
        SecretKeyFactory factory = SecretKeyFactory.getInstance("PBKDF2WithHmacSHA256");
        byte[] keyBytes = factory.generateSecret(spec).getEncoded();
        
        return new SecretKeySpec(keyBytes, KEY_ALGORITHM);
    }

    private void cleanupOldKeys() {
        long currentTime = System.currentTimeMillis();
        long rotationThreshold = KEY_ROTATION_INTERVAL_MS * 2; // Keep keys for 2 rotation cycles
        
        keyVersionMap.keySet().stream()
                .filter(version -> !version.equals(getCurrentKeyVersion()))
                .filter(version -> {
                    Long creationTime = keyCreationTimeMap.get(version);
                    return creationTime != null && (currentTime - creationTime) > rotationThreshold;
                })
                .forEach(version -> {
                    keyVersionMap.remove(version);
                    keyCreationTimeMap.remove(version);
                    logger.info("Removed old JWT key version: {}", version);
                });
    }

    public boolean shouldRotateKey() {
        if (!rotationEnabled) {
            return false;
        }
        
        String currentVersion = getCurrentKeyVersion();
        Long creationTime = keyCreationTimeMap.get(currentVersion);
        
        if (creationTime == null) {
            return true;
        }
        
        long keyAge = System.currentTimeMillis() - creationTime;
        return keyAge >= KEY_ROTATION_INTERVAL_MS;
    }
}