//! 存储模块 - 完整的KeyStore实现

use crate::key_manager::KeyPair;
use aes_gcm::{
    aead::{Aead, KeyInit, OsRng},
    Aes256Gcm, Nonce,
};
use pbkdf2::pbkdf2_hmac;
use serde::{Deserialize, Serialize};
use sha2::Sha256;
use std::path::Path;
use zeroize::Zeroize;

/// 密钥库 - 加密存储密钥对
#[derive(Debug, Clone)]
pub struct KeyStore {
    /// 地址（32字节）
    pub address: [u8; 32],
    /// 加密的密钥数据
    encrypted_data: Vec<u8>,
    /// 加密参数
    params: EncryptionParams,
}

/// 加密参数
#[derive(Debug, Clone, Serialize, Deserialize)]
struct EncryptionParams {
    /// 盐（32字节）
    salt: Vec<u8>,
    /// IV/Nonce（12字节，用于AES-GCM）
    nonce: Vec<u8>,
    /// PBKDF2迭代次数
    iterations: u32,
}

/// 密钥库文件格式
#[derive(Debug, Serialize, Deserialize)]
struct KeyStoreFile {
    /// 版本号
    version: u32,
    /// 地址（十六进制）
    address: String,
    /// 加密数据（Base64）
    ciphertext: String,
    /// 加密参数
    crypto: CryptoParamsFile,
}

/// 加密参数文件格式
#[derive(Debug, Serialize, Deserialize)]
struct CryptoParamsFile {
    /// 加密算法
    cipher: String,
    /// Nonce（Base64）
    nonce: String,
    /// KDF算法
    kdf: String,
    /// KDF参数
    kdf_params: KDFParamsFile,
}

/// KDF参数文件格式
#[derive(Debug, Serialize, Deserialize)]
struct KDFParamsFile {
    /// 盐（Base64）
    salt: String,
    /// 迭代次数
    iterations: u32,
}

impl KeyStore {
    /// 加密密钥对
    ///
    /// # Arguments
    ///
    /// * `keypair` - 密钥对
    /// * `password` - 密码
    ///
    /// # Returns
    ///
    /// 加密后的密钥库
    pub fn encrypt(keypair: &KeyPair, password: &str) -> Result<Self, Box<dyn std::error::Error>> {
        // 生成随机盐（32字节）
        let mut salt = vec![0u8; 32];
        use rand::RngCore;
        OsRng.fill_bytes(&mut salt);
        
        // 生成随机nonce（12字节，AES-GCM标准）
        let mut nonce_bytes = vec![0u8; 12];
        OsRng.fill_bytes(&mut nonce_bytes);
        
        // 使用PBKDF2派生加密密钥
        let iterations = 100_000u32;
        let mut key = [0u8; 32];
        pbkdf2_hmac::<Sha256>(password.as_bytes(), &salt, iterations, &mut key);
        
        // 序列化密钥对
        let keypair_bytes = bincode::serialize(keypair)?;
        
        // 使用AES-256-GCM加密
        let cipher = Aes256Gcm::new_from_slice(&key)?;
        let nonce = Nonce::from_slice(&nonce_bytes);
        let encrypted_data = cipher
            .encrypt(nonce, keypair_bytes.as_ref())
            .map_err(|e| format!("Encryption failed: {}", e))?;
        
        // 清除敏感数据
        key.zeroize();
        
        // 计算地址
        let pubkey_hash = keypair.public_key_hash();
        let mut address = [0u8; 32];
        address[6..].copy_from_slice(&pubkey_hash);
        
        Ok(Self {
            address,
            encrypted_data,
            params: EncryptionParams {
                salt,
                nonce: nonce_bytes,
                iterations,
            },
        })
    }
    
    /// 解密密钥对
    ///
    /// # Arguments
    ///
    /// * `password` - 密码
    ///
    /// # Returns
    ///
    /// 解密后的密钥对
    pub fn decrypt(&self, password: &str) -> Result<KeyPair, Box<dyn std::error::Error>> {
        // 使用PBKDF2派生加密密钥
        let mut key = [0u8; 32];
        pbkdf2_hmac::<Sha256>(
            password.as_bytes(),
            &self.params.salt,
            self.params.iterations,
            &mut key,
        );
        
        // 使用AES-256-GCM解密
        let cipher = Aes256Gcm::new_from_slice(&key)?;
        let nonce = Nonce::from_slice(&self.params.nonce);
        let decrypted_data = cipher
            .decrypt(nonce, self.encrypted_data.as_ref())
            .map_err(|_| "Decryption failed: wrong password or corrupted data")?;
        
        // 清除敏感数据
        key.zeroize();
        
        // 反序列化密钥对
        let keypair: KeyPair = bincode::deserialize(&decrypted_data)?;
        
        Ok(keypair)
    }
    
