"""Encrypted storage for sensitive data.

This module provides secure storage for sensitive data such as transcripts,
user preferences, and other confidential information.
"""

import json
import os
from pathlib import Path
from typing import Any, Optional, Union

from cryptography.fernet import Fernet

from src.logging import get_logger

logger = get_logger(__name__)


class EncryptedStorage:
    """Secure storage for sensitive data with encryption."""

    def __init__(self, storage_path: Path, key_path: Path):
        """Initialize encrypted storage.

        Args:
            storage_path: Directory to store encrypted files
            key_path: Path to encryption key file
        """
        self.storage_path = storage_path
        self.key_path = key_path
        self.fernet = None

        # Ensure storage directory exists
        self.storage_path.mkdir(parents=True, exist_ok=True)

        # Initialize encryption key
        self._init_encryption()

    def _init_encryption(self) -> None:
        """Initialize or load encryption key."""
        if not self.key_path.exists():
            # Generate new key
            key = Fernet.generate_key()
            with open(self.key_path, "wb") as f:
                f.write(key)
            # Set permissions to owner-only
            os.chmod(self.key_path, 0o600)
            logger.info(f"Generated new encryption key: {self.key_path}")

        # Load key
        with open(self.key_path, "rb") as f:
            key = f.read()
        self.fernet = Fernet(key)
        logger.debug("Encryption key loaded successfully")

    def store(self, key: str, data: Any) -> bool:
        """Store data with encryption.

        Args:
            key: Unique identifier for the data
            data: Data to store (string, dict, list, or bytes)

        Returns:
            True if successful, False otherwise
        """
        try:
            # Serialize data to JSON if not bytes
            if isinstance(data, bytes):
                serialized_data = data
            else:
                serialized_data = json.dumps(data, ensure_ascii=False).encode('utf-8')

            # Encrypt data
            encrypted_data = self.fernet.encrypt(serialized_data)

            # Save to file
            file_path = self.storage_path / f"{key}.enc"
            with open(file_path, "wb") as f:
                f.write(encrypted_data)

            # Set permissions to owner-only
            os.chmod(file_path, 0o600)

            logger.debug(f"Stored encrypted data for key: {key}")
            return True
        except Exception as e:
            logger.error(f"Error storing data for key {key}: {str(e)}")
            return False

    def retrieve(self, key: str) -> Optional[Any]:
        """Retrieve data with decryption.

        Args:
            key: Unique identifier for the data

        Returns:
            Retrieved data or None if not found or error
        """
        try:
            file_path = self.storage_path / f"{key}.enc"
            if not file_path.exists():
                return None

            # Read encrypted data
            with open(file_path, "rb") as f:
                encrypted_data = f.read()

            # Decrypt data
            decrypted_data = self.fernet.decrypt(encrypted_data)

            # Try to deserialize as JSON first
            try:
                return json.loads(decrypted_data.decode('utf-8'))
            except (json.JSONDecodeError, UnicodeDecodeError):
                # If JSON deserialization fails, return as bytes
                return decrypted_data

        except Exception as e:
            logger.error(f"Error retrieving data for key {key}: {str(e)}")
            return None

    def delete(self, key: str) -> bool:
        """Delete stored data.

        Args:
            key: Unique identifier for the data

        Returns:
            True if successful, False otherwise
        """
        try:
            file_path = self.storage_path / f"{key}.enc"
            if file_path.exists():
                file_path.unlink()
                logger.debug(f"Deleted data for key: {key}")
                return True
            return False
        except Exception as e:
            logger.error(f"Error deleting data for key {key}: {str(e)}")
            return False

    def list_keys(self) -> list[str]:
        """List all stored keys.

        Returns:
            List of stored keys
        """
        try:
            keys = []
            for file_path in self.storage_path.glob("*.enc"):
                key = file_path.stem  # Remove .enc extension
                keys.append(key)
            return keys
        except Exception as e:
            logger.error(f"Error listing keys: {str(e)}")
            return []

    def clear(self) -> bool:
        """Clear all stored data.

        Returns:
            True if successful, False otherwise
        """
        try:
            for file_path in self.storage_path.glob("*.enc"):
                file_path.unlink()
            logger.info("Cleared all encrypted data")
            return True
        except Exception as e:
            logger.error(f"Error clearing data: {str(e)}")
            return False

    def exists(self, key: str) -> bool:
        """Check if data exists for a key.

        Args:
            key: Unique identifier for the data

        Returns:
            True if data exists, False otherwise
        """
        file_path = self.storage_path / f"{key}.enc"
        return file_path.exists()

    def get_size(self, key: str) -> Optional[int]:
        """Get the size of stored data in bytes.

        Args:
            key: Unique identifier for the data

        Returns:
            Size in bytes or None if not found
        """
        try:
            file_path = self.storage_path / f"{key}.enc"
            if file_path.exists():
                return file_path.stat().st_size
            return None
        except Exception as e:
            logger.error(f"Error getting size for key {key}: {str(e)}")
            return None


def generate_encryption_key() -> bytes:
    """Generate a new encryption key.

    Returns:
        New encryption key
    """
    return Fernet.generate_key()


def encrypt_data(data: Any, key: bytes) -> bytes:
    """Encrypt data with the given key.

    Args:
        data: Data to encrypt
        key: Encryption key

    Returns:
        Encrypted data

    Raises:
        Exception: If encryption fails
    """
    fernet = Fernet(key)
    
    # Serialize data to JSON if not bytes
    if isinstance(data, bytes):
        serialized_data = data
    else:
        serialized_data = json.dumps(data, ensure_ascii=False).encode('utf-8')
    
    return fernet.encrypt(serialized_data)


def decrypt_data(encrypted_data: bytes, key: bytes) -> Any:
    """Decrypt data with the given key.

    Args:
        encrypted_data: Encrypted data
        key: Encryption key

    Returns:
        Decrypted data

    Raises:
        Exception: If decryption fails
    """
    fernet = Fernet(key)
    decrypted_data = fernet.decrypt(encrypted_data)
    
    # Try to deserialize as JSON first
    try:
        return json.loads(decrypted_data.decode('utf-8'))
    except (json.JSONDecodeError, UnicodeDecodeError):
        # If JSON deserialization fails, return as bytes
        return decrypted_data


def create_secure_storage(base_path: Path = Path("~/.trax/secure").expanduser()) -> EncryptedStorage:
    """Create a secure storage instance with default paths.

    Args:
        base_path: Base directory for secure storage

    Returns:
        EncryptedStorage instance
    """
    storage_path = base_path / "data"
    key_path = base_path / "key.bin"
    
    return EncryptedStorage(storage_path, key_path)
