snakebridge

"""
SnakeBridge Snakepit Adapter

Provides the Snakepit adapter interface for SnakeBridge.
This module is called by the Snakepit runtime to execute Python functions
and return results to Elixir.

Main function:
    snakebridge_call(module: str, function: str, args: dict) -> dict

The adapter:
1. Imports the specified Python module
2. Gets the specified function from the module
3. Decodes the arguments from SnakeBridge format
4. Calls the function with the decoded arguments
5. Encodes the result back to SnakeBridge format
6. Returns a success/error response

Graceful Serialization:
    When encoding results, this adapter preserves container structure (lists,
    dicts, tuples, sets) even when they contain non-serializable objects.
    Non-serializable leaf values are ref-wrapped in place. No lossy conversion
    is performed (e.g., no model_dump/to_dict calls); non-serializable leaves
    become refs that can be used for further Python operations. This enables
    users to access all serializable fields while still having refs for
    non-serializable objects.
"""

import sys
import importlib
import importlib.util
import inspect
import traceback
import uuid
import os
import glob
import hashlib
import time
import threading
import types
from contextlib import nullcontext
from typing import Any, Dict, List, Tuple, Optional

# Import the SnakeBridge type encoding system
try:
    from snakebridge_types import (
        decode,
        encode,
        encode_error,
        _is_generator_or_iterator,
        _get_stream_type,
        Atom,
        SCHEMA_VERSION,
    )
except ImportError:
    # If running as a script, try relative import
    import os
    sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
    from snakebridge_types import (
        decode,
        encode,
        encode_error,
        _is_generator_or_iterator,
        _get_stream_type,
        Atom,
        SCHEMA_VERSION,
    )

try:
    from snakepit_bridge import telemetry as snakepit_telemetry
except Exception:
    snakepit_telemetry = None


# Module cache to avoid repeated imports
_module_cache: Dict[str, Any] = {}
_instance_registry: Dict[str, Any] = {}
_helper_registry: Dict[str, Any] = {}
_helper_registry_key: Optional[Tuple[Any, ...]] = None
_helper_registry_index: List[Dict[str, Any]] = []

# Thread locks for global state
_module_cache_lock = threading.RLock()
_registry_lock = threading.RLock()
_helper_lock = threading.RLock()

PROTOCOL_VERSION = 1
MIN_SUPPORTED_VERSION = 1
REF_SCHEMA_VERSION = 1
DEFAULT_REF_TTL_SECONDS = 0.0
DEFAULT_REF_MAX_SIZE = 10000
ALLOW_LEGACY_PROTOCOL = os.getenv("SNAKEBRIDGE_ALLOW_LEGACY_PROTOCOL", "false").lower() in (
    "1",
    "true",
    "yes",
)


class SnakeBridgeHelperNotFoundError(Exception):
    pass


class SnakeBridgeHelperLoadError(Exception):
    pass


class SnakeBridgeSerializationError(Exception):
    pass


class SnakeBridgeProtocolError(Exception):
    def __init__(self, message: str, details: Optional[Dict[str, Any]] = None):
        super().__init__(message)
        self.details = details or {}


def _protocol_compatibility(arguments: Dict[str, Any]) -> None:
    protocol_version = arguments.get("protocol_version")
    min_supported_version = arguments.get("min_supported_version")

    if protocol_version is None and min_supported_version is None and ALLOW_LEGACY_PROTOCOL:
        protocol_version = MIN_SUPPORTED_VERSION
        min_supported_version = 0
    else:
        if protocol_version is None:
            protocol_version = 0
        if min_supported_version is None:
            min_supported_version = 0

    if protocol_version < MIN_SUPPORTED_VERSION or min_supported_version > PROTOCOL_VERSION:
        details = {
            "caller_protocol_version": protocol_version,
            "caller_min_supported_version": min_supported_version,
            "adapter_protocol_version": PROTOCOL_VERSION,
            "adapter_min_supported_version": MIN_SUPPORTED_VERSION,
        }
        message = (
            "SnakeBridge protocol version mismatch "
            f"(caller_protocol_version={protocol_version}, "
            f"caller_min_supported_version={min_supported_version}, "
            f"adapter_protocol_version={PROTOCOL_VERSION}, "
            f"adapter_min_supported_version={MIN_SUPPORTED_VERSION}). "
            "Ensure Elixir and Python SnakeBridge versions are compatible."
        )
        raise SnakeBridgeProtocolError(
            message,
            details=details,
        )


def _registry_limits() -> Tuple[float, int]:
    ttl_env = os.getenv("SNAKEBRIDGE_REF_TTL_SECONDS")
    max_env = os.getenv("SNAKEBRIDGE_REF_MAX")

    try:
        ttl = float(ttl_env) if ttl_env is not None else DEFAULT_REF_TTL_SECONDS
    except ValueError:
        ttl = DEFAULT_REF_TTL_SECONDS

    try:
        max_size = int(max_env) if max_env is not None else DEFAULT_REF_MAX_SIZE
    except ValueError:
        max_size = DEFAULT_REF_MAX_SIZE

    return ttl, max_size


def _entry_last_access(entry: Any) -> float:
    if isinstance(entry, dict):
        return float(entry.get("last_access") or entry.get("created_at") or 0.0)
    return 0.0


def _touch_entry(entry: Any) -> None:
    if isinstance(entry, dict):
        entry["last_access"] = time.time()


def _prune_registry() -> None:
    with _registry_lock:
        ttl_seconds, max_size = _registry_limits()
        now = time.time()

        if ttl_seconds and ttl_seconds > 0:
            for key, entry in list(_instance_registry.items()):
                if now - _entry_last_access(entry) > ttl_seconds:
                    del _instance_registry[key]

        if max_size and max_size > 0 and len(_instance_registry) > max_size:
            overflow = len(_instance_registry) - max_size
            oldest = sorted(_instance_registry.items(), key=lambda item: _entry_last_access(item[1]))
            for key, _entry in oldest[:overflow]:
                del _instance_registry[key]


