# Familiar - Self-Hosted AI Companion Platform
# Copyright (c) 2026 George Scott Foley
# ORCID: 0009-0006-4957-0540
# Email: Georgescottfoley@proton.me
# Licensed under the MIT License - see LICENSE file for details

"""
Familiar Encryption at Rest Module

Provides transparent encryption for sensitive data storage including:
- Session data
- Memory database
- Conversation history

Uses Fernet symmetric encryption (AES-128-CBC with HMAC-SHA256).

Usage:
    from familiar.core.encryption import EncryptedStorage

    storage = EncryptedStorage()  # Reads key from FAMILIAR_ENCRYPTION_KEY

    # Encrypt before writing
    encrypted = storage.encrypt(json.dumps(data))
    Path("data.json.enc").write_text(encrypted)

    # Decrypt after reading
    decrypted = storage.decrypt(Path("data.json.enc").read_text())
    data = json.loads(decrypted)
"""

import base64
import json
import logging
import os
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Optional, Union

logger = logging.getLogger(__name__)

# Optional dependency - graceful degradation if not installed
try:
    from cryptography.fernet import Fernet, InvalidToken
    from cryptography.hazmat.primitives import hashes
    from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC

    CRYPTOGRAPHY_AVAILABLE = True
except ImportError:
    CRYPTOGRAPHY_AVAILABLE = False
    logger.warning("cryptography package not installed. Encryption at rest disabled.")


class EncryptionError(Exception):
    """Raised when encryption/decryption fails."""

    pass


class EncryptionKeyError(Exception):
    """Raised when encryption key is missing or invalid."""

    pass


@dataclass
class EncryptionConfig:
    """Configuration for encryption at rest."""

    enabled: bool = False
    key_env_var: str = "FAMILIAR_ENCRYPTION_KEY"
    encrypt_sessions: bool = True
    encrypt_memory: bool = True
    encrypt_history: bool = True
    # For key derivation from password (alternative to direct key)
    use_password_derivation: bool = False
    password_env_var: str = "FAMILIAR_ENCRYPTION_PASSWORD"
    salt_file: str = str(Path.home() / ".familiar" / "data" / ".encryption_salt")


