#!/usr/bin/env python3
"""
Private Gallery Encryption Module

Provides security features for the Private Gallery:
- Password hashing with bcrypt
- Key derivation with Argon2id
- File encryption/decryption with AES-256-GCM
- Field encryption with Fernet
- Session token management
"""

import os
import secrets
import hashlib
import base64
import time
from datetime import datetime, timedelta
from typing import Optional, Dict, Tuple
from pathlib import Path
from threading import Lock

try:
    import bcrypt
except ImportError:
    bcrypt = None

try:
    from argon2 import PasswordHasher
    from argon2.low_level import hash_secret_raw, Type
    ARGON2_AVAILABLE = True
except ImportError:
    ARGON2_AVAILABLE = False

try:
    from cryptography.fernet import Fernet
    from cryptography.hazmat.primitives.ciphers.aead import AESGCM
    from cryptography.hazmat.primitives import hashes
    from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
    CRYPTO_AVAILABLE = True
except ImportError:
    CRYPTO_AVAILABLE = False

from modules.universal_logger import get_logger

logger = get_logger('PrivateGalleryCrypto')


class PrivateGalleryCrypto:
    """
    Handles all encryption operations for the Private Gallery.

    Security features:
    - Passwords hashed with bcrypt (cost factor 12)
    - Encryption key derived from password using Argon2id
    - Files encrypted with AES-256-GCM
    - Database fields encrypted with Fernet (AES-128-CBC + HMAC)
    - Session tokens with configurable timeout
    """

    # Argon2id parameters (OWASP recommended)
    ARGON2_TIME_COST = 3
    ARGON2_MEMORY_COST = 65536  # 64 MiB
    ARGON2_PARALLELISM = 4
    ARGON2_HASH_LENGTH = 32  # 256 bits for AES-256

    # AES-GCM parameters
    AES_KEY_SIZE = 32  # 256 bits
    AES_NONCE_SIZE = 12  # 96 bits (GCM recommended)
    AES_TAG_SIZE = 16  # 128 bits

    # Encryption chunk size for streaming large files
    CHUNK_SIZE = 8 * 1024 * 1024  # 8 MB chunks
    CHUNKED_THRESHOLD = 50 * 1024 * 1024  # Use chunked encryption for files > 50 MB
    CHUNKED_MAGIC = b'\x01PGCE'  # Magic bytes: version 1, Private Gallery Chunked Encryption

    def __init__(self):
        self._sessions: Dict[str, Dict] = {}  # token -> {expiry, username}
        self._session_lock = Lock()
        self._derived_key: Optional[bytes] = None
        self._fernet: Optional[Fernet] = None
        self._aesgcm: Optional[AESGCM] = None

        # Check dependencies
        if not bcrypt:
            logger.warning("bcrypt not available - password hashing will use fallback")
        if not ARGON2_AVAILABLE:
            logger.warning("argon2-cffi not available - key derivation will use PBKDF2")
        if not CRYPTO_AVAILABLE:
            raise ImportError("cryptography library required for Private Gallery")

    # =========================================================================
    # PASSWORD HASHING (bcrypt)
    # =========================================================================

    def hash_password(self, password: str) -> str:
        """
        Hash a password using bcrypt with cost factor 12.

        Args:
            password: Plain text password

        Returns:
            bcrypt hash string (includes salt)
        """
        if bcrypt:
            salt = bcrypt.gensalt(rounds=12)
            hashed = bcrypt.hashpw(password.encode('utf-8'), salt)
            return hashed.decode('utf-8')
        else:
            # Fallback to PBKDF2 if bcrypt not available
            salt = secrets.token_bytes(16)
            kdf = PBKDF2HMAC(
                algorithm=hashes.SHA256(),
                length=32,
                salt=salt,
                iterations=600000,
            )
            key = kdf.derive(password.encode('utf-8'))
            return f"pbkdf2${base64.b64encode(salt).decode()}${base64.b64encode(key).decode()}"

    def verify_password(self, password: str, password_hash: str) -> bool:
        """
        Verify a password against its hash.

        Args:
            password: Plain text password to check
            password_hash: Stored hash to verify against

        Returns:
            True if password matches
        """
        try:
            if password_hash.startswith('pbkdf2$'):
                # PBKDF2 fallback hash
                parts = password_hash.split('$')
                if len(parts) != 3:
                    return False
                salt = base64.b64decode(parts[1])
                stored_key = base64.b64decode(parts[2])
                kdf = PBKDF2HMAC(
                    algorithm=hashes.SHA256(),
                    length=32,
                    salt=salt,
                    iterations=600000,
                )
                try:
                    kdf.verify(password.encode('utf-8'), stored_key)
                    return True
                except Exception:
                    return False
            elif bcrypt:
                return bcrypt.checkpw(
                    password.encode('utf-8'),
                    password_hash.encode('utf-8')
                )
            else:
                return False
        except Exception as e:
            logger.error(f"Password verification failed: {e}")
            return False

    # =========================================================================
    # KEY DERIVATION (Argon2id or PBKDF2)
    # =========================================================================

    def derive_key(self, password: str, salt: bytes) -> bytes:
        """
        Derive an encryption key from password using Argon2id.

        Args:
            password: User's password
            salt: Random salt (should be stored)

        Returns:
            32-byte derived key for AES-256
        """
        if ARGON2_AVAILABLE:
            key = hash_secret_raw(
                secret=password.encode('utf-8'),
                salt=salt,
                time_cost=self.ARGON2_TIME_COST,
                memory_cost=self.ARGON2_MEMORY_COST,
                parallelism=self.ARGON2_PARALLELISM,
                hash_len=self.ARGON2_HASH_LENGTH,
                type=Type.ID  # Argon2id
            )
            return key
        else:
            # Fallback to PBKDF2 with high iterations
            kdf = PBKDF2HMAC(
                algorithm=hashes.SHA256(),
                length=self.AES_KEY_SIZE,
                salt=salt,
                iterations=600000,  # OWASP recommended minimum
            )
            return kdf.derive(password.encode('utf-8'))

    def generate_salt(self) -> bytes:
        """Generate a cryptographically secure random salt."""
        return secrets.token_bytes(16)

    def initialize_encryption(self, password: str, salt: bytes) -> None:
        """
        Initialize encryption with derived key.
        Must be called after successful unlock.

        Args:
            password: User's password
            salt: Stored salt for key derivation
        """
        self._derived_key = self.derive_key(password, salt)

        # Initialize Fernet for field encryption
        # Fernet requires a 32-byte key, base64-encoded
        fernet_key = base64.urlsafe_b64encode(self._derived_key)
        self._fernet = Fernet(fernet_key)

        # Initialize AES-GCM for file encryption
        self._aesgcm = AESGCM(self._derived_key)

        logger.info("Encryption initialized successfully")

    def clear_encryption(self) -> None:
        """Clear encryption keys from memory (on lock)."""
        self._derived_key = None
        self._fernet = None
        self._aesgcm = None
        logger.info("Encryption keys cleared")

    def is_initialized(self) -> bool:
        """Check if encryption is initialized (unlocked)."""
        return self._derived_key is not None

    # =========================================================================
    # FIELD ENCRYPTION (Fernet - for database fields)
    # =========================================================================

    def encrypt_field(self, plaintext: str) -> str:
        """
        Encrypt a database field value.

        Args:
            plaintext: Plain text to encrypt

        Returns:
            Base64-encoded encrypted string
        """
        if not self._fernet:
            raise RuntimeError("Encryption not initialized - call initialize_encryption first")

        if not plaintext:
            return ""

        encrypted = self._fernet.encrypt(plaintext.encode('utf-8'))
        return base64.urlsafe_b64encode(encrypted).decode('utf-8')

    def decrypt_field(self, ciphertext: str) -> str:
        """
        Decrypt a database field value.

        Args:
            ciphertext: Base64-encoded encrypted string

        Returns:
            Decrypted plain text
        """
        if not self._fernet:
            raise RuntimeError("Encryption not initialized - call initialize_encryption first")

        if not ciphertext:
            return ""

        try:
            encrypted = base64.urlsafe_b64decode(ciphertext.encode('utf-8'))
            decrypted = self._fernet.decrypt(encrypted)
            return decrypted.decode('utf-8')
        except Exception as e:
            logger.error(f"Field decryption failed: {e}")
            return "[Decryption Error]"

    # =========================================================================
    # FILE ENCRYPTION (AES-256-GCM)
    # =========================================================================

    def encrypt_file(self, input_path: Path, output_path: Path) -> bool:
        """
        Encrypt a file using AES-256-GCM.

        Small files (<=50MB): single-shot format
            [12-byte nonce][encrypted data + 16-byte tag]

        Large files (>50MB): chunked format for memory efficiency
            [5-byte magic 0x01PGCE][4-byte chunk_size BE]
            [12-byte nonce][encrypted chunk + 16-byte tag] (repeated)

        Args:
            input_path: Path to plaintext file
            output_path: Path for encrypted output

        Returns:
            True if successful
        """
        if not self._aesgcm:
            raise RuntimeError("Encryption not initialized")

        try:
            file_size = input_path.stat().st_size
            output_path.parent.mkdir(parents=True, exist_ok=True)

            if file_size <= self.CHUNKED_THRESHOLD:
                # Small file: single-shot encryption (backward compatible)
                nonce = secrets.token_bytes(self.AES_NONCE_SIZE)
                with open(input_path, 'rb') as f:
                    plaintext = f.read()
                ciphertext = self._aesgcm.encrypt(nonce, plaintext, None)
                with open(output_path, 'wb') as f:
                    f.write(nonce)
                    f.write(ciphertext)
            else:
                # Large file: chunked encryption
                import struct
                with open(input_path, 'rb') as fin, open(output_path, 'wb') as fout:
                    # Write header
                    fout.write(self.CHUNKED_MAGIC)
                    fout.write(struct.pack('>I', self.CHUNK_SIZE))

                    # Encrypt in chunks
                    while True:
                        chunk = fin.read(self.CHUNK_SIZE)
                        if not chunk:
                            break
                        nonce = secrets.token_bytes(self.AES_NONCE_SIZE)
                        encrypted_chunk = self._aesgcm.encrypt(nonce, chunk, None)
                        # Write chunk: nonce + encrypted data (includes GCM tag)
                        fout.write(nonce)
                        fout.write(struct.pack('>I', len(encrypted_chunk)))
                        fout.write(encrypted_chunk)

            return True

        except Exception as e:
            logger.error(f"File encryption failed: {e}")
            # Clean up partial output
            if output_path.exists():
                try:
                    output_path.unlink()
                except Exception:
                    pass
            return False

    def _is_chunked_format(self, input_path: Path) -> bool:
        """Check if an encrypted file uses the chunked format."""
        try:
            with open(input_path, 'rb') as f:
                magic = f.read(len(self.CHUNKED_MAGIC))
                return magic == self.CHUNKED_MAGIC
        except Exception:
            return False

    def decrypt_file(self, input_path: Path, output_path: Optional[Path] = None) -> Optional[bytes]:
        """
        Decrypt a file encrypted with AES-256-GCM.
        Handles both single-shot and chunked formats.

        Args:
            input_path: Path to encrypted file
            output_path: Optional path to write decrypted file

        Returns:
            Decrypted bytes if output_path is None, else None on success
        """
        if not self._aesgcm:
            raise RuntimeError("Encryption not initialized")

        try:
            if self._is_chunked_format(input_path):
                return self._decrypt_file_chunked(input_path, output_path)

            # Single-shot format: [nonce][ciphertext+tag]
            with open(input_path, 'rb') as f:
                nonce = f.read(self.AES_NONCE_SIZE)
                if len(nonce) != self.AES_NONCE_SIZE:
                    raise ValueError("Invalid encrypted file: missing nonce")
                ciphertext = f.read()

            plaintext = self._aesgcm.decrypt(nonce, ciphertext, None)

            if output_path:
                output_path.parent.mkdir(parents=True, exist_ok=True)
                with open(output_path, 'wb') as f:
                    f.write(plaintext)
                return None

            return plaintext

        except Exception as e:
            logger.error(f"File decryption failed: {e}")
            return None

    def _decrypt_file_chunked(self, input_path: Path, output_path: Optional[Path] = None) -> Optional[bytes]:
        """Decrypt a chunked-format encrypted file."""
        import struct

        try:
            parts = [] if output_path is None else None

            with open(input_path, 'rb') as fin:
                # Read header
                magic = fin.read(len(self.CHUNKED_MAGIC))
                if magic != self.CHUNKED_MAGIC:
                    raise ValueError("Invalid chunked file header")
                chunk_size_bytes = fin.read(4)
                # chunk_size from header (informational, actual sizes are per-chunk)
                struct.unpack('>I', chunk_size_bytes)

                fout = None
                if output_path:
                    output_path.parent.mkdir(parents=True, exist_ok=True)
                    fout = open(output_path, 'wb')

                try:
                    while True:
                        # Read chunk: [12-byte nonce][4-byte encrypted_len][encrypted data]
                        nonce = fin.read(self.AES_NONCE_SIZE)
                        if len(nonce) == 0:
                            break  # EOF
                        if len(nonce) != self.AES_NONCE_SIZE:
                            raise ValueError("Truncated chunk nonce")

                        enc_len_bytes = fin.read(4)
                        if len(enc_len_bytes) != 4:
                            raise ValueError("Truncated chunk length")
                        enc_len = struct.unpack('>I', enc_len_bytes)[0]

                        encrypted_chunk = fin.read(enc_len)
                        if len(encrypted_chunk) != enc_len:
                            raise ValueError("Truncated chunk data")

                        decrypted_chunk = self._aesgcm.decrypt(nonce, encrypted_chunk, None)

                        if fout:
                            fout.write(decrypted_chunk)
                        else:
                            parts.append(decrypted_chunk)
                finally:
                    if fout:
                        fout.close()

            if output_path:
                return None
            return b''.join(parts)

        except Exception as e:
            logger.error(f"Chunked file decryption failed for {input_path}: {type(e).__name__}: {e}")
            return None

    def re_encrypt_to_chunked(self, file_path: Path) -> bool:
        """
        Re-encrypt a single-shot encrypted file to chunked format in-place.
        Decrypts and re-encrypts in chunks to avoid loading the entire file into memory.

        Args:
            file_path: Path to the single-shot encrypted file

        Returns:
            True if successful, False if already chunked or on error
        """
        if not self._aesgcm:
            raise RuntimeError("Encryption not initialized")

        if self._is_chunked_format(file_path):
            return False  # Already chunked

        import struct
        temp_path = file_path.with_suffix(f'.enc.{secrets.token_hex(4)}.tmp')

        try:
            # Decrypt the single-shot file fully (required by AES-GCM)
            with open(file_path, 'rb') as f:
                nonce = f.read(self.AES_NONCE_SIZE)
                if len(nonce) != self.AES_NONCE_SIZE:
                    raise ValueError("Invalid encrypted file")
                ciphertext = f.read()

            plaintext = self._aesgcm.decrypt(nonce, ciphertext, None)
            del ciphertext  # Free memory

            # Write chunked format to temp file
            with open(temp_path, 'wb') as fout:
                fout.write(self.CHUNKED_MAGIC)
                fout.write(struct.pack('>I', self.CHUNK_SIZE))

                offset = 0
                while offset < len(plaintext):
                    chunk = plaintext[offset:offset + self.CHUNK_SIZE]
                    offset += len(chunk)
                    chunk_nonce = secrets.token_bytes(self.AES_NONCE_SIZE)
                    encrypted_chunk = self._aesgcm.encrypt(chunk_nonce, chunk, None)
                    fout.write(chunk_nonce)
                    fout.write(struct.pack('>I', len(encrypted_chunk)))
                    fout.write(encrypted_chunk)

            del plaintext  # Free memory

            # Atomic replace
            temp_path.replace(file_path)
            return True

        except Exception as e:
            logger.error(f"Re-encryption to chunked failed for {file_path}: {e}")
            if temp_path.exists():
                try:
                    temp_path.unlink()
                except Exception:
                    pass
            return False

    def decrypt_file_streaming(self, input_path: Path) -> Optional[bytes]:
        """
        Decrypt a file and return bytes for streaming.
        Only suitable for small files (single-shot format, ≤50MB).
        For large chunked files, use decrypt_file_generator() instead.

        Args:
            input_path: Path to encrypted file

        Returns:
            Decrypted bytes or None on error
        """
        return self.decrypt_file(input_path, output_path=None)

    def decrypt_file_generator(self, input_path: Path):
        """
        Generator that yields decrypted chunks for streaming large files.
        For chunked files, yields one decrypted chunk at a time (~8MB each).
        For single-shot files, yields the entire content at once.

        Args:
            input_path: Path to encrypted file

        Yields:
            bytes: Decrypted data chunks
        """
        import struct

        if not self._aesgcm:
            raise RuntimeError("Encryption not initialized")

        if self._is_chunked_format(input_path):
            with open(input_path, 'rb') as fin:
                # Skip header
                fin.read(len(self.CHUNKED_MAGIC))
                fin.read(4)

                while True:
                    nonce = fin.read(self.AES_NONCE_SIZE)
                    if len(nonce) == 0:
                        break
                    if len(nonce) != self.AES_NONCE_SIZE:
                        raise ValueError("Truncated chunk nonce")

                    enc_len_bytes = fin.read(4)
                    if len(enc_len_bytes) != 4:
                        raise ValueError("Truncated chunk length")
                    enc_len = struct.unpack('>I', enc_len_bytes)[0]

                    encrypted_chunk = fin.read(enc_len)
                    if len(encrypted_chunk) != enc_len:
                        raise ValueError("Truncated chunk data")

                    yield self._aesgcm.decrypt(nonce, encrypted_chunk, None)
        else:
            # Single-shot: yield everything at once (≤50MB)
            with open(input_path, 'rb') as f:
                nonce = f.read(self.AES_NONCE_SIZE)
                if len(nonce) != self.AES_NONCE_SIZE:
                    raise ValueError("Invalid encrypted file: missing nonce")
                ciphertext = f.read()
            yield self._aesgcm.decrypt(nonce, ciphertext, None)

    def decrypt_file_range_generator(self, input_path: Path, start: int, end: int):
        """
        Generator that yields only the decrypted bytes for a specific byte range.
        For chunked files, only decrypts the necessary chunks and slices them.
        For single-shot files, decrypts all and slices.

        Args:
            input_path: Path to encrypted file
            start: Start byte offset (inclusive)
            end: End byte offset (inclusive)

        Yields:
            bytes: Decrypted data for the requested range
        """
        import struct

        if not self._aesgcm:
            raise RuntimeError("Encryption not initialized")

        if not self._is_chunked_format(input_path):
            # Single-shot: decrypt all and slice (file is ≤50MB)
            with open(input_path, 'rb') as f:
                nonce = f.read(self.AES_NONCE_SIZE)
                ciphertext = f.read()
            plaintext = self._aesgcm.decrypt(nonce, ciphertext, None)
            yield plaintext[start:end + 1]
            return

        chunk_size = self.CHUNK_SIZE
        first_chunk = start // chunk_size
        last_chunk = end // chunk_size

        # Header: 5 magic + 4 chunk_size = 9 bytes
        header_size = len(self.CHUNKED_MAGIC) + 4
        # Each full encrypted chunk: 12 nonce + 4 length + (chunk_size + 16 tag)
        enc_chunk_stride = self.AES_NONCE_SIZE + 4 + chunk_size + self.AES_TAG_SIZE

        with open(input_path, 'rb') as fin:
            for chunk_idx in range(first_chunk, last_chunk + 1):
                # Seek to this chunk's position in the encrypted file
                fin.seek(header_size + chunk_idx * enc_chunk_stride)

                nonce = fin.read(self.AES_NONCE_SIZE)
                if len(nonce) == 0:
                    break
                if len(nonce) != self.AES_NONCE_SIZE:
                    raise ValueError("Truncated chunk nonce")

                enc_len_bytes = fin.read(4)
                if len(enc_len_bytes) != 4:
                    raise ValueError("Truncated chunk length")
                enc_len = struct.unpack('>I', enc_len_bytes)[0]

                encrypted_chunk = fin.read(enc_len)
                if len(encrypted_chunk) != enc_len:
                    raise ValueError("Truncated chunk data")

                decrypted_chunk = self._aesgcm.decrypt(nonce, encrypted_chunk, None)

                # Calculate which part of this chunk we need
                chunk_start_byte = chunk_idx * chunk_size
                slice_start = max(start - chunk_start_byte, 0)
                slice_end = min(end - chunk_start_byte + 1, len(decrypted_chunk))

                yield decrypted_chunk[slice_start:slice_end]

    # =========================================================================
    # SESSION MANAGEMENT
    # =========================================================================

    def create_session(self, username: str = "user", timeout_minutes: int = 30) -> str:
        """
        Create a new session token.

        Args:
            username: Username for the session
            timeout_minutes: Session timeout in minutes

        Returns:
            Session token string
        """
        token = secrets.token_urlsafe(32)
        expiry = datetime.now() + timedelta(minutes=timeout_minutes)

        with self._session_lock:
            self._sessions[token] = {
                'expiry': expiry,
                'username': username,
                'created_at': datetime.now()
            }

        logger.info(f"Created session for {username}, expires in {timeout_minutes} minutes")
        return token

    def verify_session(self, token: str) -> Optional[Dict]:
        """
        Verify a session token is valid and not expired.

        Args:
            token: Session token to verify

        Returns:
            Session info dict if valid, None otherwise
        """
        with self._session_lock:
            session = self._sessions.get(token)

            if not session:
                return None

            if datetime.now() > session['expiry']:
                # Expired - remove it
                del self._sessions[token]
                return None

            return session

    def refresh_session(self, token: str, timeout_minutes: int = 30) -> bool:
        """
        Refresh a session's expiry time.

        Args:
            token: Session token to refresh
            timeout_minutes: New timeout in minutes

        Returns:
            True if refreshed, False if token invalid
        """
        with self._session_lock:
            session = self._sessions.get(token)

            if not session:
                return False

            if datetime.now() > session['expiry']:
                del self._sessions[token]
                return False

            session['expiry'] = datetime.now() + timedelta(minutes=timeout_minutes)
            return True

    def invalidate_session(self, token: str) -> bool:
        """
        Invalidate a session token (logout/lock).

        Args:
            token: Session token to invalidate

        Returns:
            True if invalidated, False if not found
        """
        with self._session_lock:
            if token in self._sessions:
                del self._sessions[token]
                return True
            return False

    def invalidate_all_sessions(self) -> int:
        """
        Invalidate all sessions (master lock).

        Returns:
            Number of sessions invalidated
        """
        with self._session_lock:
            count = len(self._sessions)
            self._sessions.clear()
            return count

    def cleanup_expired_sessions(self) -> int:
        """
        Remove all expired sessions.

        Returns:
            Number of sessions removed
        """
        with self._session_lock:
            now = datetime.now()
            expired = [t for t, s in self._sessions.items() if now > s['expiry']]
            for token in expired:
                del self._sessions[token]
            return len(expired)

    def get_active_session_count(self) -> int:
        """Get count of active (non-expired) sessions."""
        self.cleanup_expired_sessions()
        return len(self._sessions)


