package com.secure.device;

import android.content.Context;
import android.security.keystore.KeyGenParameterSpec;
import android.security.keystore.KeyProperties;
import java.security.KeyPair;
import java.security.KeyPairGenerator;
import java.security.KeyStore;
import java.security.PrivateKey;
import java.security.PublicKey;
import java.security.Signature;
import java.util.Base64;
import javax.crypto.Cipher;
import javax.crypto.KeyGenerator;
import javax.crypto.SecretKey;
import javax.crypto.spec.GCMParameterSpec;

/**
 * 设备端安全模块（Android）
 * 
 * 核心设计：
 * - TEE 密钥对（RSA-2048）：私钥不可导出，用于解密服务端下发数据 + 签名
 * - 每次拍照生成随机 DEK（AES-256）：用后即弃
 * - DEK 明文通过 HTTPS 传给服务端，服务端用 UK 加密后存储
 * 
 * 安全边界：
 * - 私钥永不出 TEE
 * - SN 仅做身份标识，不参与加密
 * - 恢复出厂 → 密钥销毁 → 新密钥对 → 用户短信验证恢复
 */
public class DeviceCrypto {

    // ==================== 常量 ====================

    private static final String KEYSTORE_PROVIDER = "AndroidKeyStore";
    private static final String KEY_ALIAS = "device_tee_key";
    private static final String KEY_ALGO = KeyProperties.KEY_ALGORITHM_RSA;
    private static final int RSA_KEY_SIZE = 2048;
    private static final String SIGN_ALGO = "SHA256withRSA";
    private static final String RSA_TRANSFORM = "RSA/ECB/OAEPWithSHA-256AndMGF1Padding";
    private static final int AES_KEY_SIZE = 256;
    private static final int GCM_IV_LENGTH = 12;
    private static final int GCM_TAG_LENGTH = 128;

    private final Context context;
    private KeyStore keyStore;

    // ==================== 构造与初始化 ====================

    public DeviceCrypto(Context context) {
        this.context = context;
        initKeyStore();
    }

    /**
     * 初始化 AndroidKeyStore，如不存在密钥对则生成
     * 
     * 安全属性：
     * - PURPOSE_DECRYPT + PURPOSE_SIGN（私钥只用于解密下发数据 + 签名）
     * - StrongBox 优先（硬件安全模块）
     * - 不可导出
     */
    private void initKeyStore() {
        try {
            keyStore = KeyStore.getInstance(KEYSTORE_PROVIDER);
            keyStore.load(null);

            if (!keyStore.containsAlias(KEY_ALIAS)) {
                generateTeeKeyPair();
            }
        } catch (Exception e) {
            throw new RuntimeException("KeyStore init failed", e);
        }
    }

    /**
     * 在 TEE / StrongBox 中生成 RSA-2048 密钥对
     */
    private void generateTeeKeyPair() {
        try {
            KeyPairGenerator kpg = KeyPairGenerator.getInstance(
                    KEY_ALGO, KEYSTORE_PROVIDER);

            KeyGenParameterSpec spec = new KeyGenParameterSpec.Builder(
                    KEY_ALIAS,
                    KeyProperties.PURPOSE_DECRYPT |
                    KeyProperties.PURPOSE_SIGN
            )
                    .setKeySize(RSA_KEY_SIZE)
                    .setDigests(KeyProperties.DIGEST_SHA256)
                    .setEncryptionPaddings(KeyProperties.ENCRYPTION_PADDING_RSA_OAEP)
                    .setSignaturePaddings(KeyProperties.SIGNATURE_PADDING_RSA_PKCS1)
                    .setUserAuthenticationRequired(false) // 无登录体系
                    .setIsStrongBoxBacked(true)           // 优先 StrongBox
                    .build();

            kpg.initialize(spec);
            kpg.generateKeyPair();
        } catch (Exception e) {
            throw new RuntimeException("TEE key gen failed", e);
        }
    }

    // ==================== 公钥导出 ====================

