# Copyright (c) 2025 Patrick Motsch
# All rights reserved.
"""
Utility module for configuration management.

This module provides a global APP_CONFIG object for accessing configuration from both 
config.ini files and environment variables stored in .env files, using a flat structure.
"""

import os
import logging
import json
import base64
import time
import threading
from typing import Any, Dict, Optional, Tuple
from pathlib import Path
from cryptography.fernet import Fernet
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC

# Set up basic logging for configuration loading
logging.basicConfig(
    level=logging.WARNING,
    format='%(asctime)s - %(levelname)s - %(name)s - %(message)s',
    handlers=[logging.StreamHandler()]
)

# Configure logger
logger = logging.getLogger(__name__)

class Configuration:
    """
    Configuration class with attribute-style access to flattened configuration.
    """
    def __init__(self):
        """Initialize the configuration object"""
        self._data = {}
        self._configFilePath = None
        self._envFilePath = None
        self._configMtime = 0
        self._envMtime = 0
        self.refresh()
    
    def refresh(self):
        """Reload configuration from files"""
        self._loadConfig()
        self._loadEnv()
        logger.info("Configuration refreshed")
    
    def _loadConfig(self):
        """Load configuration from config.ini file in flattened format"""
        # Find config.ini file in the gateway directory
        configPath = Path(__file__).parent.parent.parent / 'config.ini'
        if not configPath.exists():
            logger.warning(f"Configuration file not found at {configPath.absolute()}")
            return
        
        self._configFilePath = configPath
        currentMtime = os.path.getmtime(configPath)
        
        # Skip if file hasn't changed
        if currentMtime <= self._configMtime:
            return
            
        self._configMtime = currentMtime

        try:
            with open(configPath, 'r', encoding='utf-8') as f:
                lines = f.readlines()
                
            i = 0
            while i < len(lines):
                line = lines[i].strip()
                
                # Skip empty lines and comments
                if not line or line.startswith('#'):
                    i += 1
                    continue
                    
                # Parse key-value pairs
                if '=' in line:
                    key, value = line.split('=', 1)
                    key = key.strip()
                    value = value.strip()
                    
                    # Check if value starts with { (JSON object)
                    if value.startswith('{'):
                        # Collect all lines until we find the closing }
                        json_lines = [value]
                        i += 1
                        brace_count = value.count('{') - value.count('}')
                        
                        while i < len(lines) and brace_count > 0:
                            json_lines.append(lines[i].rstrip('\n'))
                            brace_count += lines[i].count('{') - lines[i].count('}')
                            i += 1
                        
                        # Join all lines and parse as JSON
                        value = '\n'.join(json_lines)
                        i -= 1  # Adjust for the loop increment
                    
                    # Add to data dictionary
                    self._data[key] = value
                
                i += 1
                
                    
        except Exception as e:
            logger.error(f"Error loading configuration: {e}")
    
    def _loadEnv(self):
        """Load environment variables from .env file"""
        # Find .env file in the gateway directory
        envPath = Path(__file__).parent.parent.parent / '.env'
        if not envPath.exists():
            logger.warning(f"Environment file not found at {envPath.absolute()}")
            return
        
        self._envFilePath = envPath
        currentMtime = os.path.getmtime(envPath)
        
        # Skip if file hasn't changed
        if currentMtime <= self._envMtime:
            return
            
        self._envMtime = currentMtime
        
        try:
            with open(envPath, 'r') as f:
                lines = f.readlines()
            
            i = 0
            while i < len(lines):
                line = lines[i].strip()
                
                # Skip empty lines and comments
                if not line or line.startswith('#'):
                    i += 1
                    continue
                
                # Parse key-value pairs
                if '=' in line:
                    key, value = line.split('=', 1)
                    key = key.strip()
                    value = value.strip()
                    
                    # Check if value starts with { (JSON object)
                    if value.startswith('{'):
                        # Collect all lines until we find the closing }
                        json_lines = [value]
                        i += 1
                        brace_count = value.count('{') - value.count('}')
                        
                        while i < len(lines) and brace_count > 0:
                            json_lines.append(lines[i].rstrip('\n'))
                            brace_count += lines[i].count('{') - lines[i].count('}')
                            i += 1
                        
                        # Join all lines and create the full JSON value
                        full_json_value = '\n'.join(json_lines)
                        self._data[key] = full_json_value
                    else:
                        # Single line value
                        self._data[key] = value
                
                i += 1
            
            logger.info(f"Loaded environment variables from {envPath.absolute()}")
            
            # Also load system environment variables (don't override existing)
            for key, value in os.environ.items():
                if key not in self._data:
                    self._data[key] = value
                    
        except Exception as e:
            logger.error(f"Error loading environment variables: {e}")
    
    def checkForUpdates(self):
        """Check if configuration files have changed and reload if necessary"""
        if self._configFilePath and os.path.exists(self._configFilePath):
            currentMtime = os.path.getmtime(self._configFilePath)
            if currentMtime > self._configMtime:
                logger.info("Config file has changed, reloading...")
                self._loadConfig()
                
        if self._envFilePath and os.path.exists(self._envFilePath):
            currentMtime = os.path.getmtime(self._envFilePath)
            if currentMtime > self._envMtime:
                logger.info("Environment file has changed, reloading...")
                self._loadEnv()
    
    def get(self, key: str, default: Any = None, user_id: str = "system") -> Any:
        """Get configuration value with optional default"""
        self.checkForUpdates()  # Check for file changes
        
        if key in self._data:
            value = self._data[key]
            # Handle secrets (keys ending with _SECRET)
            if key.endswith("_SECRET"):
                # Log audit event for secret key access
                try:
                    logging.getLogger("audit.fallback").info(
                        "AUDIT | %s | %s | system | key | %s | Key: %s",
                        time.time(), user_id, "decode", key
                    )
                except Exception:
                    pass
                
                if value.startswith("{") and value.endswith("}"):
                    # Handle JSON secrets (keys ending with _API_KEY that contain JSON)
                    return handleSecretJson(value, userId=user_id, keyName=key)
                else:
                    return handleSecretText(value, userId=user_id, keyName=key)
            return value
        return default
    
    def __getattr__(self, name: str) -> Any:
        """Enable attribute-style access to configuration"""
        self.checkForUpdates()  # Check for file changes
        
        value = self.get(name, user_id="system")
        if value is None:
            raise AttributeError(f"Configuration key '{name}' not found")
        return value
    
    def __dir__(self) -> list:
        """Support auto-completion of attributes"""
        self.checkForUpdates()  # Check for file changes
        return list(self._data.keys()) + super().__dir__()
    
    def set(self, key: str, value: Any) -> None:
        """Set a configuration value (for testing/overrides)"""
        self._data[key] = value