def _store_ref(key: str, obj: Any) -> None:
    now = time.time()
    with _registry_lock:
        _instance_registry[key] = {"obj": obj, "created_at": now, "last_access": now}


def _registry_key(session_id: str, ref_id: str) -> str:
    """
    Construct a registry key from session_id and ref_id.

    This is the single source of truth for registry key format.
    All code that constructs or parses registry keys should use this function.
    """
    return f"{session_id}:{ref_id}"


def _registry_key_prefix(session_id: str) -> str:
    """
    Get the registry key prefix for a session.

    Used for session-scoped operations like release_session.
    Delegates to _registry_key to maintain single source of truth for key format.
    """
    return _registry_key(session_id, "")


def _extract_ref_identity(ref: dict, session_id: str) -> Tuple[str, str]:
    if ref.get("__type__") == "ref":
        ref_id = ref.get("id") or ref.get("ref_id")
        ref_session = ref.get("session_id") or session_id
    elif ref.get("__snakebridge_ref__"):
        ref_id = ref.get("ref_id")
        ref_session = ref.get("session_id") or session_id
    else:
        raise ValueError("Invalid SnakeBridge reference payload")

    if not ref_id:
        raise ValueError("SnakeBridge reference missing id")

    if ref_session and session_id and ref_session != session_id:
        raise ValueError("SnakeBridge reference session mismatch")

    return ref_id, ref_session

def _call_telemetry_span(metadata: Dict[str, Any]):
    if snakepit_telemetry is None:
        return nullcontext()
    try:
        return snakepit_telemetry.span("python.call", metadata)
    except Exception:
        return nullcontext()


def _call_function_name(call_type: str, function: Optional[str], arguments: Dict[str, Any]) -> str:
    if function:
        return str(function)
    if call_type == "class":
        return str(arguments.get("class") or arguments.get("class_name") or "unknown")
    if call_type in ("get_attr", "set_attr", "module_attr"):
        return str(arguments.get("attr") or "unknown")
    if call_type == "helper":
        return str(arguments.get("helper") or "unknown")
    return "unknown"


def _call_metadata(
    call_type: str,
    library: Optional[str],
    python_module: Optional[str],
    function: Optional[str],
    arguments: Dict[str, Any],
) -> Dict[str, Any]:
    metadata = {
        "library": str(library or "unknown"),
        "function": _call_function_name(call_type, function, arguments),
        "call_type": str(call_type),
    }
    if python_module:
        metadata["python_module"] = str(python_module)
    return metadata


def snakebridge_call(module: str, function: str, args: dict) -> dict:
    """
    Call a Python function from SnakeBridge.

    This is the main entry point called by the Snakepit adapter.

    Args:
        module: The Python module name (e.g., 'math', 'numpy')
        function: The function name to call (e.g., 'sqrt', 'array')
        args: Dictionary of argument names to values (in SnakeBridge encoded format)

    Returns:
        Dictionary with either:
            - {"success": True, "result": <encoded_result>}
            - {"success": False, "error": <error_message>, "error_type": <error_type>}

    Examples:
        >>> snakebridge_call('math', 'sqrt', {'x': 16})
        {'success': True, 'result': 4.0}

        >>> snakebridge_call('math', 'gcd', {'a': 48, 'b': 18})
        {'success': True, 'result': 6}

        >>> snakebridge_call('statistics', 'mean', {'data': [1, 2, 3, 4, 5]})
        {'success': True, 'result': 3.0}
    """
    try:
        # Import the module (use cache if available)
        try:
            mod = _import_module(module)
        except ImportError as e:
            return encode_error(ImportError(f"Failed to import module '{module}': {str(e)}"))

        # Get the function from the module
        if not hasattr(mod, function):
            return encode_error(AttributeError(f"Module '{module}' has no function '{function}'"))

        func = getattr(mod, function)

        # Check if it's callable
        if not callable(func):
            return encode_error(TypeError(f"'{module}.{function}' is not callable"))

        # Decode arguments from SnakeBridge format
        try:
            decoded_args = {name: decode(value) for name, value in args.items()}
        except Exception as e:
            return encode_error(ValueError(f"Failed to decode arguments: {str(e)}"))

        # Call the function
        # Try to determine if we should use positional or keyword arguments
        try:
            # First, try with keyword arguments (most flexible)
            try:
                result = func(**decoded_args)
            except TypeError as e:
                error_msg = str(e)
                # If it fails because it doesn't accept keyword arguments,
                # try with positional arguments instead
                if "keyword argument" in error_msg.lower():
                    # Try to get the signature to determine argument order
                    try:
                        sig = inspect.signature(func)
                        # Create positional args in parameter order
                        positional_args = []
                        for param_name in sig.parameters.keys():
                            if param_name in decoded_args:
                                positional_args.append(decoded_args[param_name])

                        # If we didn't find any matching parameters, it might be a *args function
                        # Fall back to using values in insertion order
                        if not positional_args:
                            positional_args = list(decoded_args.values())

                        result = func(*positional_args)
                    except (ValueError, TypeError):
                        # Can't get signature, use values in insertion order (Python 3.7+ dicts)
                        positional_args = list(decoded_args.values())
                        result = func(*positional_args)
                else:
                    # Re-raise if it's a different kind of TypeError
                    raise
        except TypeError as e:
            # Provide helpful error message for argument mismatches
            error_msg = str(e)
            return encode_error(TypeError(f"Argument error calling {module}.{function}: {error_msg}"))
        except Exception as e:
            # Return any exception from the function call
            return encode_error(e)

        # Encode and return the result
        library = module.split(".")[0] if module else "unknown"
        return {
            "success": True,
            "result": encode_result(result, "default", module, library),
        }

    except Exception as e:
        # Catch any unexpected errors
        error_info = {
            "success": False,
            "error": str(e),
            "error_type": type(e).__name__,
            "traceback": traceback.format_exc()
        }
        return error_info