# Global instance
_crypto_instance: Optional[PrivateGalleryCrypto] = None
_crypto_lock = Lock()


def get_private_gallery_crypto() -> PrivateGalleryCrypto:
    """Get or create the global crypto instance."""
    global _crypto_instance

    with _crypto_lock:
        if _crypto_instance is None:
            _crypto_instance = PrivateGalleryCrypto()
        return _crypto_instance


def export_key_to_file(path: str) -> bool:
    """
    Save the current derived key from the global crypto instance to a file.
    The file is written with mode 0600 for security.

    Args:
        path: File path to write the key material to

    Returns:
        True if successful
    """
    import json as _json

    crypto = get_private_gallery_crypto()
    if not crypto.is_initialized() or crypto._derived_key is None:
        logger.warning("Cannot export key: encryption not initialized")
        return False

    try:
        key_data = {
            'derived_key': base64.b64encode(crypto._derived_key).decode('utf-8')
        }
        key_path = Path(path)
        key_path.parent.mkdir(parents=True, exist_ok=True)

        # Write atomically via temp file
        tmp_path = key_path.with_suffix('.tmp')
        with open(tmp_path, 'w') as f:
            _json.dump(key_data, f)
        os.chmod(str(tmp_path), 0o600)
        tmp_path.replace(key_path)

        logger.info(f"Exported encryption key to {path}")
        return True
    except Exception as e:
        logger.error(f"Failed to export key to {path}: {e}")
        return False


