use std::fs::{self, OpenOptions};
use std::io::{self, Read, Write};
use std::os::unix::fs::PermissionsExt;
use std::path::{Path, PathBuf};
use rand::{rngs::OsRng, RngCore};
use thiserror::Error;
use serde::{Serialize, Deserialize};
use aes_gcm::{
    aead::{Aead, KeyInit, AeadCore},
    Aes256Gcm, Nonce, Key
};
use sha2::{Sha256, Digest};

#[derive(Error, Debug)]
pub enum StorageError {
    #[error("IO Error: {0}")]
    IoError(#[from] io::Error),
    #[error("Could not read or create salt in any allowed location")]
    SaltCreationError,
    #[error("Encryption error: {0}")]
    EncryptionError(String),
    #[error("Decryption error: {0}")]
    DecryptionError(String),
    #[error("Serialization error: {0}")]
    SerializationError(#[from] serde_json::Error),
}


#[derive(Debug, Serialize, Deserialize)]
pub struct SealedData {
    /// A human-readable label or identifier for the key.
    pub label: String,
    /// The unique nonce/IV used for this encryption operation.
    pub nonce_hex: String,
    /// The actual encrypted ciphertext (hex).
    pub ciphertext_hex: String,
}

pub trait KeyStorage {
    /// Retrieve the salt from storage.
    fn get_salt(&self) -> Result<[u8; 32], StorageError>;
    /// Persist the salt to storage.
    fn store_salt(&self, salt: &[u8; 32]) -> Result<(), StorageError>;
    /// Returns a human-readable description of the storage location.
    fn get_location_info(&self) -> String;
}

pub struct FileSystemStorage {
    primary_path: PathBuf,
    backup_path: PathBuf,
}

impl FileSystemStorage {
    pub fn new(primary: Option<&str>, backup: Option<&str>) -> Self {
        let default_primary = PathBuf::from("/boot/.device_salt");
        
        let default_backup = dirs::home_dir()
            .map(|mut p| {
                p.push(".hardware_identity");
                p.push(".device_salt_backup");
                p
            })
            .unwrap_or_else(|| PathBuf::from("./.device_salt_backup"));

        Self {
            primary_path: primary.map(PathBuf::from).unwrap_or(default_primary),
            backup_path: backup.map(PathBuf::from).unwrap_or(default_backup),
        }
    }

    fn read_salt_from_path(&self, path: &Path) -> io::Result<[u8; 32]> {
        if !path.exists() {
            return Err(io::Error::new(io::ErrorKind::NotFound, "File does not exist"));
        }

        let mut file = fs::File::open(path)?;
        let mut buffer = Vec::new();
        file.read_to_end(&mut buffer)?;

        if buffer.len() == 32 {
            let mut salt = [0u8; 32];
            salt.copy_from_slice(&buffer[..32]);
            Ok(salt)
        } else {
            Err(io::Error::new(
                io::ErrorKind::InvalidData,
                "Salt file has incorrect length",
            ))
        }
    }

    fn write_salt_to_path(&self, path: &Path, salt: &[u8; 32]) -> io::Result<()> {
        let mut file = OpenOptions::new()
            .write(true)
            .create(true)
            .truncate(true)
            .open(path)?;

        file.write_all(salt)?;
        
        // Ensure restrictive permissions (rw------- 0600)
        #[cfg(unix)]
        {
            let mut perms = fs::metadata(path)?.permissions();
            perms.set_mode(0o600);
            fs::set_permissions(path, perms)?;
        }

        Ok(())
    }
}

impl KeyStorage for FileSystemStorage {
    fn get_salt(&self) -> Result<[u8; 32], StorageError> {

        if let Ok(salt) = self.read_salt_from_path(&self.primary_path) {
            return Ok(salt);
        }

        if let Ok(salt) = self.read_salt_from_path(&self.backup_path) {
            return Ok(salt);
        }

        Err(StorageError::SaltCreationError)
    }

    fn store_salt(&self, salt: &[u8; 32]) -> Result<(), StorageError> {
     
        if self.write_salt_to_path(&self.primary_path, salt).is_ok() {
            return Ok(());
        }

        if let Some(parent) = self.backup_path.parent() {
            let _ = fs::create_dir_all(parent);
        }

        if self.write_salt_to_path(&self.backup_path, salt).is_ok() {
            return Ok(());
        }

        Err(StorageError::SaltCreationError)
    }

    fn get_location_info(&self) -> String {
        if self.primary_path.exists() {
            self.primary_path.to_string_lossy().to_string()
        } else {
            self.backup_path.to_string_lossy().to_string()
        }
    }
}


pub struct SecureEnclaveStorage {
    /// A unique identifier for the enclave/device.
    identifier: String,
    /// Where to persist the sealed data.
    persistence_path: PathBuf,
}

impl SecureEnclaveStorage {
    pub fn new(identifier: &str, persistence: Option<&str>) -> Self {
        let default_persistence = dirs::home_dir()
            .map(|mut p| {
                p.push(".hardware_identity");
                p.push("enclave_salt.json");
                p
            })
            .unwrap_or_else(|| PathBuf::from("./enclave_salt.json"));

        Self {
            identifier: identifier.to_string(),
            persistence_path: persistence.map(PathBuf::from).unwrap_or(default_persistence),
        }
    }

    fn get_hardware_key(&self) -> Key<Aes256Gcm> {
        let mut hasher = Sha256::new();
        hasher.update(self.identifier.as_bytes());
        hasher.update(b"ENCLAVE_ROOT_KEY_SALT");
        let hash = hasher.finalize();
        *Key::<Aes256Gcm>::from_slice(&hash)
    }

    fn seal(&self, plaintext: &[u8]) -> Result<SealedData, StorageError> {
        let key = self.get_hardware_key();
        let cipher = Aes256Gcm::new(&key);
        let nonce = Aes256Gcm::generate_nonce(&mut OsRng);

        let ciphertext = cipher.encrypt(&nonce, plaintext)
            .map_err(|e| StorageError::EncryptionError(e.to_string()))?;

        Ok(SealedData {
            label: format!("device-salt-{}", self.identifier),
            nonce_hex: hex::encode(nonce),
            ciphertext_hex: hex::encode(ciphertext),
        })
    }

    fn unseal(&self, sealed: &SealedData) -> Result<Vec<u8>, StorageError> {
        let key = self.get_hardware_key();
        let cipher = Aes256Gcm::new(&key);
        
        let nonce_bytes = hex::decode(&sealed.nonce_hex)
            .map_err(|_| StorageError::DecryptionError("Invalid nonce hex".to_string()))?;
        let ciphertext_bytes = hex::decode(&sealed.ciphertext_hex)
            .map_err(|_| StorageError::DecryptionError("Invalid ciphertext hex".to_string()))?;

        let nonce = Nonce::from_slice(&nonce_bytes);
        
        let plaintext = cipher.decrypt(nonce, ciphertext_bytes.as_slice())
            .map_err(|e| StorageError::DecryptionError(e.to_string()))?;

        Ok(plaintext)
    }
}

impl KeyStorage for SecureEnclaveStorage {
    fn get_salt(&self) -> Result<[u8; 32], StorageError> {
        if !self.persistence_path.exists() {
            return Err(StorageError::SaltCreationError);
        }

        let data = fs::read_to_string(&self.persistence_path)?;
        let sealed: SealedData = serde_json::from_str(&data)?;
        
        let unsealed = self.unseal(&sealed)?;
        
        if unsealed.len() == 32 {
            let mut salt = [0u8; 32];
            salt.copy_from_slice(&unsealed);
            Ok(salt)
        } else {
            Err(StorageError::DecryptionError("Unsealed data has incorrect length".to_string()))
        }
    }

    fn store_salt(&self, salt: &[u8; 32]) -> Result<(), StorageError> {
        let sealed = self.seal(salt)?;
        let json = serde_json::to_string_pretty(&sealed)?;

        if let Some(parent) = self.persistence_path.parent() {
            let _ = fs::create_dir_all(parent);
        }

        fs::write(&self.persistence_path, json)?;
        
        #[cfg(unix)]
        {
            let mut perms = fs::metadata(&self.persistence_path)?.permissions();
            perms.set_mode(0o600);
            fs::set_permissions(&self.persistence_path, perms)?;
        }

        Ok(())
    }

    fn get_location_info(&self) -> String {
        format!("SecureEnclave[{}]:{}", self.identifier, self.persistence_path.display())
    }
}
#[cfg(test)]
mod tests {
    use super::*;
    use tempfile::tempdir;

    #[test]
    fn test_secure_enclave_sealing() {
        let dir = tempdir().unwrap();
        let persistence_path = dir.path().join("test_enclave.json");
        let storage = SecureEnclaveStorage::new("test-device-123", Some(persistence_path.to_str().unwrap()));

        let salt = [0u8; 32];
        storage.store_salt(&salt).unwrap();

        let content = fs::read_to_string(&persistence_path).unwrap();
        assert!(content.contains("nonce_hex"));
        assert!(content.contains("ciphertext_hex"));


        let loaded_salt = storage.get_salt().unwrap();
        assert_eq!(salt, loaded_salt);

        let other_storage = SecureEnclaveStorage::new("wrong-device", Some(persistence_path.to_str().unwrap()));
        let result = other_storage.get_salt();
        assert!(result.is_err());
    }

    #[test]
    fn test_file_system_persistence() {
        let dir = tempdir().unwrap();
        let primary = dir.path().join("primary_salt");
        let backup = dir.path().join("backup_salt");
        
        let storage = FileSystemStorage::new(
            Some(primary.to_str().unwrap()),
            Some(backup.to_str().unwrap())
        );

        let salt = [42u8; 32];
        storage.store_salt(&salt).unwrap();

        // Verify retrieval
        let loaded_salt = storage.get_salt().unwrap();
        assert_eq!(salt, loaded_salt);
        
        // Verify primary was used
        assert!(primary.exists());
    }
}