def handleSecretText(value: str, userId: str = "system", keyName: str = "unknown") -> str:
    """
    Handle secret values with encryption/decryption support.
    
    Args:
        value: The secret value to handle (may be encrypted)
        userId: The user ID making the request (default: "system")
        keyName: The name of the key being decrypted (default: "unknown")
        
    Returns:
        str: Processed secret value (decrypted if encrypted)
    """
    if _isEncryptedValue(value):
        return decryptValue(value, userId, keyName)
    return value

def handleSecretJson(value: str, userId: str = "system", keyName: str = "unknown") -> str:
    """
    Handle JSON secret values (like Google service account keys) with encryption/decryption support.
    Validates that the value is valid JSON after decryption.
    
    Args:
        value: The JSON secret value to handle (may be encrypted)
        userId: The user ID making the request (default: "system")
        keyName: The name of the key being decrypted (default: "unknown")
        
    Returns:
        str: Processed JSON secret value (decrypted if encrypted)
        
    Raises:
        ValueError: If the value is not valid JSON after decryption
    """
    # Decrypt if encrypted
    if _isEncryptedValue(value):
        decryptedValue = decryptValue(value, userId, keyName)
    else:
        decryptedValue = value
    
    try:
        # Validate that it's valid JSON
        json.loads(decryptedValue)
        return decryptedValue
    except json.JSONDecodeError as e:
        raise ValueError(f"Invalid JSON in secret value: {e}")

# Global rate limiting tracking
# Structure: {user_id: {key_name: [timestamps]}}
_decryption_attempts = {}

# Process-wide plaintext cache for decrypted secrets.
# Key: the encrypted ciphertext (which already includes env prefix).
# Value: (expiresAtMonotonic, plaintext).
# TTL is short enough that key rotation propagates quickly, long enough that
# hot DB-init paths (every API call building a connector) don't blow the
# decryption rate limit. 60s is a deliberate compromise.
_DECRYPTION_CACHE_TTL_S = 60.0
_decryption_cache: Dict[str, Tuple[float, str]] = {}
_decryption_cache_lock = threading.Lock()