def snakebridge_batch_call(calls: list) -> list:
    """
    Execute multiple function calls in a batch.

    Args:
        calls: List of call specifications, each with 'module', 'function', and 'args'

    Returns:
        List of results corresponding to each call

    Example:
        >>> snakebridge_batch_call([
        ...     {'module': 'math', 'function': 'sqrt', 'args': {'x': 16}},
        ...     {'module': 'math', 'function': 'gcd', 'args': {'a': 48, 'b': 18}}
        ... ])
        [{'success': True, 'result': 4.0}, {'success': True, 'result': 6}]
    """
    results = []
    for call in calls:
        try:
            module = call['module']
            function = call['function']
            args = call.get('args', {})
            result = snakebridge_call(module, function, args)
            results.append(result)
        except Exception as e:
            results.append(encode_error(e))
    return results


def snakebridge_get_attribute(module: str, attribute: str) -> dict:
    """
    Get an attribute or constant from a module.

    Args:
        module: The Python module name
        attribute: The attribute name to get

    Returns:
        Dictionary with either:
            - {"success": True, "result": <encoded_value>}
            - {"success": False, "error": <error_message>}

    Example:
        >>> snakebridge_get_attribute('math', 'pi')
        {'success': True, 'result': 3.141592653589793}
    """
    try:
        # Import the module
        try:
            mod = _import_module(module)
        except ImportError as e:
            return encode_error(ImportError(f"Failed to import module '{module}': {str(e)}"))

        # Get the attribute
        if not hasattr(mod, attribute):
            return encode_error(AttributeError(f"Module '{module}' has no attribute '{attribute}'"))

        value = getattr(mod, attribute)

        # Encode and return the value
        library = module.split(".")[0] if module else "unknown"
        return {
            "success": True,
            "result": encode_result(value, "default", module, library),
        }

    except Exception as e:
        return encode_error(e)


def snakebridge_create_instance(module: str, class_name: str, args: dict) -> dict:
    """
    Create an instance of a class.

    Args:
        module: The Python module name
        class_name: The class name to instantiate
        args: Dictionary of constructor arguments

    Returns:
        Dictionary with either success or error

    Note:
        Instance objects cannot be serialized, so this is mainly useful
        for testing or when combined with a session/state system.
    """
    try:
        # Import the module
        try:
            mod = _import_module(module)
        except ImportError as e:
            return encode_error(ImportError(f"Failed to import module '{module}': {str(e)}"))

        # Get the class
        if not hasattr(mod, class_name):
            return encode_error(AttributeError(f"Module '{module}' has no class '{class_name}'"))

        cls = getattr(mod, class_name)

        # Check if it's a class
        if not isinstance(cls, type):
            return encode_error(TypeError(f"'{module}.{class_name}' is not a class"))

        # Decode arguments
        decoded_args = {name: decode(value) for name, value in args.items()}

        # Create instance
        instance = cls(**decoded_args)

        # Encode and return (note: complex objects may not serialize well)
        library = module.split(".")[0] if module else "unknown"
        return {
            "success": True,
            "result": encode_result(instance, "default", module, library),
        }

    except Exception as e:
        return encode_error(e)


def _import_module(module_name: str) -> Any:
    with _module_cache_lock:
        if module_name in _module_cache:
            return _module_cache[module_name]

        mod = importlib.import_module(module_name)
        _module_cache[module_name] = mod
        return mod


def _make_ref(session_id: str, obj: Any, python_module: str, library: str) -> dict:
    ref_id = uuid.uuid4().hex
    key = _registry_key(session_id, ref_id)
    _prune_registry()
    _store_ref(key, obj)

    type_name = type(obj).__name__

    return {
        "__type__": "ref",
        "__schema__": REF_SCHEMA_VERSION,
        "id": ref_id,
        "session_id": session_id,
        "python_module": python_module,
        "library": library,
        "type_name": type_name,
        "__type_name__": type_name,
    }

def _make_stream_ref(
    session_id: str,
    obj: Any,
    python_module: str,
    library: str,
    stream_type: str,
) -> dict:
    ref_id = uuid.uuid4().hex
    key = _registry_key(session_id, ref_id)
    _prune_registry()
    _store_ref(key, obj)

    type_name = type(obj).__name__

    return {
        "__type__": "stream_ref",
        "__schema__": REF_SCHEMA_VERSION,
        "id": ref_id,
        "session_id": session_id,
        "python_module": python_module,
        "library": library,
        "stream_type": stream_type,
        "type_name": type_name,
        "__type_name__": type_name,
    }


def _truncate_message(msg: str, max_len: int = 200) -> str:
    """Truncate a message to max_len, stripping newlines for log safety."""
    msg = msg.replace('\n', ' ').replace('\r', '')
    if len(msg) > max_len:
        return msg[:max_len] + "..."
    return msg


def _cleanup_created_refs(created_keys: list) -> None:
    """Remove refs from registry that were created during a failed encode."""
    with _registry_lock:
        for key in created_keys:
            _instance_registry.pop(key, None)


