#!/usr/bin/env python3
"""
GENESIS SECRETS MANAGER
========================
Secure secrets management with encryption and rotation.

Features:
    - AES-256 encryption at rest
    - Environment variable integration
    - Secret rotation support
    - Audit logging
    - Namespace isolation
    - Lazy loading
    - Multiple backends (file, env, vault-compatible)

Usage:
    secrets = SecretsManager()
    secrets.set("api_key", "secret_value")
    value = secrets.get("api_key")

    # With encryption
    secrets = SecretsManager(encryption_key=key)
    secrets.set("db_password", "very_secret", encrypt=True)
"""

import base64
import hashlib
import hmac
import json
import os
import secrets as py_secrets
import threading
import time
from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from datetime import datetime, timedelta
from enum import Enum
from pathlib import Path
from typing import Dict, List, Any, Optional, Callable, Set, Union

# Try to import cryptography for AES, fallback to basic obfuscation
try:
    from cryptography.fernet import Fernet
    from cryptography.hazmat.primitives import hashes
    from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
    CRYPTO_AVAILABLE = True
except ImportError:
    CRYPTO_AVAILABLE = False
    Fernet = None


class SecretType(Enum):
    """Types of secrets."""
    STRING = "string"
    API_KEY = "api_key"
    PASSWORD = "password"
    CERTIFICATE = "certificate"
    JSON = "json"
    BINARY = "binary"


@dataclass
class SecretMetadata:
    """Metadata about a secret."""
    name: str
    secret_type: SecretType = SecretType.STRING
    created_at: str = ""
    updated_at: str = ""
    expires_at: Optional[str] = None
    version: int = 1
    encrypted: bool = False
    namespace: str = "default"
    description: str = ""
    tags: List[str] = field(default_factory=list)

    def to_dict(self) -> Dict:
        return {
            "name": self.name,
            "type": self.secret_type.value,
            "created_at": self.created_at,
            "updated_at": self.updated_at,
            "expires_at": self.expires_at,
            "version": self.version,
            "encrypted": self.encrypted,
            "namespace": self.namespace,
            "description": self.description,
            "tags": self.tags
        }


@dataclass
class SecretEntry:
    """A secret entry with value and metadata."""
    metadata: SecretMetadata
    value: str  # Encrypted or plain

    def is_expired(self) -> bool:
        if not self.metadata.expires_at:
            return False
        expires = datetime.fromisoformat(self.metadata.expires_at)
        return datetime.now() > expires


@dataclass
class AuditEntry:
    """Audit log entry for secret access."""
    timestamp: str
    action: str  # get, set, delete, rotate
    secret_name: str
    namespace: str
    success: bool
    caller: Optional[str] = None
    error: Optional[str] = None


class SecretBackend(ABC):
    """Abstract backend for secret storage."""

    @abstractmethod
    def get(self, name: str, namespace: str = "default") -> Optional[str]:
        """Get a secret value."""
        pass

    @abstractmethod
    def set(self, name: str, value: str, namespace: str = "default") -> bool:
        """Set a secret value."""
        pass

    @abstractmethod
    def delete(self, name: str, namespace: str = "default") -> bool:
        """Delete a secret."""
        pass

    @abstractmethod
    def list(self, namespace: str = "default") -> List[str]:
        """List secret names in namespace."""
        pass


class EnvironmentBackend(SecretBackend):
    """Backend that uses environment variables."""

    def __init__(self, prefix: str = "GENESIS_SECRET_"):
        self.prefix = prefix

    def _make_key(self, name: str, namespace: str) -> str:
        return f"{self.prefix}{namespace.upper()}_{name.upper()}"

    def get(self, name: str, namespace: str = "default") -> Optional[str]:
        key = self._make_key(name, namespace)
        return os.environ.get(key)

    def set(self, name: str, value: str, namespace: str = "default") -> bool:
        key = self._make_key(name, namespace)
        os.environ[key] = value
        return True

    def delete(self, name: str, namespace: str = "default") -> bool:
        key = self._make_key(name, namespace)
        if key in os.environ:
            del os.environ[key]
            return True
        return False

    def list(self, namespace: str = "default") -> List[str]:
        prefix = f"{self.prefix}{namespace.upper()}_"
        names = []
        for key in os.environ:
            if key.startswith(prefix):
                names.append(key[len(prefix):].lower())
        return names