def _getMasterKey(envType: str = None) -> bytes:
    """
    Get the master key for the specified environment.
    
    Args:
        envType: The environment type (dev, int, prod, etc.). If None, uses current config.
    
    Returns:
        bytes: The master key for encryption/decryption
        
    Raises:
        ValueError: If no master key is found
    """
    # Get the key location from config
    keyLocation = APP_CONFIG.get('APP_KEY_SYSVAR')
    if envType is None:
        envType = APP_CONFIG.get('APP_ENV_TYPE', 'dev')
    
    if not keyLocation:
        raise ValueError("APP_KEY_SYSVAR not configured")
    
    # First try to get from environment variable
    masterKey = os.environ.get(keyLocation)
    
    if masterKey:
        # If found in environment, use it directly
        return masterKey.encode('utf-8')
    
    # If not in environment, try to read from file
    if os.path.exists(keyLocation):
        try:
            with open(keyLocation, 'r') as f:
                content = f.read().strip()
                
            # Parse the key file format: env = key
            lines = content.split('\n')
            for line in lines:
                line = line.strip()
                if not line or line.startswith('#'):
                    continue
                    
                if '=' in line:
                    keyEnv, keyValue = line.split('=', 1)
                    keyEnv = keyEnv.strip()
                    keyValue = keyValue.strip()
                    
                    if keyEnv == envType:
                        return keyValue.encode('utf-8')
            
            raise ValueError(f"No key found for environment '{envType}' in {keyLocation}")
            
        except Exception as e:
            raise ValueError(f"Error reading key file {keyLocation}: {e}")
    
    raise ValueError(f"Master key not found. Checked environment variable '{keyLocation}' and file path")

def _deriveEncryptionKey(masterKey: bytes) -> bytes:
    """
    Derive a 32-byte encryption key from the master key using PBKDF2.
    
    Args:
        masterKey: The master key bytes
        
    Returns:
        bytes: 32-byte derived key suitable for Fernet
    """
    # Use a fixed salt for consistency (in production, consider using a random salt stored separately)
    salt = b'poweron_config_salt_2025'
    
    kdf = PBKDF2HMAC(
        algorithm=hashes.SHA256(),
        length=32,
        salt=salt,
        iterations=100000,
    )
    
    return base64.urlsafe_b64encode(kdf.derive(masterKey))

def _isEncryptedValue(value: str) -> bool:
    """
    Check if a value is encrypted (starts with environment-specific prefix).
    
    Args:
        value: The value to check
        
    Returns:
        bool: True if encrypted, False otherwise
    """
    if not value or not isinstance(value, str):
        return False
    
    # Check for any environment-specific encryption prefixes
    return (value.startswith('DEV_ENC:') or 
            value.startswith('INT_ENC:') or 
            value.startswith('PROD_ENC:') or
            value.startswith('TEST_ENC:') or
            value.startswith('STAGING_ENC:'))

def _getEncryptionPrefix(envType: str) -> str:
    """
    Get the encryption prefix for the given environment type.
    
    Args:
        envType: The environment type (dev, int, prod, etc.)
        
    Returns:
        str: The encryption prefix
    """
    return f"{envType.upper()}_ENC:"

def _checkDecryptionRateLimit(userId: str, keyName: str, maxPerSecond: int = 10) -> bool:
    """
    Check if decryption is allowed based on rate limiting (max 10 per second per user per key).
    
    Args:
        userId: The user ID making the request
        keyName: The name of the key being decrypted
        maxPerSecond: Maximum decryptions per second (default: 10)
        
    Returns:
        bool: True if allowed, False if rate limited
    """
    currentTime = time.time()
    
    # Initialize tracking for this user if not exists
    if userId not in _decryption_attempts:
        _decryption_attempts[userId] = {}
    
    # Initialize tracking for this key if not exists
    if keyName not in _decryption_attempts[userId]:
        _decryption_attempts[userId][keyName] = []
    
    # Clean old attempts (older than 1 second)
    _decryption_attempts[userId][keyName] = [
        timestamp for timestamp in _decryption_attempts[userId][keyName]
        if currentTime - timestamp < 1.0
    ]
    
    # Check if we're within rate limit
    if len(_decryption_attempts[userId][keyName]) >= maxPerSecond:
        logger.warning(f"Decryption rate limit exceeded for user '{userId}' key '{keyName}' ({maxPerSecond}/sec)")
        return False
    
    # Record this attempt
    _decryption_attempts[userId][keyName].append(currentTime)
    return True

def encryptValue(value: str, envType: str = None, userId: str = "system", keyName: str = "unknown") -> str:
    """
    Encrypt a value using the master key for the specified environment.
    
    Args:
        value: The plain text value to encrypt
        envType: The environment type (dev, int, prod). If None, uses current environment.
        userId: The user ID making the request (default: "system")
        keyName: The name of the key being encrypted (default: "unknown")
        
    Returns:
        str: The encrypted value with prefix
        
    Raises:
        ValueError: If encryption fails
    """
    if envType is None:
        envType = APP_CONFIG.get('APP_ENV_TYPE', 'dev')
    
    try:
        masterKey = _getMasterKey(envType)
        derivedKey = _deriveEncryptionKey(masterKey)
        fernet = Fernet(derivedKey)
        
        # Encrypt the value
        encryptedBytes = fernet.encrypt(value.encode('utf-8'))
        encryptedB64 = base64.urlsafe_b64encode(encryptedBytes).decode('utf-8')
        
        # Add environment prefix
        prefix = _getEncryptionPrefix(envType)
        encryptedValue = f"{prefix}{encryptedB64}"
        
        # Log audit event for encryption
        try:
            logging.getLogger("audit.fallback").info(
                "AUDIT | %s | %s | system | key | %s | Key: %s",
                time.time(), userId, "encrypt", keyName
            )
        except Exception:
            pass
        
        return encryptedValue
        
    except Exception as e:
        raise ValueError(f"Encryption failed: {e}")