def encode_result(result: Any, session_id: str, python_module: str, library: str) -> Any:
    """
    Encode a Python result for transmission to Elixir.

    INVARIANT: Always returns JSON-safe values. Any non-serializable object
    is represented by a ref or stream_ref payload (lossless - the object
    remains accessible via Python operations).

    The returned value is ALWAYS one of:
    1. A JSON-safe primitive (None, bool, int, float, str)
    2. A JSON-safe list/dict (recursively containing the above types)
    3. A tagged value with __type__ that Elixir can decode
    4. A ref ({"__type__": "ref", ...}) for non-serializable objects
    5. A stream_ref ({"__type__": "stream_ref", ...}) for sync iterators

    This function preserves container structure: if a list/dict contains
    non-serializable items, only those items become refs, not the entire
    container. This enables "graceful serialization" where users can access
    all serializable fields even when some fields are non-serializable.

    No lossy conversion is performed (e.g., no model_dump/to_dict calls).
    Repeated references to the same non-serializable object yield the same
    ref payload; repeated references to the same iterator yield the same
    stream_ref payload.
    Async generators become refs (not stream_refs) since they cannot be
    consumed via sync iteration.

    On encoding failure, any refs created during partial encoding are cleaned
    up to prevent unreachable ref leakage in the registry.
    """
    # Use set of object ids to detect cycles
    in_progress = set()
    # Memo table: id(obj) -> ref/stream_ref payload for deduplication
    ref_memo = {}
    # Track registry keys created during this encode for cleanup on failure
    created_keys = []

    try:
        encoded = _encode_result_recursive(
            result, session_id, python_module, library, in_progress, ref_memo, created_keys
        )
    except Exception as e:
        # Any exception during encoding (RecursionError, RuntimeError from
        # concurrent mutation, weird __iter__ implementations, etc.)
        # Clean up any refs created during partial encoding
        _cleanup_created_refs(created_keys)
        # Fall back to ref-wrapping the entire result (lossless)
        _log_warning(
            f"{type(e).__name__} encoding type {type(result).__name__}: "
            f"{_truncate_message(str(e))}, wrapping in ref as fallback"
        )
        return _make_ref(session_id, result, python_module, library)

    # Safety net: validate result is actually JSON-safe
    # This catches any edge cases missed by the recursive encoder
    try:
        is_safe = _is_json_safe(encoded)
    except RecursionError:
        # _is_json_safe itself can recurse deeply on near-limit structures
        # Clean up any refs created during encoding
        _cleanup_created_refs(created_keys)
        _log_warning(
            f"RecursionError in _is_json_safe for type {type(result).__name__}, "
            "wrapping in ref as safety fallback"
        )
        return _make_ref(session_id, result, python_module, library)

    if not is_safe:
        # Clean up any refs created during encoding
        _cleanup_created_refs(created_keys)
        _log_warning(
            f"encode_result produced non-JSON-safe result for type {type(result).__name__}, "
            "wrapping in ref as safety fallback"
        )
        return _make_ref(session_id, result, python_module, library)

    return encoded