    /// 保存到文件
    ///
    /// # Arguments
    ///
    /// * `path` - 文件路径
    pub fn save_to_file<P: AsRef<Path>>(&self, path: P) -> Result<(), Box<dyn std::error::Error>> {
        let file = KeyStoreFile {
            version: 1,
            address: hex::encode(self.address),
            ciphertext: base64::encode(&self.encrypted_data),
            crypto: CryptoParamsFile {
                cipher: "aes-256-gcm".to_string(),
                nonce: base64::encode(&self.params.nonce),
                kdf: "pbkdf2-sha256".to_string(),
                kdf_params: KDFParamsFile {
                    salt: base64::encode(&self.params.salt),
                    iterations: self.params.iterations,
                },
            },
        };
        
        let json = serde_json::to_string_pretty(&file)?;
        std::fs::write(path, json)?;
        
        Ok(())
    }
    
    /// 从文件加载
    ///
    /// # Arguments
    ///
    /// * `path` - 文件路径
    pub fn load_from_file<P: AsRef<Path>>(path: P) -> Result<Self, Box<dyn std::error::Error>> {
        let json = std::fs::read_to_string(path)?;
        let file: KeyStoreFile = serde_json::from_str(&json)?;
        
        // 解析地址
        let address_bytes = hex::decode(&file.address)?;
        if address_bytes.len() != 32 {
            return Err("Invalid address length".into());
        }
        let mut address = [0u8; 32];
        address.copy_from_slice(&address_bytes);
        
        // 解析加密数据
        let encrypted_data = base64::decode(&file.ciphertext)?;
        
        // 解析加密参数
        let salt = base64::decode(&file.crypto.kdf_params.salt)?;
        let nonce = base64::decode(&file.crypto.nonce)?;
        
        Ok(Self {
            address,
            encrypted_data,
            params: EncryptionParams {
                salt,
                nonce,
                iterations: file.crypto.kdf_params.iterations,
            },
        })
    }
}

// Base64编码/解码辅助函数
mod base64 {
    use base64::{engine::general_purpose, Engine as _};
    
    pub fn encode<T: AsRef<[u8]>>(input: T) -> String {
        general_purpose::STANDARD.encode(input)
    }
    
    pub fn decode<T: AsRef<[u8]>>(input: T) -> Result<Vec<u8>, base64::DecodeError> {
        general_purpose::STANDARD.decode(input)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::key_manager::SignatureAlgorithm;
    
    #[test]
    fn test_keystore_encrypt_decrypt() {
        let keypair = KeyPair::generate(SignatureAlgorithm::Ed25519);
        let password = "test_password_123";
        
        // 加密
        let keystore = KeyStore::encrypt(&keypair, password).unwrap();
        
        // 解密
        let decrypted = keystore.decrypt(password).unwrap();
        
        // 验证
        assert_eq!(keypair.public_key_hash(), decrypted.public_key_hash());
    }
    
    #[test]
    fn test_keystore_wrong_password() {
        let keypair = KeyPair::generate(SignatureAlgorithm::Ed25519);
        let password = "correct_password";
        
        let keystore = KeyStore::encrypt(&keypair, password).unwrap();
        
        // 使用错误密码解密应该失败
        assert!(keystore.decrypt("wrong_password").is_err());
    }
    
    #[test]
    fn test_keystore_save_load() {
        use std::fs;
        use tempfile::tempdir;
        
        let keypair = KeyPair::generate(SignatureAlgorithm::Ed25519);
        let password = "test_password";
        
        let keystore = KeyStore::encrypt(&keypair, password).unwrap();
        
        // 保存到临时文件
        let dir = tempdir().unwrap();
        let file_path = dir.path().join("test_keystore.json");
        keystore.save_to_file(&file_path).unwrap();
        
        // 从文件加载
        let loaded = KeyStore::load_from_file(&file_path).unwrap();
        
        // 验证
        assert_eq!(keystore.address, loaded.address);
        
        // 解密验证
        let decrypted = loaded.decrypt(password).unwrap();
        assert_eq!(keypair.public_key_hash(), decrypted.public_key_hash());
        
        // 清理
        fs::remove_file(file_path).ok();
    }
}