def decryptValue(encryptedValue: str, userId: str = "system", keyName: str = "unknown") -> str:
    """
    Decrypt a value using the master key for the current environment.

    A short-lived plaintext cache (TTL `_DECRYPTION_CACHE_TTL_S`) is consulted
    first. The 10/sec rate-limit on cache misses still protects against
    brute-force attacks; cache HITS bypass it because they are not actual
    cryptographic operations — they just return the result of an earlier
    successful decrypt. Without this cache, hot paths like
    `mainBackgroundJobService._getDb()` (called per RAG inventory poll AND
    per walker DB call) trigger the rate limit and surface as
    "Decryption rate limit exceeded for user 'system' key 'DB_PASSWORD_SECRET'"
    ERRORs in the RAG inventory UI route.

    Args:
        encryptedValue: The encrypted value with prefix
        userId: The user ID making the request (default: "system")
        keyName: The name of the key being decrypted (default: "unknown")

    Returns:
        str: The decrypted plain text value

    Raises:
        ValueError: If decryption fails
    """
    if not _isEncryptedValue(encryptedValue):
        return encryptedValue  # Return as-is if not encrypted

    # Cache lookup BEFORE the rate-limit check: a cache hit is not a new
    # cryptographic operation and must not be throttled.
    now = time.monotonic()
    with _decryption_cache_lock:
        cached = _decryption_cache.get(encryptedValue)
        if cached is not None and cached[0] > now:
            return cached[1]

    # Cache miss → real decrypt → apply rate limit.
    if not _checkDecryptionRateLimit(userId, keyName, maxPerSecond=10):
        raise ValueError(f"Decryption rate limit exceeded for user '{userId}' key '{keyName}' (10/sec)")

    try:
        # Extract environment type from prefix
        if encryptedValue.startswith('DEV_ENC:'):
            envType = 'dev'
            prefix = 'DEV_ENC:'
        elif encryptedValue.startswith('INT_ENC:'):
            envType = 'int'
            prefix = 'INT_ENC:'
        elif encryptedValue.startswith('PROD_ENC:'):
            envType = 'prod'
            prefix = 'PROD_ENC:'
        elif encryptedValue.startswith('TEST_ENC:'):
            envType = 'test'
            prefix = 'TEST_ENC:'
        elif encryptedValue.startswith('STAGING_ENC:'):
            envType = 'staging'
            prefix = 'STAGING_ENC:'
        else:
            raise ValueError(f"Invalid encryption prefix. Expected DEV_ENC:, INT_ENC:, PROD_ENC:, TEST_ENC:, or STAGING_ENC:")
        
        encryptedPart = encryptedValue[len(prefix):]
        
        # Get master key for the specific environment and derive encryption key
        masterKey = _getMasterKey(envType)
        derivedKey = _deriveEncryptionKey(masterKey)
        fernet = Fernet(derivedKey)
        
        # Decode and decrypt
        encryptedBytes = base64.urlsafe_b64decode(encryptedPart.encode('utf-8'))
        decryptedBytes = fernet.decrypt(encryptedBytes)
        decryptedValue = decryptedBytes.decode('utf-8')

        # Log audit event for decryption
        try:
            logging.getLogger("audit.fallback").info(
                "AUDIT | %s | %s | system | key | %s | Key: %s",
                time.time(), userId, "decrypt", keyName
            )
        except Exception:
            pass

        # Populate cache so subsequent reads of the same ciphertext don't
        # re-decrypt (and don't consume rate-limit budget).
        with _decryption_cache_lock:
            _decryption_cache[encryptedValue] = (
                time.monotonic() + _DECRYPTION_CACHE_TTL_S,
                decryptedValue,
            )

        return decryptedValue

    except Exception as e:
        raise ValueError(f"Decryption failed: {e}")


def clearDecryptionCache() -> None:
    """Drop all cached plaintext secrets. Call after key rotation or in tests."""
    with _decryption_cache_lock:
        _decryption_cache.clear()

# Create the global APP_CONFIG instance
APP_CONFIG = Configuration()