def _encode_result_recursive(
    value: Any,
    session_id: str,
    python_module: str,
    library: str,
    in_progress: set,
    ref_memo: dict,
    created_keys: list,
) -> Any:
    """
    Recursively encode a Python value, creating refs for non-serializable leaves.

    Args:
        value: The Python value to encode
        session_id: Session ID for ref tracking
        python_module: Module name for ref metadata
        library: Library name for ref metadata
        in_progress: Set of object ids currently being encoded (for cycle detection)
        ref_memo: Dict mapping id(obj) -> ref payload for deduplication
        created_keys: List to append registry keys of created refs (for cleanup on failure)

    Returns:
        A JSON-safe value with refs for non-serializable objects
    """
    import base64
    import math
    from datetime import datetime, date, time

    # Handle None
    if value is None:
        return None

    # Handle booleans (must check before int - bool is subclass of int)
    if isinstance(value, bool):
        return value

    # Handle integers
    if isinstance(value, int):
        return value

    # Handle float (check for special values)
    if isinstance(value, float):
        if math.isinf(value):
            return {
                "__type__": "special_float",
                "__schema__": SCHEMA_VERSION,
                "value": "infinity" if value > 0 else "neg_infinity",
            }
        if math.isnan(value):
            return {
                "__type__": "special_float",
                "__schema__": SCHEMA_VERSION,
                "value": "nan",
            }
        return value

    # Handle strings
    if isinstance(value, str):
        return value

    # Handle bytes/bytearray - always tagged
    if isinstance(value, (bytes, bytearray)):
        return {
            "__type__": "bytes",
            "__schema__": SCHEMA_VERSION,
            "data": base64.b64encode(bytes(value)).decode("ascii"),
        }

    # Handle complex
    if isinstance(value, complex):
        return {
            "__type__": "complex",
            "__schema__": SCHEMA_VERSION,
            "real": value.real,
            "imag": value.imag,
        }

    # Handle datetime types
    if isinstance(value, datetime):
        return {
            "__type__": "datetime",
            "__schema__": SCHEMA_VERSION,
            "value": value.isoformat(),
        }
    if isinstance(value, date):
        return {
            "__type__": "date",
            "__schema__": SCHEMA_VERSION,
            "value": value.isoformat(),
        }
    if isinstance(value, time):
        return {
            "__type__": "time",
            "__schema__": SCHEMA_VERSION,
            "value": value.isoformat(),
        }

    # Handle Atom - tagged atom value for Elixir interop
    if isinstance(value, Atom):
        return {
            "__type__": "atom",
            "__schema__": SCHEMA_VERSION,
            "value": value.value,
        }

    # Check for async generators FIRST - these cannot be consumed via next()
    # so we treat them as regular refs, not stream_refs
    if hasattr(types, 'AsyncGeneratorType') and isinstance(value, types.AsyncGeneratorType):
        obj_id = id(value)
        if obj_id in ref_memo:
            return ref_memo[obj_id]
        ref_payload = _make_ref(session_id, value, python_module, library)
        created_keys.append(_registry_key(session_id, ref_payload['id']))
        ref_memo[obj_id] = ref_payload
        return ref_payload

    # Check for sync generators/iterators - create stream_ref in-place
    if _is_generator_or_iterator(value):
        obj_id = id(value)
        if obj_id in ref_memo:
            return ref_memo[obj_id]
        stream_type = _get_stream_type(value)
        stream_ref_payload = _make_stream_ref(session_id, value, python_module, library, stream_type)
        created_keys.append(_registry_key(session_id, stream_ref_payload['id']))
        ref_memo[obj_id] = stream_ref_payload
        return stream_ref_payload

    # For container types, check for cycles
    obj_id = id(value)

    # Cycle detection: if we're already encoding this object, create a ref
    # Use memo if available to ensure same ref for same object
    if obj_id in in_progress:
        if obj_id in ref_memo:
            return ref_memo[obj_id]
        ref_payload = _make_ref(session_id, value, python_module, library)
        created_keys.append(_registry_key(session_id, ref_payload['id']))
        ref_memo[obj_id] = ref_payload
        return ref_payload

    # Handle list - recursively encode elements, refs for non-encodable
    # Snapshot to avoid "list changed size during iteration" errors
    if isinstance(value, list):
        in_progress.add(obj_id)
        try:
            items = list(value)  # Snapshot
            return [
                _encode_result_recursive(item, session_id, python_module, library, in_progress, ref_memo, created_keys)
                for item in items
            ]
        finally:
            in_progress.discard(obj_id)

    # Handle tuple - recursively encode elements
    # Tuples are immutable, but snapshot for consistency
    if isinstance(value, tuple):
        in_progress.add(obj_id)
        try:
            items = tuple(value)  # Snapshot (tuples are immutable but be consistent)
            elements = [
                _encode_result_recursive(item, session_id, python_module, library, in_progress, ref_memo, created_keys)
                for item in items
            ]
            return {
                "__type__": "tuple",
                "__schema__": SCHEMA_VERSION,
                "elements": elements,
            }
        finally:
            in_progress.discard(obj_id)

    # Handle set - recursively encode elements (no sorting; sets have no order semantics)
    # Snapshot to avoid "set changed size during iteration" errors
    if isinstance(value, set):
        in_progress.add(obj_id)
        try:
            items = list(value)  # Snapshot
            elements = [
                _encode_result_recursive(item, session_id, python_module, library, in_progress, ref_memo, created_keys)
                for item in items
            ]
            return {
                "__type__": "set",
                "__schema__": SCHEMA_VERSION,
                "elements": elements,
            }
        finally:
            in_progress.discard(obj_id)

    # Handle frozenset - recursively encode elements (no sorting; sets have no order semantics)
    # Frozensets are immutable, but snapshot for consistency
    if isinstance(value, frozenset):
        in_progress.add(obj_id)
        try:
            items = list(value)  # Snapshot
            elements = [
                _encode_result_recursive(item, session_id, python_module, library, in_progress, ref_memo, created_keys)
                for item in items
            ]
            return {
                "__type__": "frozenset",
                "__schema__": SCHEMA_VERSION,
                "elements": elements,
            }
        finally:
            in_progress.discard(obj_id)

    # Handle dict - recursively encode values (and keys if non-string)
    # Snapshot to avoid "dictionary changed size during iteration" errors
    if isinstance(value, dict):
        in_progress.add(obj_id)
        try:
            items = list(value.items())  # Snapshot
            all_string_keys = all(isinstance(k, str) for k, _ in items)

            if all_string_keys:
                # Plain dict with string keys
                return {
                    k: _encode_result_recursive(v, session_id, python_module, library, in_progress, ref_memo, created_keys)
                    for k, v in items
                }
            else:
                # Tagged dict with pairs for non-string keys
                pairs = [
                    [
                        _encode_result_recursive(k, session_id, python_module, library, in_progress, ref_memo, created_keys),
                        _encode_result_recursive(v, session_id, python_module, library, in_progress, ref_memo, created_keys),
                    ]
                    for k, v in items
                ]
                return {
                    "__type__": "dict",
                    "__schema__": SCHEMA_VERSION,
                    "pairs": pairs,
                }
        finally:
            in_progress.discard(obj_id)

    # EVERYTHING ELSE - create a ref (with memoization)
    # This includes custom objects, functions, modules, etc.
    if obj_id in ref_memo:
        return ref_memo[obj_id]
    ref_payload = _make_ref(session_id, value, python_module, library)
    created_keys.append(_registry_key(session_id, ref_payload['id']))
    ref_memo[obj_id] = ref_payload
    return ref_payload


def _is_streamable(value: Any) -> bool:
    """
    Check if value can be streamed (consumed via sync iteration).

    NOTE: Async generators are explicitly EXCLUDED - they cannot be consumed
    via next() and must be treated as regular refs.
    """
    # Explicitly exclude async generators - they cannot be consumed via next()
    if hasattr(types, 'AsyncGeneratorType') and isinstance(value, types.AsyncGeneratorType):
        return False
    if _is_generator_or_iterator(value):
        return True
    if hasattr(value, "__iter__"):
        if isinstance(value, (str, bytes, bytearray, list, tuple, dict, set, frozenset)):
            return False
        # Prefer regular refs for context managers (e.g. file objects)
        if hasattr(value, "__enter__") and hasattr(value, "__exit__"):
            return False
        return True
    return False


def _dynamic_stream_iterator(
    result: Any,
    session_id: str,
    python_module: str,
    library: str,
):
    if _is_streamable(result):
        iterator = result
        if not hasattr(iterator, "__next__") and hasattr(iterator, "__iter__"):
            iterator = iter(iterator)
        for item in iterator:
            yield encode_result(item, session_id, python_module, library)
    else:
        yield encode_result(result, session_id, python_module, library)


