package tech.easyflow.auth.service.impl;

import com.alibaba.fastjson2.JSON;
import com.alibaba.fastjson2.JSONObject;
import com.mybatisflex.core.query.QueryWrapper;
import com.mybatisflex.spring.service.impl.ServiceImpl;
import jakarta.annotation.PostConstruct;
import org.springframework.data.redis.core.StringRedisTemplate;
import org.springframework.data.redis.core.script.DefaultRedisScript;
import org.springframework.scheduling.annotation.Scheduled;
import org.springframework.stereotype.Service;
import org.springframework.util.StringUtils;
import tech.easyflow.auth.config.CredentialKeyProperties;
import tech.easyflow.auth.constant.CredentialKeyStatus;
import tech.easyflow.auth.entity.AuthCredentialKey;
import tech.easyflow.auth.entity.CredentialKeyTransitionVO;
import tech.easyflow.auth.entity.CredentialKeyVO;
import tech.easyflow.auth.entity.EncryptedCredentialDTO;
import tech.easyflow.auth.mapper.AuthCredentialKeyMapper;
import tech.easyflow.auth.service.AuthCredentialKeyService;
import tech.easyflow.common.cache.RedisLockExecutor;
import tech.easyflow.common.web.exceptions.BusinessException;

import javax.crypto.Cipher;
import javax.crypto.spec.GCMParameterSpec;
import javax.crypto.spec.SecretKeySpec;
import java.nio.charset.StandardCharsets;
import java.security.*;
import java.security.spec.MGF1ParameterSpec;
import java.security.spec.PKCS8EncodedKeySpec;
import java.time.Duration;
import java.util.*;

/**
 * 认证传输密钥服务实现。
 */
