"""
Token encryption utility for securing OAuth tokens at rest
"""
import os
import base64
import json
from cryptography.fernet import Fernet
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2


class TokenEncryption:
    """Handles encryption/decryption of sensitive tokens"""
    
    def __init__(self):
        # Get encryption key from environment or generate one
        encryption_key = os.getenv('TOKEN_ENCRYPTION_KEY')
        
        if not encryption_key:
            # Generate a key from a password if not provided
            password = os.getenv('SECRET_KEY', 'default-dev-key').encode()
            salt = os.getenv('TOKEN_SALT', 'rateright-tokens-2025').encode()
            kdf = PBKDF2(
                algorithm=hashes.SHA256(),
                length=32,
                salt=salt,
                iterations=100000,
            )
            key = base64.urlsafe_b64encode(kdf.derive(password))
            self.cipher = Fernet(key)
        else:
            self.cipher = Fernet(encryption_key.encode())
    
    def encrypt_token(self, token_data):
        """
        Encrypt token data (dict) to encrypted string
        
        Args:
            token_data: Dictionary containing OAuth tokens
            
        Returns:
            Encrypted string safe for database storage
        """
        if not token_data:
            return None
            
        try:
            # Convert dict to JSON string
            json_str = json.dumps(token_data)
            # Encrypt the JSON string
            encrypted = self.cipher.encrypt(json_str.encode())
            # Return base64 encoded string for safe storage
            return base64.urlsafe_b64encode(encrypted).decode('utf-8')
        except Exception as e:
            print(f"Error encrypting token: {e}")
            return None
    
    def decrypt_token(self, encrypted_token):
        """
        Decrypt encrypted token string back to dict
        
        Args:
            encrypted_token: Encrypted string from database
            
        Returns:
            Dictionary containing OAuth tokens
        """
        if not encrypted_token:
            return None
            
        try:
            # Decode from base64
            encrypted_bytes = base64.urlsafe_b64decode(encrypted_token.encode('utf-8'))
            # Decrypt the data
            decrypted = self.cipher.decrypt(encrypted_bytes)
            # Parse JSON back to dict
            return json.loads(decrypted.decode('utf-8'))
        except Exception as e:
            print(f"Error decrypting token: {e}")
            return None


# Global instance
token_encryptor = TokenEncryption()