def _is_json_safe(value: Any) -> bool:
    """
    Verify a value is safe to serialize as JSON.

    This is a safety check after encoding - if encode() is correct,
    this should always return True.
    """
    import math

    if value is None:
        return True
    if isinstance(value, bool):
        return True
    if isinstance(value, (int, float)):
        # Check for non-JSON floats (should be tagged already)
        if isinstance(value, float) and (math.isinf(value) or math.isnan(value)):
            return False
        return True
    if isinstance(value, str):
        return True
    if isinstance(value, list):
        return all(_is_json_safe(item) for item in value)
    if isinstance(value, dict):
        # Check for valid tagged types
        type_tag = value.get("__type__")
        if type_tag in (
            "bytes", "tuple", "set", "frozenset", "complex",
            "datetime", "date", "time", "special_float",
            "atom", "dict", "ref", "stream_ref", "callback",
            "stop_iteration"
        ):
            # Tagged values are safe - validate values recursively
            return all(_is_json_safe(v) for v in value.values())
        # Regular dict - keys must be strings
        if not all(isinstance(k, str) for k in value.keys()):
            return False
        return all(_is_json_safe(v) for v in value.values())
    return False


def _log_warning(message: str) -> None:
    """Log a warning message."""
    print(f"[SnakeBridge WARNING] {message}", file=sys.stderr)


def _is_ref_payload(value: Any) -> bool:
    if not isinstance(value, dict):
        return False
    type_tag = value.get("__type__")
    if type_tag is not None:
        return type_tag == "ref"
    return "id" in value and ("session_id" in value or "ref_id" in value)


def _resolve_refs(value: Any, session_id: str) -> Any:
    if isinstance(value, dict):
        if _is_ref_payload(value):
            return _resolve_ref(value, session_id)
        return {k: _resolve_refs(v, session_id) for k, v in value.items()}
    if isinstance(value, list):
        return [_resolve_refs(item, session_id) for item in value]
    if isinstance(value, tuple):
        return tuple(_resolve_refs(item, session_id) for item in value)
    return value


def _resolve_ref(ref: dict, session_id: str) -> Any:
    if not isinstance(ref, dict):
        raise ValueError("Invalid SnakeBridge reference payload")

    with _registry_lock:
        _prune_registry()
        ref_id, ref_session = _extract_ref_identity(ref, session_id)
        key = _registry_key(ref_session, ref_id)

        if key not in _instance_registry:
            raise KeyError(f"Unknown SnakeBridge reference: {ref_id}")

        entry = _instance_registry[key]
        if isinstance(entry, dict):
            _touch_entry(entry)
            return entry.get("obj")
        return entry


def _release_ref(ref: dict, session_id: str) -> bool:
    if not isinstance(ref, dict):
        raise ValueError("Invalid SnakeBridge reference payload")

    with _registry_lock:
        _prune_registry()
        ref_id, ref_session = _extract_ref_identity(ref, session_id)
        key = _registry_key(ref_session, ref_id)

        if key in _instance_registry:
            del _instance_registry[key]
            return True
        return False


def _release_session(session_id: str) -> int:
    if not session_id:
        return 0

    with _registry_lock:
        removed = 0
        prefix = _registry_key_prefix(session_id)
        for key in list(_instance_registry.keys()):
            if key.startswith(prefix):
                del _instance_registry[key]
                removed += 1

        return removed


def _default_helper_config() -> Dict[str, Any]:
    return {
        "helper_paths": ["priv/python/helpers"],
        "helper_pack_enabled": True,
        "helper_allowlist": "all"
    }


def _normalize_helper_config(config: Optional[Dict[str, Any]]) -> Dict[str, Any]:
    normalized = _default_helper_config()
    if isinstance(config, dict):
        normalized.update(config)

    helper_paths = normalized.get("helper_paths") or []
    if isinstance(helper_paths, str):
        helper_paths = [helper_paths]

    helper_paths = [os.path.abspath(path) for path in helper_paths if path]
    normalized["helper_paths"] = helper_paths
    normalized["helper_pack_enabled"] = bool(normalized.get("helper_pack_enabled", True))

    allowlist = normalized.get("helper_allowlist", "all")
    if allowlist in [None, "all", ":all"]:
        normalized["helper_allowlist"] = "all"
    elif isinstance(allowlist, (list, tuple, set)):
        normalized["helper_allowlist"] = [str(item) for item in allowlist]
    else:
        normalized["helper_allowlist"] = [str(allowlist)]

    return normalized


def _helper_config_key(config: Dict[str, Any]) -> Tuple[Any, ...]:
    helper_paths = tuple(config.get("helper_paths", []))
    allowlist = config.get("helper_allowlist", "all")
    allowlist_key = "all" if allowlist == "all" else tuple(allowlist)

    return (helper_paths, allowlist_key, bool(config.get("helper_pack_enabled", True)))


def _resolve_helper_paths(config: Dict[str, Any]) -> List[str]:
    paths: List[str] = []
    if config.get("helper_pack_enabled", True):
        pack_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "helpers")
        paths.append(pack_path)

    paths.extend(config.get("helper_paths", []))
    return paths


def _list_helper_files(paths: List[str]) -> List[str]:
    files: List[str] = []
    for path in paths:
        if not path:
            continue
        if os.path.isdir(path):
            for file_path in sorted(glob.glob(os.path.join(path, "*.py"))):
                base = os.path.basename(file_path)
                if base == "__init__.py" or base.startswith("_"):
                    continue
                files.append(file_path)
        elif os.path.isfile(path):
            base = os.path.basename(path)
            if base != "__init__.py" and not base.startswith("_"):
                files.append(path)
    return files


def _import_helper_module(path: str) -> Any:
    module_name = f"snakebridge_helper_{hashlib.md5(path.encode('utf-8')).hexdigest()}"
    if module_name in sys.modules:
        return sys.modules[module_name]

    spec = importlib.util.spec_from_file_location(module_name, path)
    if spec is None or spec.loader is None:
        raise SnakeBridgeHelperLoadError(f"Unable to load helper module: {path}")

    module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)
    sys.modules[module_name] = module
    return module