    /**
     * 获取设备公钥（Base64 编码）
     * 公钥可导出，用于：
     * - 上传服务端存储
     * - 服务端加密下发数据
     */
    public String getPublicKeyBase64() {
        try {
            KeyStore.PrivateKeyEntry entry =
                    (KeyStore.PrivateKeyEntry) keyStore.getEntry(KEY_ALIAS, null);
            PublicKey publicKey = entry.getCertificate().getPublicKey();
            return Base64.getEncoder().encodeToString(publicKey.getEncoded());
        } catch (Exception e) {
            throw new RuntimeException("Get public key failed", e);
        }
    }

    /**
     * 获取设备 SN
     * 实际实现：android.os.Build.getSerial()
     */
    public String getDeviceSN() {
        // return android.os.Build.getSerial();
        return "SN-DEMO-001";
    }

    // ==================== 数据信封加密 ====================

    /**
     * 加密敏感数据（如照片）
     * 
     * 流程：
     * 1. 生成随机 AES-256 DEK
     * 2. 生成随机 12 字节 IV
     * 3. AES-256-GCM 加密数据
     * 4. 返回密文 + IV + 明文DEK（DEK 由服务端用 UK 加密存储）
     * 
     * @return EncryptedPayload
     */
    public EncryptedPayload encryptData(byte[] plaintext) {
        try {
            // 1. 随机 DEK
            KeyGenerator kg = KeyGenerator.getInstance("AES");
            kg.init(AES_KEY_SIZE);
            SecretKey dek = kg.generateKey();

            // 2. 随机 IV
            byte[] iv = new byte[GCM_IV_LENGTH];
            java.security.SecureRandom random = new java.security.SecureRandom();
            random.nextBytes(iv);

            // 3. AES-256-GCM 加密
            Cipher cipher = Cipher.getInstance("AES/GCM/NoPadding");
            GCMParameterSpec gcmSpec = new GCMParameterSpec(GCM_TAG_LENGTH, iv);
            cipher.init(Cipher.ENCRYPT_MODE, dek, gcmSpec);
            byte[] ciphertext = cipher.doFinal(plaintext);

            return new EncryptedPayload(
                    Base64.getEncoder().encodeToString(ciphertext),
                    Base64.getEncoder().encodeToString(iv),
                    Base64.getEncoder().encodeToString(dek.getEncoded())
            );
        } catch (Exception e) {
            throw new RuntimeException("Encrypt failed", e);
        }
    }

    // ==================== 元数据签名 ====================

    /**
     * 对元数据做 RSA 签名（防伪造、防篡改）
     * 
     * 签名内容示例："SN-xxx|timestamp|photoId"
     * 服务端用设备公钥验签
     */
    public String signMetadata(String metadata) {
        try {
            KeyStore.PrivateKeyEntry entry =
                    (KeyStore.PrivateKeyEntry) keyStore.getEntry(KEY_ALIAS, null);
            PrivateKey privateKey = entry.getPrivateKey();

            Signature sig = Signature.getInstance(SIGN_ALGO);
            sig.initSign(privateKey);
            sig.update(metadata.getBytes("UTF-8"));
            byte[] signature = sig.sign();

            return Base64.getEncoder().encodeToString(signature);
        } catch (Exception e) {
            throw new RuntimeException("Sign failed", e);
        }
    }

    // ==================== 恢复流程 ====================

    /**
     * 恢复出厂后，服务端下发用新公钥加密的 Recovery Token
     * 本方法用 TEE 私钥解密，获取 Token
     * 
     * @param encryptedToken 服务端用设备公钥加密的 Token
     * @return 明文 Token（格式：userId|deviceId|timestamp|nonce）
     */
    public String decryptRecoveryToken(byte[] encryptedToken) {
        try {
            KeyStore.PrivateKeyEntry entry =
                    (KeyStore.PrivateKeyEntry) keyStore.getEntry(KEY_ALIAS, null);
            PrivateKey privateKey = entry.getPrivateKey();

            Cipher cipher = Cipher.getInstance(RSA_TRANSFORM);
            cipher.init(Cipher.DECRYPT_MODE, privateKey);
            byte[] decrypted = cipher.doFinal(encryptedToken);

            return new String(decrypted, "UTF-8");
        } catch (Exception e) {
            throw new RuntimeException("Decrypt recovery token failed", e);
        }
    }