class FileBackend(SecretBackend):
    """Backend that uses encrypted file storage."""

    def __init__(self, secrets_dir: Path, encryptor: 'Encryptor' = None):
        self.secrets_dir = secrets_dir
        self.secrets_dir.mkdir(parents=True, exist_ok=True)
        self.encryptor = encryptor
        self._cache: Dict[str, Dict[str, str]] = {}  # namespace -> {name: value}
        self._lock = threading.RLock()
        self._load_all()

    def _get_file(self, namespace: str) -> Path:
        return self.secrets_dir / f"{namespace}.secrets.json"

    def _load_all(self):
        """Load all secrets from files."""
        for file in self.secrets_dir.glob("*.secrets.json"):
            namespace = file.stem.replace(".secrets", "")
            try:
                with open(file, 'r') as f:
                    data = json.load(f)
                self._cache[namespace] = data
            except Exception:
                self._cache[namespace] = {}

    def _save(self, namespace: str):
        """Save namespace to file."""
        file = self._get_file(namespace)
        with open(file, 'w') as f:
            json.dump(self._cache.get(namespace, {}), f, indent=2)

    def get(self, name: str, namespace: str = "default") -> Optional[str]:
        with self._lock:
            ns_data = self._cache.get(namespace, {})
            return ns_data.get(name)

    def set(self, name: str, value: str, namespace: str = "default") -> bool:
        with self._lock:
            if namespace not in self._cache:
                self._cache[namespace] = {}
            self._cache[namespace][name] = value
            self._save(namespace)
            return True

    def delete(self, name: str, namespace: str = "default") -> bool:
        with self._lock:
            if namespace in self._cache and name in self._cache[namespace]:
                del self._cache[namespace][name]
                self._save(namespace)
                return True
            return False

    def list(self, namespace: str = "default") -> List[str]:
        with self._lock:
            return list(self._cache.get(namespace, {}).keys())


class MemoryBackend(SecretBackend):
    """In-memory backend (for testing)."""

    def __init__(self):
        self._secrets: Dict[str, Dict[str, str]] = {}
        self._lock = threading.RLock()

    def get(self, name: str, namespace: str = "default") -> Optional[str]:
        with self._lock:
            return self._secrets.get(namespace, {}).get(name)

    def set(self, name: str, value: str, namespace: str = "default") -> bool:
        with self._lock:
            if namespace not in self._secrets:
                self._secrets[namespace] = {}
            self._secrets[namespace][name] = value
            return True

    def delete(self, name: str, namespace: str = "default") -> bool:
        with self._lock:
            if namespace in self._secrets and name in self._secrets[namespace]:
                del self._secrets[namespace][name]
                return True
            return False

    def list(self, namespace: str = "default") -> List[str]:
        with self._lock:
            return list(self._secrets.get(namespace, {}).keys())


class Encryptor:
    """Handles encryption/decryption of secrets."""

    def __init__(self, key: Union[str, bytes] = None):
        self._fernet: Optional['Fernet'] = None

        if key and CRYPTO_AVAILABLE:
            if isinstance(key, str):
                # Derive key from password
                key_bytes = self._derive_key(key.encode())
            else:
                key_bytes = key

            # Ensure key is valid Fernet key (32 bytes, base64 encoded)
            if len(key_bytes) != 32:
                key_bytes = hashlib.sha256(key_bytes).digest()

            fernet_key = base64.urlsafe_b64encode(key_bytes)
            self._fernet = Fernet(fernet_key)

    def _derive_key(self, password: bytes, salt: bytes = None) -> bytes:
        """Derive encryption key from password."""
        if not CRYPTO_AVAILABLE:
            return hashlib.sha256(password).digest()

        salt = salt or b'genesis_salt_v1'  # Static salt for simplicity
        kdf = PBKDF2HMAC(
            algorithm=hashes.SHA256(),
            length=32,
            salt=salt,
            iterations=100000
        )
        return kdf.derive(password)

    def encrypt(self, data: str) -> str:
        """Encrypt string data."""
        if self._fernet:
            encrypted = self._fernet.encrypt(data.encode())
            return base64.urlsafe_b64encode(encrypted).decode()
        else:
            # Basic obfuscation (NOT secure, just for testing)
            return base64.b64encode(data.encode()).decode()

    def decrypt(self, data: str) -> str:
        """Decrypt string data."""
        if self._fernet:
            encrypted = base64.urlsafe_b64decode(data.encode())
            return self._fernet.decrypt(encrypted).decode()
        else:
            # Basic de-obfuscation
            return base64.b64decode(data.encode()).decode()

    @staticmethod
    def generate_key() -> bytes:
        """Generate a new encryption key."""
        if CRYPTO_AVAILABLE:
            return Fernet.generate_key()
        else:
            return py_secrets.token_bytes(32)


