"""
Key Provider Abstraction for HieraChain.

This module defines the interface for cryptographic key operations, allowing
for different storage backends (Local Memory, File Vault, HSM, KMS) without
changing the core consensus logic.
"""

import json
import base64
import os
from abc import ABC, abstractmethod
from cryptography.fernet import Fernet
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC

from hierachain.security.security_utils import KeyPair, CryptoError


class KeyProvider(ABC):
    """
    Abstract interface for key operations.
    Implementations handle how the private key is stored and accessed.
    """
    
    @property
    @abstractmethod
    def public_key_hex(self) -> str:
        """Return the public key in hex format."""
        pass
        
    @abstractmethod
    def sign(self, data: bytes) -> str:
        """
        Sign data and return hex signature.
        
        Args:
            data: Bytes to sign.
            
        Returns:
            Hex-encoded signature.
        """
        pass


class LocalKeyProvider(KeyProvider):
    """
    Standard provider that holds the KeyPair in memory.
    Used for development and backward compatibility.
    """
    
    def __init__(self, keypair: KeyPair):
        self._keypair = keypair
        
    @property
    def public_key_hex(self) -> str:
        return self._keypair.public_key
        
    def sign(self, data: bytes) -> str:
        return self._keypair.sign(data)
        
    @classmethod
    def generate(cls) -> 'LocalKeyProvider':
        """Generate a new random key provider."""
        return cls(KeyPair.generate())

    @classmethod
    def from_file(cls, path: str) -> 'LocalKeyProvider':
        """
        Load keypair from a validator identity JSON file.
        
        Args:
            path: Path to the validator_key.json file.
            
        Returns:
            LocalKeyProvider instance.
        """
        if not os.path.exists(path):
            raise CryptoError(f"Identity file not found: {path}")
            
        try:
            with open(path, "r", encoding="utf-8") as f:
                data = json.load(f)
                
            private_key_hex = data.get("private_key")
            if not private_key_hex:
                raise CryptoError("Missing 'private_key' in identity file")
                
            keypair = KeyPair.from_private_key(private_key_hex)
            return cls(keypair)
        except json.JSONDecodeError:
            raise CryptoError(f"Invalid JSON format in {path}")
        except Exception as e:
            raise CryptoError(f"Failed to load identity from {path}: {str(e)}")


class FileVaultProvider(KeyProvider):
    """
    Local file-based vault for development/testing.
    
    ⚠️ Production: Use HSM or cloud KMS instead.
    See: MasterKeyProvider for production key management.
    
    The private key is only decrypted briefly during the signing operation 
    (or held in protected memory depending on implementation), relying on 
    a password/master key provided at runtime.
    """
    
    def __init__(self, vault_path: str, password: str):
        """
        Initialize the vault provider.
        
        Args:
            vault_path: Path to the .vault file.
            password: Password to unlock the vault.
        """
        self.vault_path = vault_path
        self._password = password
        self._public_key: str | None = None
        
        # Verify easy access by loading public key immediately
        self._load_public_key()
        
    def _derive_key(self, salt: bytes) -> bytes:
        """Derive AES key from password."""
        kdf = PBKDF2HMAC(
            algorithm=hashes.SHA256(),
            length=32,
            salt=salt,
            iterations=100000,
        )
        return base64.urlsafe_b64encode(kdf.derive(self._password.encode()))

    def _load_public_key(self):
        """Load public key from vault to verify password and cache it."""
        if not os.path.exists(self.vault_path):
            raise CryptoError(f"Vault not found: {self.vault_path}")
            
        with open(self.vault_path, "rb") as f:
            data = json.load(f)
            
        encrypted_blob = base64.b64decode(data['blob'])
        salt = base64.b64decode(data['salt'])
        
        key = self._derive_key(salt)
        f = Fernet(key)
        
        try:
            decrypted = f.decrypt(encrypted_blob)
            key_data = json.loads(decrypted)
            self._public_key = key_data['public_key']
        except Exception:
            raise CryptoError("Invalid vault password or corrupted vault")

    @property
    def public_key_hex(self) -> str:
        if not self._public_key:
            self._load_public_key()
        return self._public_key  # type: ignore

    def sign(self, data: bytes) -> str:
        """
        Decrypt private key, sign, and return signature.

        Security Note:
            This method uses ephemeral decryption where the private key is held
            in memory briefly during signing. However, due to Python's memory
            management, key material may remain in memory after this call returns.

            For high-security production environments, consider using:
            - Hardware Security Module (HSM)
            - Cloud Key Management Service (AWS KMS, GCP KMS, Azure Key Vault)

            FileVaultProvider is suitable for development, testing, and
            moderate-security deployments where HSM is not available.

        Args:
            data: Bytes to sign.

        Returns:
            Hex-encoded signature string.
        """
        # 1. Load and Decrypt
        with open(self.vault_path, "rb") as f:
            vault_data = json.load(f)
            
        salt = base64.b64decode(vault_data['salt'])
        key = self._derive_key(salt)
        f = Fernet(key)
        
        decrypted = f.decrypt(base64.b64decode(vault_data['blob']))
        key_data = json.loads(decrypted)
        
        # 2. Re-construct KeyPair ephemeral
        kp = KeyPair.from_private_key(key_data['private_key'])
        
        # 3. Sign
        signature = kp.sign(data)
        
        # 4. Cleanup (Variables go out of scope)
        del kp
        del key_data
        del decrypted
        
        return signature

    @classmethod
    def create_vault(cls, vault_path: str, password: str) -> 'FileVaultProvider':
        """Create a new vault with a fresh keypair."""
        # Generate fresh keys
        kp = KeyPair.generate()
        key_data = {
            'public_key': kp.public_key,
            'private_key': kp.private_key
        }
        
        # Encrypt
        salt = os.urandom(16)
        kdf = PBKDF2HMAC(
            algorithm=hashes.SHA256(),
            length=32,
            salt=salt,
            iterations=100000,
        )
        key = base64.urlsafe_b64encode(kdf.derive(password.encode()))
        f = Fernet(key)
        
        encrypted_blob = f.encrypt(json.dumps(key_data).encode())
        
        # Save
        data = {
            'salt': base64.b64encode(salt).decode('utf-8'),
            'blob': base64.b64encode(encrypted_blob).decode('utf-8')
        }
        
        with open(vault_path, "w") as f_out:
            json.dump(data, f_out)
            
        return cls(vault_path, password)