    /**
     * 解密服务端下发的加密 DEK
     * 恢复流程中，服务端用设备公钥加密每个照片的 DEK
     * 
     * @param encryptedDek 加密的 DEK（Base64）
     * @return 明文 DEK 字节
     */
    public byte[] decryptDek(String encryptedDekBase64) {
        try {
            KeyStore.PrivateKeyEntry entry =
                    (KeyStore.PrivateKeyEntry) keyStore.getEntry(KEY_ALIAS, null);
            PrivateKey privateKey = entry.getPrivateKey();

            Cipher cipher = Cipher.getInstance(RSA_TRANSFORM);
            cipher.init(Cipher.DECRYPT_MODE, privateKey);
            return cipher.doFinal(Base64.getDecoder().decode(encryptedDekBase64));
        } catch (Exception e) {
            throw new RuntimeException("Decrypt DEK failed", e);
        }
    }

    /**
     * 用解出的 DEK 解密照片
     */
    public byte[] decryptPhoto(String ciphertextBase64, String ivBase64, byte[] dekBytes) {
        try {
            SecretKey dek = new javax.crypto.spec.SecretKeySpec(dekBytes, "AES");
            byte[] iv = Base64.getDecoder().decode(ivBase64);
            byte[] ciphertext = Base64.getDecoder().decode(ciphertextBase64);

            Cipher cipher = Cipher.getInstance("AES/GCM/NoPadding");
            GCMParameterSpec spec = new GCMParameterSpec(GCM_TAG_LENGTH, iv);
            cipher.init(Cipher.DECRYPT_MODE, dek, spec);

            return cipher.doFinal(ciphertext);
        } catch (Exception e) {
            throw new RuntimeException("Decrypt photo failed", e);
        }
    }

    // ==================== 数据模型 ====================

    /**
     * 加密后的数据载体
     */
    public static class EncryptedPayload {
        public final String ciphertextBase64;  // AES-GCM 密文
        public final String ivBase64;          // GCM IV
        public final String dekBase64;         // 明文 DEK（HTTPS 传输给服务端）

        public EncryptedPayload(String ciphertextBase64, String ivBase64, String dekBase64) {
            this.ciphertextBase64 = ciphertextBase64;
            this.ivBase64 = ivBase64;
            this.dekBase64 = dekBase64;
        }
    }

    /**
     * 设备注册请求体
     */
    public static class DeviceRegisterRequest {
        public String sn;
        public String publicKeyBase64;

        public DeviceRegisterRequest(String sn, String publicKeyBase64) {
            this.sn = sn;
            this.publicKeyBase64 = publicKeyBase64;
        }
    }

    /**
     * 上传加密照片请求体
     */
    public static class UploadPhotoRequest {
        public String sn;
        public String photoId;
        public String ciphertextBase64;
        public String ivBase64;
        public String dekBase64;
        public String metadataSignature;
        public String metadata;

        public UploadPhotoRequest(String sn, String photoId, String ciphertextBase64,
                                  String ivBase64, String dekBase64,
                                  String metadataSignature, String metadata) {
            this.sn = sn;
            this.photoId = photoId;
            this.ciphertextBase64 = ciphertextBase64;
            this.ivBase64 = ivBase64;
            this.dekBase64 = dekBase64;
            this.metadataSignature = metadataSignature;
            this.metadata = metadata;
        }
    }

    /**
     * 恢复请求体（用户已通过短信验证）
     */
    public static class RecoveryRequest {
        public String sn;
        public String userId;
        public String smsCode;
        public String newPublicKeyBase64;

        public RecoveryRequest(String sn, String userId, String smsCode, String newPublicKeyBase64) {
            this.sn = sn;
            this.userId = userId;
            this.smsCode = smsCode;
            this.newPublicKeyBase64 = newPublicKeyBase64;
        }
    }
}