namespace.py

"""Namespace setup code generation for SubprocessExecutor.

Generates Python code that sets up tools, workflows, artifacts, and deps namespaces
in the kernel subprocess using full py-code-mode functionality.
"""

from __future__ import annotations

from py_code_mode.execution.protocol import FileStorageAccess, RedisStorageAccess, StorageAccess


def build_namespace_setup_code(
    storage_access: StorageAccess | None,
    allow_runtime_deps: bool = True,
) -> str:
    """Generate Python code that sets up namespaces in the kernel.

    The generated code imports from py-code-mode (which must be installed
    in the kernel's venv) and creates real namespace objects with full
    functionality including tool invocation, workflow creation, and semantic search.

    Args:
        storage_access: Storage access descriptor with paths or connection info.
                       If None, returns empty string.
        allow_runtime_deps: Whether to allow deps.add() and deps.sync() calls.
                           If False, these methods raise RuntimeError.

    Returns:
        Python code string to execute in the kernel to set up namespaces.
    """
    if storage_access is None:
        return ""

    if isinstance(storage_access, FileStorageAccess):
        return _build_file_storage_setup_code(storage_access, allow_runtime_deps)

    if isinstance(storage_access, RedisStorageAccess):
        return _build_redis_storage_setup_code(storage_access, allow_runtime_deps)

    # Unknown storage types not supported
    return ""


def _build_vector_store_setup_code(vectors_path_str: str) -> str:
    """Generate vector store setup code if vectors_path is provided.

    Args:
        vectors_path_str: String representation of vectors path or "None".

    Returns:
        Code string for vector store setup, or empty string if not needed.
    """
    if vectors_path_str == "None":
        return """
_vector_store = None"""

    return f"""
_vectors_path = Path({vectors_path_str})

# Setup vector store if vectors_path provided
_vector_store = None
try:
    from py_code_mode.workflows.vector_stores.chroma import ChromaVectorStore
    from py_code_mode.workflows import Embedder
    _vectors_path.mkdir(parents=True, exist_ok=True)
    _embedder = Embedder()
    _vector_store = ChromaVectorStore(path=_vectors_path, embedder=_embedder)
except ImportError:
    _vector_store = None"""


def _build_vector_store_cleanup_code(has_vectors_path: bool) -> str:
    """Generate cleanup code for vector store variables.

    Args:
        has_vectors_path: Whether vectors_path was provided.

    Returns:
        Code string to cleanup vector store-related variables.
    """
    if not has_vectors_path:
        return "del _vector_store"

    return """del _vectors_path, _vector_store
try:
    del ChromaVectorStore, Embedder, _embedder
except NameError:
    pass"""