class SecretsManager:
    """
    Central secrets management system.
    """

    def __init__(
        self,
        backend: SecretBackend = None,
        encryption_key: Union[str, bytes] = None,
        secrets_dir: Path = None,
        enable_audit: bool = True
    ):
        # Initialize encryptor
        self.encryptor = Encryptor(encryption_key) if encryption_key else None

        # Initialize backend
        if backend:
            self.backend = backend
        elif secrets_dir:
            self.backend = FileBackend(secrets_dir, self.encryptor)
        else:
            # Default: environment + memory fallback
            self.backend = MemoryBackend()

        self._metadata: Dict[str, SecretMetadata] = {}
        self._audit_log: List[AuditEntry] = []
        self._enable_audit = enable_audit
        self._lock = threading.RLock()

        # Rotation callbacks
        self._rotation_callbacks: Dict[str, Callable[[str], str]] = {}

    def get(
        self,
        name: str,
        namespace: str = "default",
        default: str = None,
        caller: str = None
    ) -> Optional[str]:
        """Get a secret value."""
        with self._lock:
            try:
                value = self.backend.get(name, namespace)

                if value is None:
                    self._audit("get", name, namespace, False, caller, "not found")
                    return default

                # Check expiration
                meta_key = f"{namespace}:{name}"
                if meta_key in self._metadata:
                    meta = self._metadata[meta_key]
                    if meta.expires_at:
                        if datetime.now() > datetime.fromisoformat(meta.expires_at):
                            self._audit("get", name, namespace, False, caller, "expired")
                            return default

                    # Decrypt if needed
                    if meta.encrypted and self.encryptor:
                        value = self.encryptor.decrypt(value)

                self._audit("get", name, namespace, True, caller)
                return value

            except Exception as e:
                self._audit("get", name, namespace, False, caller, str(e))
                return default

    def set(
        self,
        name: str,
        value: str,
        namespace: str = "default",
        secret_type: SecretType = SecretType.STRING,
        encrypt: bool = False,
        ttl_seconds: int = None,
        description: str = "",
        tags: List[str] = None,
        caller: str = None
    ) -> bool:
        """Set a secret value."""
        with self._lock:
            try:
                meta_key = f"{namespace}:{name}"
                now = datetime.now().isoformat()

                # Get or create metadata
                if meta_key in self._metadata:
                    meta = self._metadata[meta_key]
                    meta.version += 1
                    meta.updated_at = now
                else:
                    meta = SecretMetadata(
                        name=name,
                        secret_type=secret_type,
                        created_at=now,
                        updated_at=now,
                        namespace=namespace,
                        description=description,
                        tags=tags or []
                    )

                # Set expiration
                if ttl_seconds:
                    meta.expires_at = (datetime.now() + timedelta(seconds=ttl_seconds)).isoformat()

                # Encrypt if requested
                store_value = value
                if encrypt and self.encryptor:
                    store_value = self.encryptor.encrypt(value)
                    meta.encrypted = True
                else:
                    meta.encrypted = False

                # Store
                success = self.backend.set(name, store_value, namespace)
                if success:
                    self._metadata[meta_key] = meta

                self._audit("set", name, namespace, success, caller)
                return success

            except Exception as e:
                self._audit("set", name, namespace, False, caller, str(e))
                return False

    def delete(
        self,
        name: str,
        namespace: str = "default",
        caller: str = None
    ) -> bool:
        """Delete a secret."""
        with self._lock:
            try:
                success = self.backend.delete(name, namespace)

                meta_key = f"{namespace}:{name}"
                if meta_key in self._metadata:
                    del self._metadata[meta_key]

                self._audit("delete", name, namespace, success, caller)
                return success

            except Exception as e:
                self._audit("delete", name, namespace, False, caller, str(e))
                return False

    def list(self, namespace: str = "default") -> List[str]:
        """List secret names in namespace."""
        return self.backend.list(namespace)

    def get_metadata(self, name: str, namespace: str = "default") -> Optional[SecretMetadata]:
        """Get secret metadata."""
        meta_key = f"{namespace}:{name}"
        return self._metadata.get(meta_key)

    def exists(self, name: str, namespace: str = "default") -> bool:
        """Check if a secret exists."""
        return self.backend.get(name, namespace) is not None

    def rotate(
        self,
        name: str,
        namespace: str = "default",
        new_value: str = None,
        caller: str = None
    ) -> bool:
        """Rotate a secret to a new value."""
        with self._lock:
            try:
                meta_key = f"{namespace}:{name}"
                if meta_key not in self._metadata:
                    self._audit("rotate", name, namespace, False, caller, "not found")
                    return False

                meta = self._metadata[meta_key]

                # Get new value
                if new_value is None:
                    if name in self._rotation_callbacks:
                        new_value = self._rotation_callbacks[name](name)
                    else:
                        # Generate random token
                        new_value = py_secrets.token_urlsafe(32)

                # Store with same encryption settings
                store_value = new_value
                if meta.encrypted and self.encryptor:
                    store_value = self.encryptor.encrypt(new_value)

                success = self.backend.set(name, store_value, namespace)
                if success:
                    meta.version += 1
                    meta.updated_at = datetime.now().isoformat()

                self._audit("rotate", name, namespace, success, caller)
                return success

            except Exception as e:
                self._audit("rotate", name, namespace, False, caller, str(e))
                return False

    def register_rotation_callback(self, name: str, callback: Callable[[str], str]):
        """Register a callback for generating new secret values during rotation."""
        self._rotation_callbacks[name] = callback

    def check_expired(self, namespace: str = "default") -> List[str]:
        """Check for expired secrets."""
        expired = []
        now = datetime.now()

        for meta_key, meta in self._metadata.items():
            if not meta_key.startswith(f"{namespace}:"):
                continue
            if meta.expires_at:
                if now > datetime.fromisoformat(meta.expires_at):
                    expired.append(meta.name)

        return expired

    def _audit(
        self,
        action: str,
        name: str,
        namespace: str,
        success: bool,
        caller: str = None,
        error: str = None
    ):
        """Log an audit entry."""
        if not self._enable_audit:
            return

        entry = AuditEntry(
            timestamp=datetime.now().isoformat(),
            action=action,
            secret_name=name,
            namespace=namespace,
            success=success,
            caller=caller,
            error=error
        )
        self._audit_log.append(entry)

        # Keep last 1000 entries
        if len(self._audit_log) > 1000:
            self._audit_log = self._audit_log[-1000:]

    def get_audit_log(
        self,
        limit: int = 100,
        namespace: str = None,
        action: str = None
    ) -> List[AuditEntry]:
        """Get audit log entries."""
        log = self._audit_log.copy()

        if namespace:
            log = [e for e in log if e.namespace == namespace]
        if action:
            log = [e for e in log if e.action == action]

        return log[-limit:]

    def get_status(self) -> Dict:
        """Get secrets manager status."""
        namespaces = set()
        for meta_key in self._metadata:
            ns = meta_key.split(":")[0]
            namespaces.add(ns)

        return {
            "encryption_available": CRYPTO_AVAILABLE,
            "encryption_enabled": self.encryptor is not None,
            "backend_type": type(self.backend).__name__,
            "secret_count": len(self._metadata),
            "namespaces": list(namespaces),
            "audit_entries": len(self._audit_log)
        }


