from cryptography.fernet import Fernet
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
from cryptography.hazmat.backends import default_backend
import base64
import os
from typing import Optional
import logging
from ..config.settings import settings

logger = logging.getLogger(__name__)

class EncryptionService:
    """Service for data encryption at rest"""
    
    def __init__(self, encryption_key: Optional[str] = None):
        """
        Initialize encryption service
        
        Args:
            encryption_key: Base64-encoded encryption key. If not provided, will use ENCRYPTION_KEY from settings or env var
        """
        if encryption_key:
            self.key = encryption_key.encode()
        else:
            # Try to get key from settings first (loads from .env), then fall back to os.getenv
            key_str = getattr(settings, 'ENCRYPTION_KEY', None) or os.getenv('ENCRYPTION_KEY')
            if not key_str or key_str.strip() == '':
                # Generate a key if not provided (for development only)
                logger.warning("ENCRYPTION_KEY not set. Generating temporary key. This should be set in production!")
                key = Fernet.generate_key()
                self.key = key
            else:
                self.key = key_str.encode()
        
        try:
            self.cipher = Fernet(self.key)
        except Exception as e:
            logger.error(f"Failed to initialize encryption: {str(e)}")
            raise
    
    def encrypt(self, data: str) -> str:
        """Encrypt a string"""
        try:
            if not data:
                return data
            encrypted = self.cipher.encrypt(data.encode())
            return base64.urlsafe_b64encode(encrypted).decode()
        except Exception as e:
            logger.error(f"Encryption failed: {str(e)}")
            raise
    
    def decrypt(self, encrypted_data: str) -> str:
        """Decrypt a string"""
        try:
            if not encrypted_data:
                return encrypted_data
            decoded = base64.urlsafe_b64decode(encrypted_data.encode())
            decrypted = self.cipher.decrypt(decoded)
            return decrypted.decode()
        except Exception as e:
            logger.error(f"Decryption failed: {str(e)}")
            raise
    
    def encrypt_dict(self, data: dict) -> dict:
        """Encrypt sensitive fields in a dictionary"""
        encrypted = {}
        sensitive_fields = ['password', 'token', 'secret', 'key', 'api_key', 'access_token', 'refresh_token']
        
        for key, value in data.items():
            if any(sensitive in key.lower() for sensitive in sensitive_fields):
                if isinstance(value, str):
                    encrypted[key] = self.encrypt(value)
                else:
                    encrypted[key] = value
            else:
                encrypted[key] = value
        
        return encrypted
    
    @staticmethod
    def generate_key() -> str:
        """Generate a new encryption key"""
        key = Fernet.generate_key()
        return key.decode()

# Global instance
encryption_service = EncryptionService()