class EncryptedStorage:
    """
    Provides encryption at rest for Familiar data files.

    Supports two modes:
    1. Direct Fernet key (recommended): Set FAMILIAR_ENCRYPTION_KEY to a Fernet key
    2. Password derivation: Set FAMILIAR_ENCRYPTION_PASSWORD and a salt file is created

    Example:
        # Generate a key
        python -c "from cryptography.fernet import Fernet; print(Fernet.generate_key().decode())"

        # Set in environment
        export FAMILIAR_ENCRYPTION_KEY="your-generated-key"

        # Use in code
        storage = EncryptedStorage()
        encrypted = storage.encrypt("sensitive data")
        decrypted = storage.decrypt(encrypted)
    """

    def __init__(self, config: Optional[EncryptionConfig] = None, key: Optional[str] = None):
        """
        Initialize encrypted storage.

        Args:
            config: Encryption configuration. Defaults to reading from environment.
            key: Direct encryption key. Overrides environment variable.
        """
        self.config = config or EncryptionConfig()
        self._fernet: Optional[Any] = None
        self._key = key
        self._initialized = False

        if not CRYPTOGRAPHY_AVAILABLE:
            logger.warning("Encryption unavailable - cryptography package not installed")
            return

        self._initialize_encryption()

    def _initialize_encryption(self) -> None:
        """Initialize the Fernet cipher with the configured key."""
        if self._initialized:
            return

        try:
            key = self._get_key()
            if key:
                self._fernet = Fernet(key.encode() if isinstance(key, str) else key)
                self._initialized = True
                logger.info("Encryption at rest initialized successfully")
            else:
                logger.debug("No encryption key configured - encryption disabled")
        except (ValueError, TypeError, AttributeError) as e:
            logger.error(f"Failed to initialize encryption: {e}")
            raise EncryptionKeyError(f"Invalid encryption key: {e}")

    def _get_key(self) -> Optional[str]:
        """Get the encryption key from configured source."""
        # Direct key takes precedence
        if self._key:
            return self._key

        # Try environment variable for direct key
        key = os.environ.get(self.config.key_env_var)
        if key:
            return key

        # Try password derivation
        if self.config.use_password_derivation:
            password = os.environ.get(self.config.password_env_var)
            if password:
                return self._derive_key_from_password(password)

        return None

    def _derive_key_from_password(self, password: str) -> str:
        """
        Derive a Fernet key from a password using PBKDF2.

        Creates or reads a salt file for consistent key derivation.
        """
        salt_path = Path(self.config.salt_file)

        if salt_path.exists():
            salt = salt_path.read_bytes()
        else:
            salt = os.urandom(16)
            salt_path.write_bytes(salt)
            salt_path.chmod(0o600)
            logger.info(f"Created encryption salt file: {salt_path}")

        kdf = PBKDF2HMAC(
            algorithm=hashes.SHA256(),
            length=32,
            salt=salt,
            iterations=480000,  # OWASP 2023 recommendation
        )

        key = base64.urlsafe_b64encode(kdf.derive(password.encode()))
        return key.decode()

    @property
    def is_available(self) -> bool:
        """Check if encryption is available and properly configured."""
        return CRYPTOGRAPHY_AVAILABLE and self._fernet is not None

    def encrypt(self, data: Union[str, bytes]) -> str:
        """
        Encrypt data and return base64-encoded ciphertext.

        Args:
            data: String or bytes to encrypt

        Returns:
            Base64-encoded encrypted string

        Raises:
            EncryptionError: If encryption fails
            EncryptionKeyError: If encryption is not configured
        """
        if not self.is_available:
            if not CRYPTOGRAPHY_AVAILABLE:
                raise EncryptionError("cryptography package not installed")
            raise EncryptionKeyError("Encryption key not configured")

        # Explicit None check for type safety (mypy)
        if self._fernet is None:
            raise EncryptionKeyError("Encryption not initialized")

        try:
            if isinstance(data, str):
                data = data.encode("utf-8")

            encrypted = self._fernet.encrypt(data)
            return encrypted.decode("utf-8")

        except Exception as e:
            logger.error(f"Encryption failed: {e}")
            raise EncryptionError(f"Failed to encrypt data: {e}")

    def decrypt(self, encrypted_data: str) -> str:
        """
        Decrypt base64-encoded ciphertext.

        Args:
            encrypted_data: Base64-encoded encrypted string

        Returns:
            Decrypted string

        Raises:
            EncryptionError: If decryption fails (wrong key, corrupted data)
            EncryptionKeyError: If encryption is not configured
        """
        if not self.is_available:
            if not CRYPTOGRAPHY_AVAILABLE:
                raise EncryptionError("cryptography package not installed")
            raise EncryptionKeyError("Encryption key not configured")

        # Explicit None check for type safety (mypy)
        if self._fernet is None:
            raise EncryptionKeyError("Encryption not initialized")

        try:
            decrypted = self._fernet.decrypt(encrypted_data.encode("utf-8"))
            return decrypted.decode("utf-8")

        except InvalidToken:
            logger.error("Decryption failed - invalid token (wrong key or corrupted data)")
            raise EncryptionError(
                "Decryption failed. This may indicate: "
                "1) Wrong encryption key, "
                "2) Corrupted data, or "
                "3) Data was not encrypted with this key"
            )
        except Exception as e:
            logger.error(f"Decryption failed: {e}")
            raise EncryptionError(f"Failed to decrypt data: {e}")

    def encrypt_dict(self, data: dict) -> str:
        """Encrypt a dictionary as JSON."""
        return self.encrypt(json.dumps(data, default=str))

    def decrypt_dict(self, encrypted_data: str) -> dict:
        """Decrypt to a dictionary."""
        return json.loads(self.decrypt(encrypted_data))

    def encrypt_file(self, source_path: Path, dest_path: Optional[Path] = None) -> Path:
        """
        Encrypt a file.

        Args:
            source_path: Path to file to encrypt
            dest_path: Destination path. Defaults to source_path + '.enc'

        Returns:
            Path to encrypted file
        """
        if dest_path is None:
            dest_path = source_path.with_suffix(source_path.suffix + ".enc")

        content = source_path.read_text()
        encrypted = self.encrypt(content)

        dest_path.write_text(encrypted)
        dest_path.chmod(0o600)

        logger.debug(f"Encrypted {source_path} -> {dest_path}")
        return dest_path

    def decrypt_file(self, source_path: Path, dest_path: Optional[Path] = None) -> Path:
        """
        Decrypt a file.

        Args:
            source_path: Path to encrypted file
            dest_path: Destination path. Defaults to source_path without '.enc'

        Returns:
            Path to decrypted file
        """
        if dest_path is None:
            if source_path.suffix == ".enc":
                dest_path = source_path.with_suffix("")
            else:
                dest_path = source_path.with_suffix(".dec")

        encrypted = source_path.read_text()
        decrypted = self.decrypt(encrypted)

        dest_path.write_text(decrypted)
        dest_path.chmod(0o600)

        logger.debug(f"Decrypted {source_path} -> {dest_path}")
        return dest_path