def _build_file_storage_setup_code(
    storage_access: FileStorageAccess,
    allow_runtime_deps: bool,
) -> str:
    """Generate namespace setup code for FileStorageAccess.

    Note: Tools are owned by executors (via config.tools_path), not storage.
    The generated code creates an empty ToolRegistry since tools loading is
    handled separately by the executor.
    """
    workflows_path_str = (
        repr(str(storage_access.workflows_path)) if storage_access.workflows_path else "None"
    )
    artifacts_path_str = repr(str(storage_access.artifacts_path))
    base_path = storage_access.root_path or storage_access.artifacts_path.parent
    base_path_str = repr(str(base_path))
    allow_deps_str = "True" if allow_runtime_deps else "False"
    vectors_path_str = (
        repr(str(storage_access.vectors_path)) if storage_access.vectors_path else "None"
    )
    has_vectors_path = storage_access.vectors_path is not None

    # Generate vector store setup code conditionally
    vector_store_setup = _build_vector_store_setup_code(vectors_path_str)
    vector_store_cleanup = _build_vector_store_cleanup_code(has_vectors_path)

    return f'''# Auto-generated namespace setup for SubprocessExecutor
# This code sets up full py-code-mode namespaces in the kernel

from pathlib import Path
import asyncio
import nest_asyncio

# Enable nested event loops (required for sync tool calls in Jupyter kernel)
nest_asyncio.apply()

# =============================================================================
# Tools Namespace (with sync wrapper for subprocess context)
# =============================================================================
# NOTE: Tools are owned by executors (via config.tools_path), not storage.
# The subprocess executor handles tool loading separately if tools_path is configured.
# Here we create an empty registry as a placeholder.

from py_code_mode.tools import ToolRegistry, ToolsNamespace
from py_code_mode.tools.adapters import CLIAdapter

_registry = ToolRegistry()

# Create the base namespace
_base_tools = ToolsNamespace(_registry)

# Wrapper that forces sync execution in Jupyter kernel context
class _SyncToolsWrapper:
    """Wrapper that ensures tools execute synchronously in subprocess."""

    def __init__(self, namespace):
        self._namespace = namespace

    def __getattr__(self, name):
        attr = getattr(self._namespace, name)
        if hasattr(attr, '_tool'):
            # It's a ToolProxy - wrap it
            return _SyncToolProxy(attr)
        return attr

    def list(self):
        return self._namespace.list()

    def search(self, query, limit=5):
        return self._namespace.search(query, limit)


class _SyncToolProxy:
    """Wrapper that forces sync execution for tool proxies."""

    def __init__(self, proxy):
        self._proxy = proxy

    def __call__(self, **kwargs):
        result = self._proxy(**kwargs)
        if asyncio.iscoroutine(result):
            return asyncio.get_event_loop().run_until_complete(result)
        return result

    def __getattr__(self, name):
        attr = getattr(self._proxy, name)
        if callable(attr):
            return _SyncCallableWrapper(attr)
        return attr

    def list(self):
        return self._proxy.list()


class _SyncCallableWrapper:
    """Wrapper that forces sync execution for callable proxies."""

    def __init__(self, callable_proxy):
        self._callable = callable_proxy

    def __call__(self, **kwargs):
        result = self._callable(**kwargs)
        if asyncio.iscoroutine(result):
            return asyncio.get_event_loop().run_until_complete(result)
        return result


tools = _SyncToolsWrapper(_base_tools)

# =============================================================================
# Workflows Namespace (with optional vector store)
# =============================================================================

from py_code_mode.workflows import FileWorkflowStore, create_workflow_library
from py_code_mode.execution.in_process.workflows_namespace import WorkflowsNamespace

_workflows_path = Path({workflows_path_str}) if {workflows_path_str} else None
{vector_store_setup}

if _workflows_path is not None:
    _workflows_path.mkdir(parents=True, exist_ok=True)
    _store = FileWorkflowStore(_workflows_path)
    _library = create_workflow_library(store=_store, vector_store=_vector_store)
else:
    from py_code_mode.workflows import MemoryWorkflowStore, MockEmbedder, WorkflowLibrary
    _store = MemoryWorkflowStore()
    _library = WorkflowLibrary(embedder=MockEmbedder(), store=_store, vector_store=_vector_store)

# WorkflowsNamespace now takes a namespace dict directly (no executor needed).
# Create the namespace dict first, then wire up circular references.
_workflows_ns_dict = {{}}
workflows = WorkflowsNamespace(_library, _workflows_ns_dict)

# Wire up the namespace so workflows can access tools/workflows/artifacts
_workflows_ns_dict["tools"] = tools
_workflows_ns_dict["workflows"] = workflows

# =============================================================================
# Artifacts Namespace (with simplified API for agent usage)
# =============================================================================

from py_code_mode.artifacts import FileArtifactStore

_artifacts_path = Path({artifacts_path_str})
_artifacts_path.mkdir(parents=True, exist_ok=True)
_base_artifacts = FileArtifactStore(_artifacts_path)


class _SimpleArtifactStore:
    """Wrapper providing simplified artifacts API for agents.

    Wraps FileArtifactStore to provide:
    - save(name, data) with optional description (defaults to empty string)
    - All other methods pass through unchanged
    """

    def __init__(self, store):
        self._store = store

    def save(self, name, data, description=""):
        """Save artifact with optional description."""
        return self._store.save(name, data, description)

    def load(self, name):
        """Load artifact by name."""
        return self._store.load(name)

    def list(self):
        """List all artifacts."""
        return self._store.list()

    def exists(self, name):
        """Check if artifact exists."""
        return self._store.exists(name)

    def delete(self, name):
        """Delete artifact."""
        return self._store.delete(name)

    def get(self, name):
        """Get artifact metadata."""
        return self._store.get(name)

    @property
    def path(self):
        """Base path for raw file access."""
        return self._store.path


artifacts = _SimpleArtifactStore(_base_artifacts)

# Complete the namespace wiring for workflows
_workflows_ns_dict["artifacts"] = artifacts


# =============================================================================
# Deps Namespace (with optional runtime deps control)
# =============================================================================

from py_code_mode.deps import DepsNamespace, FileDepsStore, PackageInstaller

_base_path = Path({base_path_str})
_deps_store = FileDepsStore(_base_path)
_installer = PackageInstaller()
_base_deps = DepsNamespace(_deps_store, _installer)

_allow_runtime_deps = {allow_deps_str}


class _RuntimeDepsDisabledError(RuntimeError):
    """Raised when runtime deps are disabled and a blocked operation is attempted."""
    pass


class _ControlledDepsNamespace:
    """Wrapper that optionally blocks add() and remove() calls.

    When allow_runtime_deps=False, add() and remove() raise error
    to prevent runtime package modification. list() and sync() always work.
    sync() is allowed because it only installs pre-configured packages.

    Security: Access to internal attributes is blocked via __getattribute__
    to prevent bypass attacks like deps._namespace.add().
    """

    _ALLOWED_ATTRS = frozenset({{
        "add", "list", "remove", "sync", "__repr__", "__class__", "__doc__"
    }})

    def __init__(self, namespace, allow_runtime):
        # Use object.__setattr__ to bypass __getattribute__
        object.__setattr__(self, "_namespace", namespace)
        object.__setattr__(self, "_allow_runtime", allow_runtime)

    def __getattribute__(self, name):
        """Control access to attributes - block internal attrs to prevent bypass."""
        allowed = object.__getattribute__(self, "_ALLOWED_ATTRS")
        if name in allowed:
            return object.__getattribute__(self, name)
        if name.startswith("_"):
            raise AttributeError(
                f"Cannot access internal attribute '{{name}}'. Runtime deps are disabled."
            )
        return object.__getattribute__(self, name)

    def add(self, package):
        """Add a package (blocked if runtime deps disabled)."""
        allow_runtime = object.__getattribute__(self, "_allow_runtime")
        if not allow_runtime:
            raise _RuntimeDepsDisabledError(
                "RuntimeDepsDisabledError: Runtime dependency installation is disabled. "
                "Dependencies must be pre-configured before session start."
            )
        namespace = object.__getattribute__(self, "_namespace")
        return namespace.add(package)

    def sync(self):
        """Sync packages (always allowed - only installs pre-configured deps)."""
        namespace = object.__getattribute__(self, "_namespace")
        return namespace.sync()

    def list(self):
        """List packages (always allowed)."""
        namespace = object.__getattribute__(self, "_namespace")
        return namespace.list()

    def remove(self, package):
        """Remove a package from config (blocked if runtime deps disabled)."""
        allow_runtime = object.__getattribute__(self, "_allow_runtime")
        if not allow_runtime:
            raise _RuntimeDepsDisabledError(
                "RuntimeDepsDisabledError: Runtime dependency modification is disabled. "
                "Dependencies must be pre-configured before session start."
            )
        namespace = object.__getattribute__(self, "_namespace")
        return namespace.remove(package)

    def __repr__(self):
        allow_runtime = object.__getattribute__(self, "_allow_runtime")
        status = "enabled" if allow_runtime else "disabled"
        return f"<ControlledDepsNamespace: runtime={{status}}>"


deps = _ControlledDepsNamespace(_base_deps, _allow_runtime_deps)

# Complete the namespace wiring for workflows to include deps
_workflows_ns_dict["deps"] = deps

# =============================================================================
# Cleanup temporary variables (keep wrapper classes for runtime use)
# =============================================================================

del _registry, _base_tools
del _workflows_path, _store, _library, _workflows_ns_dict
del _artifacts_path, _base_artifacts
{vector_store_cleanup}
del _base_path, _deps_store, _installer, _base_deps, _allow_runtime_deps
del Path
del ToolRegistry, ToolsNamespace, CLIAdapter
del FileWorkflowStore, create_workflow_library, WorkflowsNamespace
try:
    del MemoryWorkflowStore, MockEmbedder, WorkflowLibrary
except NameError:
    pass
del FileArtifactStore
del DepsNamespace, FileDepsStore, PackageInstaller

# Note: Wrapper classes (_SyncToolsWrapper, _SyncToolProxy, _SyncCallableWrapper,
# _SimpleArtifactStore, _ControlledDepsNamespace) and asyncio/nest_asyncio are kept for runtime use
'''