@Service
public class AuthCredentialKeyServiceImpl
        extends ServiceImpl<AuthCredentialKeyMapper, AuthCredentialKey>
        implements AuthCredentialKeyService {

    private static final String MASTER_KEY_ENV = "EASYFLOW_CREDENTIAL_MASTER_KEY";
    private static final String ALGORITHM = "RSA-OAEP-256 + AES-256-GCM";
    private static final String LOCK_KEY = "easyflow:auth:credential-key:lock";
    private static final String NONCE_KEY_PREFIX = "easyflow:auth:credential-key:nonce:";
    private static final Duration LOCK_WAIT = Duration.ofSeconds(5);
    private static final Duration LOCK_LEASE = Duration.ofSeconds(30);
    private static final int AES_GCM_TAG_BITS = 128;
    private static final int RSA_KEY_SIZE = 2048;
    private static final DefaultRedisScript<String> CONSUME_NONCE_SCRIPT;
    private static final SecureRandom SECURE_RANDOM = new SecureRandom();

    static {
        CONSUME_NONCE_SCRIPT = new DefaultRedisScript<>();
        CONSUME_NONCE_SCRIPT.setScriptText(
                "local value = redis.call('get', KEYS[1]); " +
                        "if value then redis.call('del', KEYS[1]); return value; end; " +
                        "return nil;"
        );
        CONSUME_NONCE_SCRIPT.setResultType(String.class);
    }

    private final CredentialKeyProperties properties;
    private final RedisLockExecutor redisLockExecutor;
    private final StringRedisTemplate stringRedisTemplate;
    private byte[] masterKey;

    /**
     * 创建认证传输密钥服务。
     *
     * @param properties 密钥配置
     * @param redisLockExecutor Redis 分布式锁执行器
     * @param stringRedisTemplate Redis 字符串模板
     */
    public AuthCredentialKeyServiceImpl(CredentialKeyProperties properties,
                                        RedisLockExecutor redisLockExecutor,
                                        StringRedisTemplate stringRedisTemplate) {
        this.properties = properties;
        this.redisLockExecutor = redisLockExecutor;
        this.stringRedisTemplate = stringRedisTemplate;
    }

    /**
     * 初始化主密钥并确保密钥环可用。
     */
    @PostConstruct
    public void init() {
        this.masterKey = loadMasterKey();
        ensureActiveKey();
    }

    /**
     * 周期维护密钥环。
     */
    @Override
    @Scheduled(fixedDelayString = "#{@credentialKeyProperties.checkInterval.toMillis()}",
            initialDelay = 60000L)
    public void maintainKeyRing() {
        ensureActiveKey();
        retireExpiredPreviousKeys();
        AuthCredentialKey active = getActiveKey();
        if (active != null && active.getRotateAt() != null
                && active.getRotateAt().getTime() <= System.currentTimeMillis()) {
            rotate(false);
        }
    }

    /**
     * 获取当前公钥材料。
     *
     * @return 公钥材料
     */
    @Override
    public CredentialKeyVO getCurrentCredentialKey() {
        AuthCredentialKey active = ensureActiveKey();
        String nonce = issueNonce(active.getKeyId());
        CredentialKeyVO vo = new CredentialKeyVO();
        vo.setKeyId(active.getKeyId());
        vo.setPublicKey(active.getPublicKeyPem());
        vo.setFingerprint(active.getFingerprint());
        vo.setAlgorithm(ALGORITHM);
        vo.setNonce(nonce);
        vo.setExpiresAt(new Date(System.currentTimeMillis() + safeDuration(properties.getMaterialTtl()).toMillis()));
        vo.setTransitions(getTransitions());
        return vo;
    }

    /**
     * 解密前端加密信封。
     *
     * @param encryptedCredential 加密信封
     * @return 明文 JSON
     */
    @Override
    public JSONObject decryptPayload(EncryptedCredentialDTO encryptedCredential) {
        validateEnvelope(encryptedCredential);
        AuthCredentialKey key = getByKeyId(encryptedCredential.getKeyId());
        validateDecryptableKey(key);
        consumeNonce(encryptedCredential.getNonce(), encryptedCredential.getKeyId());
        try {
            PrivateKey privateKey = parsePrivateKey(decryptPrivateKeyPem(key.getEncryptedPrivateKey()));
            byte[] aesKey = decryptAesKey(privateKey, decodeBase64(encryptedCredential.getEncryptedKey()));
            byte[] plaintext = decryptCiphertext(aesKey, decodeBase64(encryptedCredential.getIv()),
                    decodeBase64(encryptedCredential.getCiphertext()));
            JSONObject payload = JSON.parseObject(new String(plaintext, StandardCharsets.UTF_8));
            if (payload == null || payload.isEmpty()) {
                throw new BusinessException("登录安全载荷为空");
            }
            if (!encryptedCredential.getNonce().equals(payload.getString("nonce"))) {
                throw new BusinessException("登录安全随机数不匹配");
            }
            return payload;
        } catch (BusinessException e) {
            throw e;
        } catch (Exception e) {
            throw new BusinessException("登录安全载荷解密失败");
        }
    }

    /**
     * 手动触发密钥轮换。
     *
     * @param emergency 是否紧急吊销旧密钥
     */
    @Override
    public void rotate(boolean emergency) {
        redisLockExecutor.executeWithLock(LOCK_KEY, LOCK_WAIT, LOCK_LEASE, () -> {
            AuthCredentialKey active = getActiveKey();
            if (active == null) {
                ensureActiveKey();
                return;
            }
            AuthCredentialKey newKey = createCredentialKey(active, active.getRotateSeq() + 1);
            save(newKey);
            if (emergency) {
                markOldKey(active, CredentialKeyStatus.REVOKED, null);
            } else {
                markOldKey(active, CredentialKeyStatus.PREVIOUS,
                        new Date(System.currentTimeMillis() + safeDuration(properties.getPreviousRetain()).toMillis()));
            }
        });
    }

    private AuthCredentialKey ensureActiveKey() {
        AuthCredentialKey active = getActiveKey();
        if (active != null) {
            return active;
        }
        redisLockExecutor.executeWithLock(LOCK_KEY, LOCK_WAIT, LOCK_LEASE, () -> {
            if (getActiveKey() == null) {
                save(createCredentialKey(null, nextRotateSeq()));
            }
        });
        active = getActiveKey();
        if (active == null) {
            throw new IllegalStateException("认证传输密钥初始化失败");
        }
        return active;
    }

    private AuthCredentialKey createCredentialKey(AuthCredentialKey parent, long rotateSeq) {
        try {
            KeyPairGenerator generator = KeyPairGenerator.getInstance("RSA");
            generator.initialize(RSA_KEY_SIZE);
            KeyPair keyPair = generator.generateKeyPair();
            Date now = new Date();
            Date expiresAt = new Date(now.getTime() + safeDuration(properties.getLifetime()).toMillis());
            Date rotateAt = new Date(expiresAt.getTime() - safeDuration(properties.getRotateBefore()).toMillis());
            String publicKeyPem = toPem("PUBLIC KEY", keyPair.getPublic().getEncoded());
            String privateKeyPem = toPem("PRIVATE KEY", keyPair.getPrivate().getEncoded());
            AuthCredentialKey key = new AuthCredentialKey();
            key.setKeyId(UUID.randomUUID().toString().replace("-", ""));
            key.setPublicKeyPem(publicKeyPem);
            key.setEncryptedPrivateKey(encryptPrivateKeyPem(privateKeyPem));
            key.setFingerprint(fingerprint(keyPair.getPublic()));
            key.setStatus(CredentialKeyStatus.ACTIVE);
            key.setRotateSeq(rotateSeq);
            key.setNotBefore(now);
            key.setExpiresAt(expiresAt);
            key.setRotateAt(rotateAt);
            key.setCreated(now);
            key.setModified(now);
            if (parent != null) {
                key.setParentKeyId(parent.getKeyId());
                String payload = buildTransitionPayload(parent, key);
                key.setTransitionPayload(payload);
                key.setTransitionSignature(signTransition(parent, payload));
            }
            return key;
        } catch (Exception e) {
            throw new IllegalStateException("生成认证传输密钥失败", e);
        }
    }

    private String buildTransitionPayload(AuthCredentialKey parent, AuthCredentialKey key) {
        Map<String, Object> payload = new LinkedHashMap<>();
        payload.put("oldKeyId", parent.getKeyId());
        payload.put("newKeyId", key.getKeyId());
        payload.put("newPublicKey", key.getPublicKeyPem());
        payload.put("newFingerprint", key.getFingerprint());
        payload.put("issuedAt", new Date().getTime());
        payload.put("expiresAt", key.getExpiresAt().getTime());
        payload.put("rotateSeq", key.getRotateSeq());
        return Base64.getUrlEncoder().withoutPadding()
                .encodeToString(JSON.toJSONString(payload).getBytes(StandardCharsets.UTF_8));
    }

    private String signTransition(AuthCredentialKey parent, String payload) throws Exception {
        PrivateKey privateKey = parsePrivateKey(decryptPrivateKeyPem(parent.getEncryptedPrivateKey()));
        Signature signature = Signature.getInstance("SHA256withRSA");
        signature.initSign(privateKey);
        signature.update(payload.getBytes(StandardCharsets.UTF_8));
        return Base64.getEncoder().encodeToString(signature.sign());
    }

    private AuthCredentialKey getActiveKey() {
        QueryWrapper wrapper = QueryWrapper.create()
                .eq(AuthCredentialKey::getStatus, CredentialKeyStatus.ACTIVE)
                .orderBy("rotate_seq desc");
        List<AuthCredentialKey> keys = list(wrapper);
        return keys == null || keys.isEmpty() ? null : keys.get(0);
    }

    private AuthCredentialKey getByKeyId(String keyId) {
        QueryWrapper wrapper = QueryWrapper.create().eq(AuthCredentialKey::getKeyId, keyId);
        AuthCredentialKey key = getOne(wrapper);
        if (key == null) {
            throw new BusinessException("登录安全密钥不存在");
        }
        return key;
    }

    private List<CredentialKeyTransitionVO> getTransitions() {
        QueryWrapper wrapper = QueryWrapper.create()
                .isNotNull(AuthCredentialKey::getTransitionPayload)
                .ne(AuthCredentialKey::getStatus, CredentialKeyStatus.REVOKED)
                .orderBy("rotate_seq asc");
        List<AuthCredentialKey> keys = list(wrapper);
        List<CredentialKeyTransitionVO> transitions = new ArrayList<>();
        if (keys == null) {
            return transitions;
        }
        for (AuthCredentialKey key : keys) {
            CredentialKeyTransitionVO transition = new CredentialKeyTransitionVO();
            transition.setFromKeyId(key.getParentKeyId());
            transition.setToKeyId(key.getKeyId());
            transition.setPayload(key.getTransitionPayload());
            transition.setSignature(key.getTransitionSignature());
            transitions.add(transition);
        }
        return transitions;
    }

    private void markOldKey(AuthCredentialKey oldKey, String status, Date previousUntil) {
        AuthCredentialKey update = new AuthCredentialKey();
        update.setId(oldKey.getId());
        update.setStatus(status);
        update.setPreviousUntil(previousUntil);
        update.setModified(new Date());
        updateById(update);
    }

    private void retireExpiredPreviousKeys() {
        QueryWrapper wrapper = QueryWrapper.create()
                .eq(AuthCredentialKey::getStatus, CredentialKeyStatus.PREVIOUS)
                .isNotNull(AuthCredentialKey::getPreviousUntil);
        List<AuthCredentialKey> keys = list(wrapper);
        if (keys == null || keys.isEmpty()) {
            return;
        }
        long now = System.currentTimeMillis();
        for (AuthCredentialKey key : keys) {
            if (key.getPreviousUntil() != null && key.getPreviousUntil().getTime() <= now) {
                markOldKey(key, CredentialKeyStatus.RETIRED, key.getPreviousUntil());
            }
        }
    }

    private long nextRotateSeq() {
        List<AuthCredentialKey> keys = list();
        if (keys == null || keys.isEmpty()) {
            return 1L;
        }
        return keys.stream()
                .map(AuthCredentialKey::getRotateSeq)
                .filter(item -> item != null)
                .max(Long::compareTo)
                .orElse(0L) + 1L;
    }

    private void validateEnvelope(EncryptedCredentialDTO encryptedCredential) {
        if (encryptedCredential == null
                || !StringUtils.hasText(encryptedCredential.getKeyId())
                || !StringUtils.hasText(encryptedCredential.getEncryptedKey())
                || !StringUtils.hasText(encryptedCredential.getIv())
                || !StringUtils.hasText(encryptedCredential.getCiphertext())
                || !StringUtils.hasText(encryptedCredential.getNonce())) {
            throw new BusinessException("登录安全载荷不完整");
        }
    }

    private void validateDecryptableKey(AuthCredentialKey key) {
        long now = System.currentTimeMillis();
        if (CredentialKeyStatus.ACTIVE.equals(key.getStatus())) {
            if (key.getExpiresAt() != null && key.getExpiresAt().getTime() <= now) {
                throw new BusinessException("登录安全密钥已过期");
            }
            return;
        }
        if (CredentialKeyStatus.PREVIOUS.equals(key.getStatus())
                && key.getPreviousUntil() != null
                && key.getPreviousUntil().getTime() > now) {
            return;
        }
        throw new BusinessException("登录安全密钥不可用");
    }

    private String issueNonce(String keyId) {
        String nonce = UUID.randomUUID().toString().replace("-", "");
        stringRedisTemplate.opsForValue().set(nonceKey(nonce), keyId, safeDuration(properties.getNonceTtl()));
        return nonce;
    }

    private void consumeNonce(String nonce, String keyId) {
        String storedKeyId = stringRedisTemplate.execute(CONSUME_NONCE_SCRIPT,
                Collections.singletonList(nonceKey(nonce)));
        if (!keyId.equals(storedKeyId)) {
            throw new BusinessException("登录安全随机数无效或已过期");
        }
    }

    private String nonceKey(String nonce) {
        return NONCE_KEY_PREFIX + nonce;
    }

    private byte[] loadMasterKey() {
        String raw = System.getenv(MASTER_KEY_ENV);
        if (!StringUtils.hasText(raw)) {
            throw new IllegalStateException(MASTER_KEY_ENV + " 未配置");
        }
        byte[] decoded;
        try {
            decoded = Base64.getDecoder().decode(raw);
        } catch (IllegalArgumentException e) {
            throw new IllegalStateException(MASTER_KEY_ENV + " 必须是 Base64 编码", e);
        }
        if (decoded.length != 32) {
            throw new IllegalStateException(MASTER_KEY_ENV + " 必须解码为 32 字节");
        }
        return decoded;
    }

    private String encryptPrivateKeyPem(String privateKeyPem) throws Exception {
        byte[] iv = new byte[12];
        SECURE_RANDOM.nextBytes(iv);
        Cipher cipher = Cipher.getInstance("AES/GCM/NoPadding");
        cipher.init(Cipher.ENCRYPT_MODE, new SecretKeySpec(masterKey, "AES"), new GCMParameterSpec(AES_GCM_TAG_BITS, iv));
        byte[] ciphertext = cipher.doFinal(privateKeyPem.getBytes(StandardCharsets.UTF_8));
        return Base64.getEncoder().encodeToString(iv) + ":" + Base64.getEncoder().encodeToString(ciphertext);
    }

    private String decryptPrivateKeyPem(String encryptedPrivateKey) throws Exception {
        String[] parts = encryptedPrivateKey.split(":", 2);
        if (parts.length != 2) {
            throw new IllegalArgumentException("私钥密文格式无效");
        }
        byte[] iv = decodeBase64(parts[0]);
        byte[] ciphertext = decodeBase64(parts[1]);
        Cipher cipher = Cipher.getInstance("AES/GCM/NoPadding");
        cipher.init(Cipher.DECRYPT_MODE, new SecretKeySpec(masterKey, "AES"), new GCMParameterSpec(AES_GCM_TAG_BITS, iv));
        return new String(cipher.doFinal(ciphertext), StandardCharsets.UTF_8);
    }

    private byte[] decryptAesKey(PrivateKey privateKey, byte[] encryptedKey) throws Exception {
        Cipher cipher = Cipher.getInstance("RSA/ECB/OAEPWithSHA-256AndMGF1Padding");
        cipher.init(Cipher.DECRYPT_MODE, privateKey,
                new javax.crypto.spec.OAEPParameterSpec("SHA-256", "MGF1",
                        MGF1ParameterSpec.SHA256, javax.crypto.spec.PSource.PSpecified.DEFAULT));
        return cipher.doFinal(encryptedKey);
    }

    private byte[] decryptCiphertext(byte[] aesKey, byte[] iv, byte[] ciphertext) throws Exception {
        Cipher cipher = Cipher.getInstance("AES/GCM/NoPadding");
        cipher.init(Cipher.DECRYPT_MODE, new SecretKeySpec(aesKey, "AES"), new GCMParameterSpec(AES_GCM_TAG_BITS, iv));
        return cipher.doFinal(ciphertext);
    }

    private PrivateKey parsePrivateKey(String privateKeyPem) throws Exception {
        byte[] der = parsePem(privateKeyPem);
        return KeyFactory.getInstance("RSA").generatePrivate(new PKCS8EncodedKeySpec(der));
    }

    private String fingerprint(PublicKey publicKey) throws Exception {
        byte[] digest = MessageDigest.getInstance("SHA-256").digest(publicKey.getEncoded());
        StringBuilder builder = new StringBuilder(digest.length * 2);
        for (byte item : digest) {
            builder.append(String.format("%02x", item));
        }
        return builder.toString();
    }

    private String toPem(String type, byte[] der) {
        String body = Base64.getMimeEncoder(64, "\n".getBytes(StandardCharsets.UTF_8)).encodeToString(der);
        return "-----BEGIN " + type + "-----\n" + body + "\n-----END " + type + "-----";
    }

    private byte[] parsePem(String pem) {
        String base64 = pem
                .replace("-----BEGIN PRIVATE KEY-----", "")
                .replace("-----END PRIVATE KEY-----", "")
                .replace("-----BEGIN PUBLIC KEY-----", "")
                .replace("-----END PUBLIC KEY-----", "")
                .replaceAll("\\s", "");
        return Base64.getDecoder().decode(base64);
    }

    private byte[] decodeBase64(String value) {
        try {
            return Base64.getDecoder().decode(value);
        } catch (IllegalArgumentException e) {
            return Base64.getUrlDecoder().decode(value);
        }
    }

    private Duration safeDuration(Duration duration) {
        if (duration == null || duration.isNegative() || duration.isZero()) {
            throw new IllegalStateException("认证传输密钥配置必须为正数");
        }
        return duration;
    }
}