class EncryptedJSONStorage:
    """
    Drop-in replacement for JSON file storage with transparent encryption.

    Usage:
        storage = EncryptedJSONStorage(Path("data"), encryption_enabled=True)

        # Write (encrypts automatically)
        storage.write("sessions/user123.json", {"trust_level": "known"})

        # Read (decrypts automatically)
        data = storage.read("sessions/user123.json")
    """

    def __init__(
        self,
        base_path: Path,
        encryption_enabled: bool = False,
        encrypted_storage: Optional[EncryptedStorage] = None,
    ):
        """
        Initialize encrypted JSON storage.

        Args:
            base_path: Base directory for storage
            encryption_enabled: Whether to encrypt files
            encrypted_storage: Custom EncryptedStorage instance
        """
        self.base_path = Path(base_path)
        self.encryption_enabled = encryption_enabled
        self._storage = encrypted_storage or EncryptedStorage()

        # Check if encryption is requested but unavailable — refuse to proceed
        # silently storing plaintext when the user expects encryption.
        if encryption_enabled and not self._storage.is_available:
            raise RuntimeError(
                "Encryption was requested but is not available. "
                "Set FAMILIAR_ENCRYPTION_KEY environment variable or install "
                "the 'cryptography' package. Refusing to store data in plaintext."
            )

        self.base_path.mkdir(parents=True, exist_ok=True)

    def _get_path(self, relative_path: str) -> Path:
        """Get full path, adding .enc suffix if encryption enabled."""
        path = self.base_path / relative_path
        if self.encryption_enabled:
            path = path.with_suffix(path.suffix + ".enc")
        return path

    def write(self, relative_path: str, data: dict) -> Path:
        """
        Write data to file, encrypting if enabled.

        Args:
            relative_path: Path relative to base_path
            data: Dictionary to write

        Returns:
            Path to written file
        """
        path = self._get_path(relative_path)
        path.parent.mkdir(parents=True, exist_ok=True)

        json_str = json.dumps(data, indent=2, default=str)

        if self.encryption_enabled:
            content = self._storage.encrypt(json_str)
        else:
            content = json_str

        # Atomic write
        tmp_path = path.with_suffix(path.suffix + ".tmp")
        try:
            tmp_path.write_text(content)
            tmp_path.chmod(0o600)
            tmp_path.replace(path)
        except Exception:
            if tmp_path.exists():
                tmp_path.unlink()
            raise

        return path

    def read(self, relative_path: str) -> Optional[dict]:
        """
        Read data from file, decrypting if needed.

        Args:
            relative_path: Path relative to base_path

        Returns:
            Dictionary or None if file doesn't exist
        """
        path = self._get_path(relative_path)

        # Also check for unencrypted version (migration case)
        unencrypted_path = self.base_path / relative_path

        if not path.exists():
            if unencrypted_path.exists() and self.encryption_enabled:
                # Migrate: encrypt existing unencrypted file
                logger.info(f"Migrating {unencrypted_path} to encrypted storage")
                data = json.loads(unencrypted_path.read_text())
                self.write(relative_path, data)
                unencrypted_path.unlink()  # Remove unencrypted version
                return data
            elif unencrypted_path.exists():
                return json.loads(unencrypted_path.read_text())
            return None

        content = path.read_text()

        if self.encryption_enabled:
            json_str = self._storage.decrypt(content)
        else:
            json_str = content

        return json.loads(json_str)

    def delete(self, relative_path: str) -> bool:
        """
        Delete a file.

        Args:
            relative_path: Path relative to base_path

        Returns:
            True if file was deleted, False if it didn't exist
        """
        path = self._get_path(relative_path)

        if path.exists():
            path.unlink()
            return True
        return False

    def exists(self, relative_path: str) -> bool:
        """Check if file exists."""
        return self._get_path(relative_path).exists()

    def list_files(self, relative_dir: str = "") -> list[str]:
        """List files in directory."""
        dir_path = self.base_path / relative_dir
        if not dir_path.exists():
            return []

        files = []
        for path in dir_path.iterdir():
            if path.is_file():
                name = path.name
                if self.encryption_enabled and name.endswith(".enc"):
                    name = name[:-4]  # Remove .enc suffix
                files.append(name)

        return files