def _build_redis_storage_setup_code(
    storage_access: RedisStorageAccess,
    allow_runtime_deps: bool,
) -> str:
    """Generate namespace setup code for RedisStorageAccess.

    Note: Tools are owned by executors (via config.tools_path), not storage.
    The generated code creates an empty ToolRegistry since tools loading is
    handled separately by the executor.
    """
    redis_url_str = repr(storage_access.redis_url)
    workflows_prefix_str = repr(storage_access.workflows_prefix)
    artifacts_prefix_str = repr(storage_access.artifacts_prefix)
    vectors_prefix_str = (
        repr(storage_access.vectors_prefix) if storage_access.vectors_prefix else "None"
    )
    deps_prefix = storage_access.root_prefix or storage_access.workflows_prefix.rsplit(":", 1)[0]
    deps_prefix_str = repr(deps_prefix)
    allow_deps_str = "True" if allow_runtime_deps else "False"

    return f'''# Auto-generated namespace setup for SubprocessExecutor (Redis)
# This code sets up full py-code-mode namespaces in the kernel

import asyncio
import nest_asyncio

# Enable nested event loops (required for sync tool calls in Jupyter kernel)
nest_asyncio.apply()

from redis import Redis

_redis_client = Redis.from_url({redis_url_str}, decode_responses=False)

# =============================================================================
# Tools Namespace (with sync wrapper for subprocess context)
# =============================================================================
# NOTE: Tools are owned by executors (via config.tools_path), not storage.
# The subprocess executor handles tool loading separately if tools_path is configured.
# Here we create an empty registry as a placeholder.

from py_code_mode.tools import ToolRegistry, ToolsNamespace
from py_code_mode.tools.adapters import CLIAdapter

_registry = ToolRegistry()

# Create the base namespace
_base_tools = ToolsNamespace(_registry)

# Wrapper that forces sync execution in Jupyter kernel context
class _SyncToolsWrapper:
    """Wrapper that ensures tools execute synchronously in subprocess."""

    def __init__(self, namespace):
        self._namespace = namespace

    def __getattr__(self, name):
        attr = getattr(self._namespace, name)
        if hasattr(attr, '_tool'):
            # It's a ToolProxy - wrap it
            return _SyncToolProxy(attr)
        return attr

    def list(self):
        return self._namespace.list()

    def search(self, query, limit=5):
        return self._namespace.search(query, limit)


class _SyncToolProxy:
    """Wrapper that forces sync execution for tool proxies."""

    def __init__(self, proxy):
        self._proxy = proxy

    def __call__(self, **kwargs):
        result = self._proxy(**kwargs)
        if asyncio.iscoroutine(result):
            return asyncio.get_event_loop().run_until_complete(result)
        return result

    def __getattr__(self, name):
        attr = getattr(self._proxy, name)
        if callable(attr):
            return _SyncCallableWrapper(attr)
        return attr

    def list(self):
        return self._proxy.list()


class _SyncCallableWrapper:
    """Wrapper that forces sync execution for callable proxies."""

    def __init__(self, callable_proxy):
        self._callable = callable_proxy

    def __call__(self, **kwargs):
        result = self._callable(**kwargs)
        if asyncio.iscoroutine(result):
            return asyncio.get_event_loop().run_until_complete(result)
        return result


tools = _SyncToolsWrapper(_base_tools)

# =============================================================================
# Workflows Namespace (with optional vector store)
# =============================================================================

from py_code_mode.workflows import RedisWorkflowStore, create_workflow_library
from py_code_mode.execution.in_process.workflows_namespace import WorkflowsNamespace

_workflows_prefix = {workflows_prefix_str}

_vector_store = None
if {vectors_prefix_str} is not None:
    try:
        from py_code_mode.workflows.vector_stores.redis_store import RedisVectorStore
        from py_code_mode.workflows import Embedder
        _embedder = Embedder()
        _vector_store = RedisVectorStore(
            redis=_redis_client,
            embedder=_embedder,
            prefix={vectors_prefix_str},
        )
    except ImportError:
        _vector_store = None

_store = RedisWorkflowStore(_redis_client, prefix=_workflows_prefix)
_library = create_workflow_library(store=_store, vector_store=_vector_store)

# WorkflowsNamespace now takes a namespace dict directly (no executor needed).
# Create the namespace dict first, then wire up circular references.
_workflows_ns_dict = {{}}
workflows = WorkflowsNamespace(_library, _workflows_ns_dict)

# Wire up the namespace so workflows can access tools/workflows/artifacts
_workflows_ns_dict["tools"] = tools
_workflows_ns_dict["workflows"] = workflows

# =============================================================================
# Artifacts Namespace (with simplified API for agent usage)
# =============================================================================

from py_code_mode.artifacts import RedisArtifactStore

_artifacts_prefix = {artifacts_prefix_str}
_base_artifacts = RedisArtifactStore(_redis_client, prefix=_artifacts_prefix)


class _SimpleArtifactStore:
    """Wrapper providing simplified artifacts API for agents.

    Wraps RedisArtifactStore to provide:
    - save(name, data) with optional description (defaults to empty string)
    - All other methods pass through unchanged
    """

    def __init__(self, store):
        self._store = store

    def save(self, name, data, description=""):
        """Save artifact with optional description."""
        return self._store.save(name, data, description)

    def load(self, name):
        """Load artifact by name."""
        return self._store.load(name)

    def list(self):
        """List all artifacts."""
        return self._store.list()

    def exists(self, name):
        """Check if artifact exists."""
        return self._store.exists(name)

    def delete(self, name):
        """Delete artifact."""
        return self._store.delete(name)

    def get(self, name):
        """Get artifact metadata."""
        return self._store.get(name)


artifacts = _SimpleArtifactStore(_base_artifacts)

# Complete the namespace wiring for workflows
_workflows_ns_dict["artifacts"] = artifacts

# =============================================================================
# Deps Namespace (with optional runtime deps control)
# =============================================================================

from py_code_mode.deps import DepsNamespace, RedisDepsStore, PackageInstaller

_deps_prefix = {deps_prefix_str}
_deps_store = RedisDepsStore(_redis_client, prefix=_deps_prefix)
_installer = PackageInstaller()
_base_deps = DepsNamespace(_deps_store, _installer)

_allow_runtime_deps = {allow_deps_str}


class _RuntimeDepsDisabledError(RuntimeError):
    """Raised when runtime deps are disabled and a blocked operation is attempted."""
    pass


class _ControlledDepsNamespace:
    """Wrapper that optionally blocks add() and remove() calls.

    When allow_runtime_deps=False, add() and remove() raise error
    to prevent runtime package modification. list() and sync() always work.
    sync() is allowed because it only installs pre-configured packages.

    Security: Access to internal attributes is blocked via __getattribute__
    to prevent bypass attacks like deps._namespace.add().
    """

    _ALLOWED_ATTRS = frozenset({{
        "add", "list", "remove", "sync", "__repr__", "__class__", "__doc__"
    }})

    def __init__(self, namespace, allow_runtime):
        # Use object.__setattr__ to bypass __getattribute__
        object.__setattr__(self, "_namespace", namespace)
        object.__setattr__(self, "_allow_runtime", allow_runtime)

    def __getattribute__(self, name):
        """Control access to attributes - block internal attrs to prevent bypass."""
        allowed = object.__getattribute__(self, "_ALLOWED_ATTRS")
        if name in allowed:
            return object.__getattribute__(self, name)
        if name.startswith("_"):
            raise AttributeError(
                f"Cannot access internal attribute '{{name}}'. Runtime deps are disabled."
            )
        return object.__getattribute__(self, name)

    def add(self, package):
        """Add a package (blocked if runtime deps disabled)."""
        allow_runtime = object.__getattribute__(self, "_allow_runtime")
        if not allow_runtime:
            raise _RuntimeDepsDisabledError(
                "RuntimeDepsDisabledError: Runtime dependency installation is disabled. "
                "Dependencies must be pre-configured before session start."
            )
        namespace = object.__getattribute__(self, "_namespace")
        return namespace.add(package)

    def sync(self):
        """Sync packages (always allowed - only installs pre-configured deps)."""
        namespace = object.__getattribute__(self, "_namespace")
        return namespace.sync()

    def list(self):
        """List packages (always allowed)."""
        namespace = object.__getattribute__(self, "_namespace")
        return namespace.list()

    def remove(self, package):
        """Remove a package from config (blocked if runtime deps disabled)."""
        allow_runtime = object.__getattribute__(self, "_allow_runtime")
        if not allow_runtime:
            raise _RuntimeDepsDisabledError(
                "RuntimeDepsDisabledError: Runtime dependency modification is disabled. "
                "Dependencies must be pre-configured before session start."
            )
        namespace = object.__getattribute__(self, "_namespace")
        return namespace.remove(package)

    def __repr__(self):
        allow_runtime = object.__getattribute__(self, "_allow_runtime")
        status = "enabled" if allow_runtime else "disabled"
        return f"<ControlledDepsNamespace: runtime={{status}}>"


deps = _ControlledDepsNamespace(_base_deps, _allow_runtime_deps)

# Complete the namespace wiring for workflows to include deps
_workflows_ns_dict["deps"] = deps

# =============================================================================
# Cleanup temporary variables (keep wrapper classes for runtime use)
# =============================================================================

if {vectors_prefix_str} is not None:
    del _vector_store
    try:
        del RedisVectorStore, Embedder, _embedder
    except NameError:
        pass
else:
    del _vector_store

del _registry, _base_tools
del _workflows_prefix, _store, _library, _workflows_ns_dict
del _artifacts_prefix, _base_artifacts
del _deps_prefix, _deps_store, _installer, _base_deps, _allow_runtime_deps
del ToolRegistry, ToolsNamespace, CLIAdapter
del RedisWorkflowStore, create_workflow_library, WorkflowsNamespace
try:
    del MockEmbedder, WorkflowLibrary
except NameError:
    pass
del RedisArtifactStore
del DepsNamespace, RedisDepsStore, PackageInstaller
del Redis
# Note: Wrapper classes (_SyncToolsWrapper, _SyncToolProxy, _SyncCallableWrapper,
# _SimpleArtifactStore, _ControlledDepsNamespace), asyncio/nest_asyncio, and
# _redis_client are kept for runtime use
'''