def _extract_helpers_from_module(module: Any) -> Dict[str, Any]:
    if hasattr(module, "__snakebridge_helpers__"):
        helpers = getattr(module, "__snakebridge_helpers__")
    elif hasattr(module, "snakebridge_helpers"):
        helpers = module.snakebridge_helpers()
    elif hasattr(module, "HELPERS"):
        helpers = getattr(module, "HELPERS")
    else:
        return {}

    if not isinstance(helpers, dict):
        raise SnakeBridgeHelperLoadError("Helper registry must be a dict of name => callable")

    for name, func in helpers.items():
        if not isinstance(name, str):
            raise SnakeBridgeHelperLoadError("Helper names must be strings")
        if not callable(func):
            raise SnakeBridgeHelperLoadError(f"Helper '{name}' is not callable")

    return helpers


def _apply_allowlist(helpers: Dict[str, Any], allowlist: Any) -> Dict[str, Any]:
    if allowlist == "all":
        return helpers
    if not allowlist:
        return {}

    return {name: func for name, func in helpers.items() if name in allowlist}


def _load_helper_registry(config: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
    global _helper_registry, _helper_registry_key, _helper_registry_index

    with _helper_lock:
        normalized = _normalize_helper_config(config)
        key = _helper_config_key(normalized)

        if key == _helper_registry_key:
            return _helper_registry

        registry: Dict[str, Any] = {}
        for path in _resolve_helper_paths(normalized):
            if not path or not os.path.exists(path):
                continue

            for file_path in _list_helper_files([path]):
                module = _import_helper_module(file_path)
                helpers = _extract_helpers_from_module(module)
                registry.update(helpers)

        registry = _apply_allowlist(registry, normalized.get("helper_allowlist", "all"))
        _helper_registry = registry
        _helper_registry_key = key
        _helper_registry_index = _build_helper_index(registry)

        return registry


def _format_annotation(annotation: Any) -> Optional[str]:
    if annotation is inspect.Signature.empty:
        return None
    if hasattr(annotation, "__name__"):
        return annotation.__name__
    return str(annotation)


def _param_info(param: inspect.Parameter) -> Dict[str, Any]:
    info: Dict[str, Any] = {"name": param.name, "kind": param.kind.name}
    if param.default is not inspect.Parameter.empty:
        info["default"] = repr(param.default)
    if param.annotation is not inspect.Parameter.empty:
        info["annotation"] = _format_annotation(param.annotation)
    return info


def _build_helper_index(helpers: Dict[str, Any]) -> List[Dict[str, Any]]:
    index: List[Dict[str, Any]] = []
    for name, func in helpers.items():
        entry = {"name": name}

        try:
            sig = inspect.signature(func)
            entry["parameters"] = [_param_info(p) for p in sig.parameters.values()]
        except (ValueError, TypeError):
            entry["parameters"] = []

        doc = inspect.getdoc(func) or ""
        if doc:
            entry["docstring"] = doc[:8000]

        index.append(entry)

    index.sort(key=lambda item: item.get("name", ""))
    return index


def helper_registry_index(config: Optional[Dict[str, Any]] = None) -> List[Dict[str, Any]]:
    _load_helper_registry(config or {})
    return _helper_registry_index


class SnakeBridgeAdapter:
    def __init__(self):
        self.session_context = None

    def set_session_context(self, session_context):
        self.session_context = session_context

    def execute_tool(self, tool_name: str, arguments: dict, context):
        if tool_name == "snakebridge.helpers":
            helper_config = {}
            if isinstance(arguments, dict):
                helper_config = arguments.get("helper_config") or arguments
            return helper_registry_index(helper_config)

        if tool_name not in [
            "snakebridge.call",
            "snakebridge.stream",
            "snakebridge.release_ref",
            "snakebridge.release_session",
        ]:
            raise AttributeError(f"Tool '{tool_name}' not supported by SnakeBridgeAdapter")

        if isinstance(arguments, dict):
            _protocol_compatibility(arguments)

        session_id = None
        if isinstance(arguments, dict) and arguments.get("session_id"):
            session_id = arguments.get("session_id")
        elif context is not None and hasattr(context, "session_id"):
            session_id = context.session_id
        elif self.session_context is not None:
            session_id = self.session_context.session_id
        else:
            session_id = "default"

        if tool_name == "snakebridge.release_ref":
            ref = arguments.get("ref") if isinstance(arguments, dict) else None
            if ref is None:
                raise ValueError("snakebridge.release_ref requires ref")
            return _release_ref(ref, session_id)

        if tool_name == "snakebridge.release_session":
            return _release_session(session_id)

        call_type = arguments.get("call_type") or "function"
        module_path = arguments.get("module_path")
        python_module = arguments.get("python_module") or arguments.get("module")
        function = arguments.get("function")
        args = arguments.get("args") or []
        kwargs = arguments.get("kwargs") or {}
        if call_type in ("dynamic", "dynamic_stream") and not python_module:
            python_module = module_path
        library = arguments.get("library") or (python_module.split(".")[0] if python_module else None)
        if call_type not in ("helper", "stream_next") and not python_module:
            raise ValueError("snakebridge.call requires python_module")
        if not python_module:
            python_module = library or "unknown"
        if not library:
            library = python_module.split(".")[0] if python_module else "unknown"
        metadata = _call_metadata(call_type, library, python_module, function, arguments)
        decoded_args = [decode(item, session_id=session_id, context=context) for item in args]
        decoded_kwargs = {
            key: decode(value, session_id=session_id, context=context) for key, value in kwargs.items()
        }
        decoded_args = [_resolve_refs(item, session_id) for item in decoded_args]
        decoded_kwargs = {key: _resolve_refs(value, session_id) for key, value in decoded_kwargs.items()}

        with _call_telemetry_span(metadata):
            if call_type == "stream_next":
                stream_ref_payload = arguments.get("stream_ref")
                if stream_ref_payload is None:
                    raise ValueError("snakebridge.call requires stream_ref for stream_next")

                stream_ref = decode(stream_ref_payload, session_id=session_id, context=context)
                iterator = _resolve_ref(stream_ref, session_id)
                python_module = ""
                library = ""

                if isinstance(stream_ref_payload, dict):
                    python_module = stream_ref_payload.get("python_module", "") or ""
                    library = stream_ref_payload.get("library", "") or ""

                try:
                    item = next(iterator)
                    return encode_result(item, session_id, python_module, library)
                except StopIteration:
                    return {"__type__": "stop_iteration"}

            if call_type == "dynamic_stream":
                module_path = arguments.get("module_path") or python_module
                if not module_path:
                    raise ValueError("snakebridge.call requires module_path for dynamic calls")
                mod = _import_module(module_path)
                func = getattr(mod, function)
                result = func(*decoded_args, **decoded_kwargs)
                return _dynamic_stream_iterator(result, session_id, module_path, library)

            if call_type == "dynamic":
                module_path = arguments.get("module_path") or python_module
                if not module_path:
                    raise ValueError("snakebridge.call requires module_path for dynamic calls")
                mod = _import_module(module_path)
                func = getattr(mod, function)
                result = func(*decoded_args, **decoded_kwargs)
                return encode_result(result, session_id, module_path, library)

            if call_type == "class":
                class_name = arguments.get("class") or arguments.get("class_name")
                mod = _import_module(python_module)
                cls = getattr(mod, class_name)
                instance = cls(*decoded_args, **decoded_kwargs)
                return encode_result(instance, session_id, python_module, library)

            if call_type == "method":
                instance_payload = arguments.get("instance")
                instance = _resolve_ref(decode(instance_payload), session_id)
                method = getattr(instance, function)
                result = method(*decoded_args, **decoded_kwargs)
                return encode_result(result, session_id, python_module, library)

            if call_type == "get_attr":
                instance_payload = arguments.get("instance")
                instance = _resolve_ref(decode(instance_payload), session_id)
                attr = arguments.get("attr") or function
                result = getattr(instance, attr)
                return encode_result(result, session_id, python_module, library)

            if call_type == "module_attr":
                attr = arguments.get("attr") or function
                mod = _import_module(python_module)
                result = getattr(mod, attr)
                return encode_result(result, session_id, python_module, library)

            if call_type == "set_attr":
                instance_payload = arguments.get("instance")
                instance = _resolve_ref(decode(instance_payload), session_id)
                attr = arguments.get("attr") or function
                value = decoded_args[0] if decoded_args else None
                setattr(instance, attr, value)
                return encode_result(True, session_id, python_module, library)

            if call_type == "helper":
                helper_name = arguments.get("helper") or function
                helper_config = arguments.get("helper_config") or {}

                if not helper_name:
                    raise SnakeBridgeHelperNotFoundError("Helper name is required")

                registry = _load_helper_registry(helper_config)
                if helper_name not in registry:
                    raise SnakeBridgeHelperNotFoundError(f"Helper '{helper_name}' not found")

                result = registry[helper_name](*decoded_args, **decoded_kwargs)
                return encode_result(result, session_id, python_module, library)

            mod = _import_module(python_module)
            func = getattr(mod, function)
            result = func(*decoded_args, **decoded_kwargs)
            return encode_result(result, session_id, python_module, library)


# Make the module callable for testing
if __name__ == "__main__":
    import json

    # Simple test runner
    if len(sys.argv) > 1:
        # Test with command-line arguments
        # Usage: python snakebridge_adapter.py <module> <function> <json_args>
        if len(sys.argv) >= 4:
            module = sys.argv[1]
            function = sys.argv[2]
            args_json = sys.argv[3]
            args = json.loads(args_json)

            result = snakebridge_call(module, function, args)
            print(json.dumps(result, indent=2))
        else:
            print("Usage: python snakebridge_adapter.py <module> <function> <json_args>")
    else:
        # Run built-in tests
        print("Running SnakeBridge adapter tests...\n")

        # Test 1: math.sqrt
        print("Test 1: math.sqrt(16)")
        result = snakebridge_call('math', 'sqrt', {'x': 16})
        print(json.dumps(result, indent=2))
        assert result['success'] == True
        assert result['result'] == 4.0
        print("PASS\n")

        # Test 2: math.gcd
        print("Test 2: math.gcd(48, 18)")
        result = snakebridge_call('math', 'gcd', {'a': 48, 'b': 18})
        print(json.dumps(result, indent=2))
        assert result['success'] == True
        assert result['result'] == 6
        print("PASS\n")

        # Test 3: Error handling - module not found
        print("Test 3: Error handling - nonexistent module")
        result = snakebridge_call('nonexistent_module', 'func', {})
        print(json.dumps(result, indent=2))
        assert result['success'] == False
        print("PASS\n")

        # Test 4: Error handling - function not found
        print("Test 4: Error handling - nonexistent function")
        result = snakebridge_call('math', 'nonexistent_function', {})
        print(json.dumps(result, indent=2))
        assert result['success'] == False
        print("PASS\n")

        # Test 5: Complex types - tuple encoding
        print("Test 5: Complex types - math.gcd with large numbers")
        result = snakebridge_call('math', 'gcd', {'a': 1071, 'b': 462})
        print(json.dumps(result, indent=2))
        assert result['success'] == True
        assert result['result'] == 21
        print("PASS\n")

        # Test 6: Get attribute
        print("Test 6: Get math.pi")
        result = snakebridge_get_attribute('math', 'pi')
        print(json.dumps(result, indent=2))
        assert result['success'] == True
        assert abs(result['result'] - 3.141592653589793) < 0.0001
        print("PASS\n")

        print("All tests passed!")