def generate_key() -> str:
    """
    Generate a new Fernet encryption key.

    Returns:
        Base64-encoded Fernet key suitable for FAMILIAR_ENCRYPTION_KEY
    """
    if not CRYPTOGRAPHY_AVAILABLE:
        raise EncryptionError("cryptography package required.")

    return Fernet.generate_key().decode()


def rotate_key(
    old_storage: EncryptedStorage, new_storage: EncryptedStorage, data_dir: Path
) -> dict:
    """
    Re-encrypt all data with a new key.

    Args:
        old_storage: Storage with old key
        new_storage: Storage with new key
        data_dir: Directory containing encrypted files

    Returns:
        Summary of rotated files
    """
    results = {"success": [], "failed": [], "skipped": []}

    for path in data_dir.rglob("*.enc"):
        try:
            # Decrypt with old key
            decrypted = old_storage.decrypt(path.read_text())

            # Encrypt with new key
            encrypted = new_storage.encrypt(decrypted)

            # Write back
            path.write_text(encrypted)
            results["success"].append(str(path))

        except EncryptionError as e:
            logger.error(f"Failed to rotate {path}: {e}")
            results["failed"].append(str(path))
        except Exception as e:
            logger.error(f"Unexpected error rotating {path}: {e}")
            results["failed"].append(str(path))

    logger.info(
        f"Key rotation complete: {len(results['success'])} succeeded, "
        f"{len(results['failed'])} failed"
    )

    return results


# CLI interface for key management
if __name__ == "__main__":
    import sys

    if len(sys.argv) < 2:
        print("Familiar Encryption Utility")
        print()
        print("Commands:")
        print("  generate-key      Generate a new encryption key")
        print("  test-key          Test if configured key works")
        print("  encrypt <file>    Encrypt a file")
        print("  decrypt <file>    Decrypt a file")
        print()
        print("Environment:")
        print("  FAMILIAR_ENCRYPTION_KEY  Set this to your Fernet key")
        sys.exit(0)

    command = sys.argv[1]

    if command == "generate-key":
        key = generate_key()
        print("Generated encryption key:")
        print(key)
        print()
        print("Add to your .env file:")
        print(f"FAMILIAR_ENCRYPTION_KEY={key}")

    elif command == "test-key":
        try:
            storage = EncryptedStorage()
            if storage.is_available:
                # Test round-trip
                test_data = "Familiar encryption test"
                encrypted = storage.encrypt(test_data)
                decrypted = storage.decrypt(encrypted)
                assert decrypted == test_data
                print("✓ Encryption key is valid and working")
            else:
                print("✗ No encryption key configured")
                print("  Set FAMILIAR_ENCRYPTION_KEY environment variable")
                sys.exit(1)
        except Exception as e:
            print(f"✗ Encryption test failed: {e}")
            sys.exit(1)

    elif command == "encrypt" and len(sys.argv) > 2:
        storage = EncryptedStorage()
        source = Path(sys.argv[2])
        dest = storage.encrypt_file(source)
        print(f"Encrypted: {source} -> {dest}")

    elif command == "decrypt" and len(sys.argv) > 2:
        storage = EncryptedStorage()
        source = Path(sys.argv[2])
        dest = storage.decrypt_file(source)
        print(f"Decrypted: {source} -> {dest}")

    else:
        print(f"Unknown command: {command}")
        sys.exit(1)