# Convenience functions
_manager: Optional[SecretsManager] = None


def get_secrets_manager() -> SecretsManager:
    """Get global secrets manager."""
    global _manager
    if _manager is None:
        # Try to use env var for encryption key
        key = os.environ.get("GENESIS_SECRETS_KEY")
        secrets_dir = Path(os.environ.get("GENESIS_SECRETS_DIR", ".secrets"))
        _manager = SecretsManager(encryption_key=key, secrets_dir=secrets_dir)
    return _manager


def get_secret(name: str, default: str = None, namespace: str = "default") -> Optional[str]:
    """Get a secret from the global manager."""
    return get_secrets_manager().get(name, namespace, default)


def set_secret(
    name: str,
    value: str,
    encrypt: bool = True,
    namespace: str = "default"
) -> bool:
    """Set a secret in the global manager."""
    return get_secrets_manager().set(name, value, namespace=namespace, encrypt=encrypt)


def main():
    """CLI and demo for secrets manager."""
    import argparse
    parser = argparse.ArgumentParser(description="Genesis Secrets Manager")
    parser.add_argument("command", choices=["demo", "status", "list", "get", "set"])
    parser.add_argument("--name", help="Secret name")
    parser.add_argument("--value", help="Secret value")
    parser.add_argument("--namespace", default="default", help="Namespace")
    args = parser.parse_args()

    if args.command == "demo":
        print("Secrets Manager Demo")
        print("=" * 40)

        # Create manager with encryption
        print(f"\nCrypto library available: {CRYPTO_AVAILABLE}")

        # Memory backend demo
        print("\n1. Basic secrets (memory backend):")
        manager = SecretsManager(backend=MemoryBackend())

        manager.set("api_key", "sk-12345")
        manager.set("db_password", "secret123", encrypt=True)

        print(f"  api_key: {manager.get('api_key')}")
        print(f"  db_password: {manager.get('db_password')}")
        print(f"  missing: {manager.get('missing', 'DEFAULT')}")

        # Metadata
        print("\n2. Secret metadata:")
        manager.set(
            "oauth_token",
            "token_value",
            secret_type=SecretType.API_KEY,
            ttl_seconds=3600,
            description="OAuth access token",
            tags=["auth", "oauth"]
        )

        meta = manager.get_metadata("oauth_token")
        print(f"  {json.dumps(meta.to_dict(), indent=4)}")

        # Namespaces
        print("\n3. Namespaces:")
        manager.set("api_key", "prod_key", namespace="production")
        manager.set("api_key", "dev_key", namespace="development")

        print(f"  default: {manager.get('api_key')}")
        print(f"  production: {manager.get('api_key', namespace='production')}")
        print(f"  development: {manager.get('api_key', namespace='development')}")

        # Rotation
        print("\n4. Secret rotation:")
        manager.set("rotating_key", "original_value")
        print(f"  Before: {manager.get('rotating_key')}")
        manager.rotate("rotating_key")
        print(f"  After rotation: {manager.get('rotating_key')[:20]}...")

        # Custom rotation callback
        def generate_api_key(name):
            return f"custom_key_{py_secrets.token_hex(8)}"

        manager.register_rotation_callback("custom_key", generate_api_key)
        manager.set("custom_key", "initial")
        manager.rotate("custom_key")
        print(f"  Custom rotation: {manager.get('custom_key')}")

        # Audit log
        print("\n5. Audit log:")
        for entry in manager.get_audit_log(5):
            status = "OK" if entry.success else "FAIL"
            print(f"  [{entry.timestamp[:19]}] {entry.action} {entry.secret_name} [{status}]")

        # Expiration
        print("\n6. TTL/Expiration:")
        manager.set("temp_token", "expires_soon", ttl_seconds=1)
        print(f"  Before expiry: {manager.get('temp_token')}")
        time.sleep(1.5)
        print(f"  After expiry: {manager.get('temp_token', 'EXPIRED')}")

        expired = manager.check_expired()
        print(f"  Expired secrets: {expired}")

        # Encryption demo (if available)
        if CRYPTO_AVAILABLE:
            print("\n7. Encryption (with cryptography):")
            encrypted_manager = SecretsManager(
                backend=MemoryBackend(),
                encryption_key="my_secret_password"
            )
            encrypted_manager.set("encrypted_secret", "very_sensitive", encrypt=True)

            # Get raw value from backend
            raw = encrypted_manager.backend.get("encrypted_secret")
            decrypted = encrypted_manager.get("encrypted_secret")

            print(f"  Raw (encrypted): {raw[:40]}...")
            print(f"  Decrypted: {decrypted}")
        else:
            print("\n7. Note: cryptography library not installed")
            print("   Install with: pip install cryptography")

        # Status
        print("\n8. Status:")
        print(f"  {json.dumps(manager.get_status(), indent=4)}")

    elif args.command == "status":
        manager = get_secrets_manager()
        print(json.dumps(manager.get_status(), indent=2))

    elif args.command == "list":
        manager = get_secrets_manager()
        secrets = manager.list(args.namespace)
        print(f"Secrets in '{args.namespace}':")
        for name in secrets:
            meta = manager.get_metadata(name, args.namespace)
            encrypted = "[encrypted]" if meta and meta.encrypted else ""
            print(f"  - {name} {encrypted}")

    elif args.command == "get":
        if not args.name:
            print("Error: --name required")
            return
        manager = get_secrets_manager()
        value = manager.get(args.name, args.namespace)
        if value:
            print(value)
        else:
            print(f"Secret '{args.name}' not found")

    elif args.command == "set":
        if not args.name or not args.value:
            print("Error: --name and --value required")
            return
        manager = get_secrets_manager()
        if manager.set(args.name, args.value, namespace=args.namespace, encrypt=True):
            print(f"Secret '{args.name}' set successfully")
        else:
            print(f"Failed to set secret '{args.name}'")


if __name__ == "__main__":
    main()