def load_key_from_file(path: str) -> Optional[PrivateGalleryCrypto]:
    """
    Load a derived key from a file and return an initialized crypto instance.

    Args:
        path: File path containing the key material

    Returns:
        Initialized PrivateGalleryCrypto instance, or None if unavailable
    """
    import json as _json

    key_path = Path(path)
    if not key_path.exists():
        return None

    try:
        with open(key_path, 'r') as f:
            key_data = _json.load(f)

        derived_key = base64.b64decode(key_data['derived_key'])

        crypto = PrivateGalleryCrypto()
        crypto._derived_key = derived_key

        # Initialize Fernet for field encryption
        fernet_key = base64.urlsafe_b64encode(derived_key)
        crypto._fernet = Fernet(fernet_key)

        # Initialize AES-GCM for file encryption
        crypto._aesgcm = AESGCM(derived_key)

        return crypto
    except Exception as e:
        logger.error(f"Failed to load key from {path}: {e}")
        return None


def delete_key_file(path: str) -> bool:
    """Delete the key file if it exists."""
    try:
        key_path = Path(path)
        if key_path.exists():
            key_path.unlink()
            logger.info(f"Deleted key file {path}")
        return True
    except Exception as e:
        logger.error(f"Failed to delete key file {path}: {e}")
        return False