snakebridge

"""
Tests for SnakeBridge adapter encode_result graceful serialization.

These tests verify that containers with non-serializable items preserve
their structure, with only the non-serializable leaf objects becoming refs.
"""

import sys
import os

# Add parent directory to path for imports
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

from snakebridge_adapter import (
    encode_result,
    _instance_registry,
    _resolve_ref,
    _is_json_safe,
    _registry_key_prefix,
)
from snakebridge_types import Atom, SCHEMA_VERSION


class CustomObject:
    """A custom class that cannot be JSON serialized."""
    def __init__(self, value=42):
        self.value = value


class HashableCustom:
    """A hashable custom class for set testing."""
    def __init__(self, value=42):
        self.value = value

    def __hash__(self):
        return hash(self.value)

    def __eq__(self, other):
        return isinstance(other, HashableCustom) and self.value == other.value


class TestEncodeResultGraceful:
    """Test that encode_result preserves container structure with nested refs."""

    def test_list_with_custom_object_preserves_structure(self):
        """List with custom object should preserve list, ref-wrap only the object."""
        obj = CustomObject()
        result = encode_result([1, obj, 3], "test-session", "test", "test")

        # Result should be a list, not a ref
        assert isinstance(result, list), f"Expected list, got {type(result)}"
        assert len(result) == 3

        # First and third elements should be primitives
        assert result[0] == 1
        assert result[2] == 3

        # Second element should be a ref payload
        ref_payload = result[1]
        assert isinstance(ref_payload, dict)
        assert ref_payload.get("__type__") == "ref"
        assert "id" in ref_payload
        assert ref_payload.get("session_id") == "test-session"
        assert ref_payload.get("type_name") == "CustomObject"
        assert ref_payload.get("__type_name__") == "CustomObject"

    def test_dict_with_custom_object_preserves_structure(self):
        """Dict with custom object value should preserve dict, ref-wrap only the value."""
        obj = CustomObject()
        result = encode_result({"a": 1, "b": obj}, "test-session", "test", "test")

        # Result should be a dict, not a ref
        assert isinstance(result, dict), f"Expected dict, got {type(result)}"
        assert "__type__" not in result or result.get("__type__") != "ref"

        # Key "a" should have primitive value
        assert result["a"] == 1

        # Key "b" should have ref payload
        ref_payload = result["b"]
        assert isinstance(ref_payload, dict)
        assert ref_payload.get("__type__") == "ref"
        assert ref_payload.get("type_name") == "CustomObject"

    def test_nested_structure_dspy_history_like(self):
        """DSPy history-like structure should preserve list/dict with only response as ref."""
        response_obj = CustomObject()
        history = [
            {
                "model": "gpt-4",
                "response": response_obj,
                "cost": 1.23,
                "prompt_tokens": 100,
                "completion_tokens": 50,
            }
        ]

        result = encode_result(history, "test-session", "test", "test")

        # Result should be a list
        assert isinstance(result, list), f"Expected list, got {type(result)}"
        assert len(result) == 1

        # Entry should be a dict
        entry = result[0]
        assert isinstance(entry, dict), f"Expected dict, got {type(entry)}"

        # Serializable fields should be preserved
        assert entry["model"] == "gpt-4"
        assert entry["cost"] == 1.23
        assert entry["prompt_tokens"] == 100
        assert entry["completion_tokens"] == 50

        # Response should be a ref
        ref_payload = entry["response"]
        assert isinstance(ref_payload, dict)
        assert ref_payload.get("__type__") == "ref"
        assert ref_payload.get("type_name") == "CustomObject"

    def test_cycle_detection(self):
        """Cyclic structures should not cause infinite recursion."""
        lst = []
        lst.append(lst)  # Self-referential list

        result = encode_result(lst, "test-session", "test", "test")

        # Result should be a list containing a ref (not infinite recursion)
        assert isinstance(result, list), f"Expected list, got {type(result)}"
        assert len(result) == 1

        # The nested element should be a ref pointing to the same list
        ref_payload = result[0]
        assert isinstance(ref_payload, dict)
        assert ref_payload.get("__type__") == "ref"

    def test_iterator_nested_in_dict(self):
        """Iterator nested in dict should become stream_ref, dict preserved."""
        gen = (x for x in range(3))
        result = encode_result({"stream": gen, "ok": 1}, "test-session", "test", "test")

        # Result should be a dict, not a ref
        assert isinstance(result, dict)
        assert "__type__" not in result or result.get("__type__") != "ref"

        # "ok" should be preserved
        assert result["ok"] == 1

        # "stream" should be a stream_ref
        stream_ref = result["stream"]
        assert isinstance(stream_ref, dict)
        assert stream_ref.get("__type__") == "stream_ref"

    def test_ref_stored_in_registry(self):
        """Ref payloads should have objects stored in registry."""
        obj = CustomObject(value=999)
        result = encode_result([obj], "test-session", "test", "test")

        # Get the ref payload
        ref_payload = result[0]
        ref_id = ref_payload["id"]
        session_id = ref_payload["session_id"]

        # Verify registry contains the object
        key = f"{session_id}:{ref_id}"
        assert key in _instance_registry

        # Resolve and verify identity
        resolved = _resolve_ref(ref_payload, session_id)
        assert resolved is obj
        assert resolved.value == 999

    def test_deeply_nested_custom_object(self):
        """Deeply nested custom objects should become refs at correct depth."""
        obj = CustomObject()
        nested = {"level1": {"level2": {"level3": [1, 2, obj, 4]}}}

        result = encode_result(nested, "test-session", "test", "test")

        # Verify structure is preserved
        assert isinstance(result, dict)
        assert isinstance(result["level1"], dict)
        assert isinstance(result["level1"]["level2"], dict)
        assert isinstance(result["level1"]["level2"]["level3"], list)

        # Verify the custom object is a ref
        inner_list = result["level1"]["level2"]["level3"]
        assert inner_list[0] == 1
        assert inner_list[1] == 2
        assert inner_list[3] == 4
        assert inner_list[2].get("__type__") == "ref"

    def test_tuple_with_custom_object(self):
        """Tuple with custom object should preserve tuple structure."""
        obj = CustomObject()
        result = encode_result((1, obj, 3), "test-session", "test", "test")

        # Result should be a tagged tuple
        assert isinstance(result, dict)
        assert result.get("__type__") == "tuple"

        elements = result["elements"]
        assert len(elements) == 3
        assert elements[0] == 1
        assert elements[2] == 3

        # Middle element should be ref
        assert elements[1].get("__type__") == "ref"

    def test_set_with_custom_object(self):
        """Set with custom object should preserve set structure."""
        obj = HashableCustom()
        result = encode_result({1, 2, obj}, "test-session", "test", "test")

        # Result should be a tagged set
        assert isinstance(result, dict)
        assert result.get("__type__") == "set"

        elements = result["elements"]
        # Should have 3 elements
        assert len(elements) == 3

        # One element should be a ref
        refs = [e for e in elements if isinstance(e, dict) and e.get("__type__") == "ref"]
        assert len(refs) == 1

    def test_frozenset_with_custom_object(self):
        """Frozenset with custom object should preserve frozenset structure."""
        obj = HashableCustom()
        result = encode_result(frozenset([1, 2, obj]), "test-session", "test", "test")

        # Result should be a tagged frozenset
        assert isinstance(result, dict)
        assert result.get("__type__") == "frozenset"

        elements = result["elements"]
        refs = [e for e in elements if isinstance(e, dict) and e.get("__type__") == "ref"]
        assert len(refs) == 1

    def test_dict_with_non_string_keys(self):
        """Dict with non-string keys should use tagged dict format with nested refs."""
        obj = CustomObject()
        result = encode_result({1: "one", 2: obj}, "test-session", "test", "test")

        # Result should be a tagged dict
        assert isinstance(result, dict)
        assert result.get("__type__") == "dict"

        pairs = result["pairs"]
        assert len(pairs) == 2

        # Find the pair with value=obj (should be ref)
        for pair in pairs:
            key, val = pair[0], pair[1]
            if key == 2:
                assert isinstance(val, dict)
                assert val.get("__type__") == "ref"

    def test_multiple_custom_objects_in_list(self):
        """Multiple custom objects should each become separate refs."""
        obj1 = CustomObject(1)
        obj2 = CustomObject(2)
        result = encode_result([obj1, "middle", obj2], "test-session", "test", "test")

        assert isinstance(result, list)
        assert len(result) == 3
        assert result[1] == "middle"

        # Both objects should be refs with different IDs
        ref1 = result[0]
        ref2 = result[2]
        assert ref1.get("__type__") == "ref"
        assert ref2.get("__type__") == "ref"
        assert ref1["id"] != ref2["id"]

    def test_result_is_json_safe(self):
        """encode_result output should always be JSON-safe."""
        obj = CustomObject()
        cases = [
            [1, obj, 3],
            {"a": obj, "b": 2},
            {"nested": {"deep": [obj]}},
            (1, obj),
            {1, HashableCustom()},
        ]

        for case in cases:
            result = encode_result(case, "test-session", "test", "test")
            assert _is_json_safe(result), f"Result not JSON-safe for {case}"

    def test_empty_containers(self):
        """Empty containers should encode correctly."""
        assert encode_result([], "test", "test", "test") == []
        assert encode_result({}, "test", "test", "test") == {}

        tuple_result = encode_result((), "test", "test", "test")
        assert tuple_result.get("__type__") == "tuple"
        assert tuple_result["elements"] == []

    def test_primitives_pass_through(self):
        """Primitive values should pass through unchanged."""
        assert encode_result(None, "test", "test", "test") is None
        assert encode_result(True, "test", "test", "test") is True
        assert encode_result(False, "test", "test", "test") is False
        assert encode_result(42, "test", "test", "test") == 42
        assert encode_result(3.14, "test", "test", "test") == 3.14
        assert encode_result("hello", "test", "test", "test") == "hello"

    def test_bytes_encoding(self):
        """Bytes should be tagged correctly."""
        result = encode_result(b"hello", "test", "test", "test")
        assert result.get("__type__") == "bytes"
        assert result.get("data") == "aGVsbG8="

    def test_datetime_encoding(self):
        """Datetime should be tagged correctly."""
        from datetime import datetime
        dt = datetime(2024, 1, 15, 12, 30, 45)
        result = encode_result(dt, "test", "test", "test")
        assert result.get("__type__") == "datetime"
        assert "2024-01-15" in result.get("value", "")


class TestEncodeResultRefMetadata:
    """Test that ref payloads include required metadata."""

    def test_ref_has_type_name(self):
        """Ref payload should include type_name."""
        obj = CustomObject()
        result = encode_result([obj], "test-session", "test", "test")

        ref_payload = result[0]
        assert "type_name" in ref_payload
        assert ref_payload["type_name"] == "CustomObject"

    def test_ref_has_double_underscore_type_name(self):
        """Ref payload should include __type_name__ for SnakeBridge.Ref compatibility."""
        obj = CustomObject()
        result = encode_result([obj], "test-session", "test", "test")

        ref_payload = result[0]
        assert "__type_name__" in ref_payload
        assert ref_payload["__type_name__"] == "CustomObject"

    def test_ref_has_session_id(self):
        """Ref payload should include session_id."""
        obj = CustomObject()
        result = encode_result([obj], "my-session", "test", "test")

        ref_payload = result[0]
        assert ref_payload.get("session_id") == "my-session"

    def test_ref_has_python_module(self):
        """Ref payload should include python_module."""
        obj = CustomObject()
        result = encode_result([obj], "test", "mymodule", "mylib")

        ref_payload = result[0]
        assert ref_payload.get("python_module") == "mymodule"

    def test_ref_has_library(self):
        """Ref payload should include library."""
        obj = CustomObject()
        result = encode_result([obj], "test", "mymodule", "mylib")

        ref_payload = result[0]
        assert ref_payload.get("library") == "mylib"


class TestCycleDetection:
    """Test cycle detection in various structures."""

    def test_self_referential_list(self):
        """Self-referential list should not cause stack overflow."""
        lst = [1, 2]
        lst.append(lst)

        result = encode_result(lst, "test", "test", "test")
        assert isinstance(result, list)
        # The third element should be a ref to avoid infinite recursion
        assert result[2].get("__type__") == "ref"

    def test_mutually_referential_dicts(self):
        """Mutually referential dicts should not cause stack overflow."""
        a = {"name": "a"}
        b = {"name": "b"}
        a["other"] = b
        b["other"] = a

        result = encode_result(a, "test", "test", "test")
        assert isinstance(result, dict)
        assert result["name"] == "a"

        # b should be a dict with "other" being a ref back to a
        b_result = result["other"]
        assert isinstance(b_result, dict)
        assert b_result["name"] == "b"
        assert b_result["other"].get("__type__") == "ref"

    def test_deeply_nested_cycle(self):
        """Deeply nested cycles should be detected."""
        root = {"level": 0}
        current = root
        for i in range(1, 10):
            current["child"] = {"level": i}
            current = current["child"]
        # Create cycle back to root
        current["child"] = root

        result = encode_result(root, "test", "test", "test")

        # Navigate to the deepest level and verify cycle is broken with ref
        current_result = result
        for i in range(10):
            current_result = current_result["child"]

        assert current_result.get("__type__") == "ref"


class TestRecursionErrorFallback:
    """Test RecursionError fallback behavior for deeply nested structures."""

    def test_deeply_nested_list_returns_ref_on_recursion_error(self):
        """Deeply nested structure that would cause RecursionError returns ref."""
        # Create a deeply nested structure that exceeds Python's recursion limit
        # Python default recursion limit is ~1000, so 2000 levels should trigger it
        import sys
        current_limit = sys.getrecursionlimit()

        # Create nested structure deeper than recursion limit
        nested = [1]
        current = nested
        for _ in range(current_limit + 100):
            new_list = [1]
            current.append(new_list)
            current = new_list

        result = encode_result(nested, "test-session", "test", "test")

        # Instead of raising RecursionError, should return a ref
        assert isinstance(result, dict), f"Expected dict (ref), got {type(result)}"
        assert result.get("__type__") == "ref"
        assert "id" in result
        assert result.get("session_id") == "test-session"

    def test_moderate_nesting_encodes_normally(self):
        """Moderately nested structures should encode normally without fallback."""
        # 50 levels of nesting should work fine
        nested = {"value": 1}
        current = nested
        for i in range(50):
            current["child"] = {"value": i + 2}
            current = current["child"]

        result = encode_result(nested, "test-session", "test", "test")

        # Should be a regular dict, not a ref fallback
        assert isinstance(result, dict)
        assert result.get("__type__") != "ref"
        assert result["value"] == 1
        assert "child" in result

    def test_fallback_cleans_up_partial_refs(self):
        """On fallback, refs created during partial encoding are cleaned up."""
        import sys

        test_session = "cleanup-test-session"

        # Count refs for our specific session before
        def count_session_refs():
            prefix = _registry_key_prefix(test_session)
            return sum(1 for k in _instance_registry.keys() if k.startswith(prefix))

        initial_count = count_session_refs()

        # Create a structure with a custom object that will create a ref,
        # followed by deep nesting that will trigger RecursionError
        obj = CustomObject(value=42)
        current_limit = sys.getrecursionlimit()

        # Structure: [custom_object, deeply_nested_list]
        # The custom_object will create a ref, then deep nesting will fail
        deeply_nested = [1]
        current = deeply_nested
        for _ in range(current_limit + 100):
            new_list = [1]
            current.append(new_list)
            current = new_list

        test_structure = [obj, deeply_nested]

        # This should fall back to wrapping the whole structure
        result = encode_result(test_structure, test_session, "test", "test")

        # Should have fallen back to a single ref for the whole structure
        assert result.get("__type__") == "ref"

        # Count refs for our session after - should be exactly 1 (the fallback ref)
        # The partial ref from obj should have been cleaned up
        final_count = count_session_refs()
        refs_created = final_count - initial_count

        # Should only have 1 new ref (the fallback), not 2 (obj ref was cleaned)
        assert refs_created == 1, f"Expected 1 ref (fallback only), got {refs_created}"


class TestAsyncGeneratorHandling:
    """Test that async generators become refs, not stream_refs."""

    def test_async_generator_becomes_ref(self):
        """Async generator should become ref, not stream_ref (cannot be consumed via next)."""
        import sys
        import asyncio
        if sys.version_info < (3, 6):
            # Async generators not available before Python 3.6
            return

        # Create an async generator
        async def async_gen():
            for i in range(3):
                yield i

        gen = async_gen()

        result = encode_result(gen, "test-session", "test", "test")

        # Should be a ref, NOT a stream_ref
        assert isinstance(result, dict)
        assert result.get("__type__") == "ref", f"Expected ref, got {result.get('__type__')}"
        assert "id" in result
        assert result.get("session_id") == "test-session"
        assert result.get("type_name") == "async_generator"

        # Cleanup - aclose() returns a coroutine that must be awaited
        asyncio.run(gen.aclose())

    def test_sync_generator_still_becomes_stream_ref(self):
        """Sync generators should still become stream_refs."""
        gen = (x for x in range(3))

        result = encode_result(gen, "test-session", "test", "test")

        # Should be a stream_ref
        assert isinstance(result, dict)
        assert result.get("__type__") == "stream_ref"

    def test_async_generator_nested_in_dict_becomes_ref(self):
        """Async generator nested in dict should become ref, dict preserved."""
        import sys
        import asyncio
        if sys.version_info < (3, 6):
            return

        async def async_gen():
            for i in range(3):
                yield i

        gen = async_gen()

        result = encode_result({"stream": gen, "ok": 1}, "test-session", "test", "test")

        # Dict should be preserved
        assert isinstance(result, dict)
        assert result.get("__type__") != "ref"  # Dict itself is not a ref
        assert result["ok"] == 1

        # Async generator should be a ref (not stream_ref)
        stream_result = result["stream"]
        assert isinstance(stream_result, dict)
        assert stream_result.get("__type__") == "ref"

        # Cleanup - aclose() returns a coroutine that must be awaited
        asyncio.run(gen.aclose())


class TestRefMemoization:
    """Test that repeated objects yield the same ref payload (deduplication)."""

    def test_same_object_yields_same_ref_id(self):
        """Same object appearing multiple times should yield same ref id."""
        obj = CustomObject(value=42)
        result = encode_result([obj, "separator", obj], "test-session", "test", "test")

        assert isinstance(result, list)
        assert len(result) == 3
        assert result[1] == "separator"

        # Both refs should have the same id since it's the same object
        ref1 = result[0]
        ref2 = result[2]
        assert ref1.get("__type__") == "ref"
        assert ref2.get("__type__") == "ref"
        assert ref1["id"] == ref2["id"], "Same object should yield same ref id"

    def test_different_objects_yield_different_ref_ids(self):
        """Different objects should yield different ref ids."""
        obj1 = CustomObject(value=1)
        obj2 = CustomObject(value=2)
        result = encode_result([obj1, obj2], "test-session", "test", "test")

        ref1 = result[0]
        ref2 = result[1]
        assert ref1["id"] != ref2["id"], "Different objects should have different ref ids"

    def test_same_object_in_nested_structure_yields_same_ref(self):
        """Same object appearing at different nesting levels should yield same ref."""
        obj = CustomObject(value=99)
        nested = {
            "shallow": obj,
            "deep": {
                "deeper": {
                    "deepest": obj
                }
            }
        }

        result = encode_result(nested, "test-session", "test", "test")

        shallow_ref = result["shallow"]
        deepest_ref = result["deep"]["deeper"]["deepest"]

        assert shallow_ref.get("__type__") == "ref"
        assert deepest_ref.get("__type__") == "ref"
        assert shallow_ref["id"] == deepest_ref["id"]

    def test_generator_memoization(self):
        """Same generator appearing multiple times should yield same stream_ref id."""
        gen = (x for x in range(3))
        # Put the same generator object in multiple places in a dict
        # This tests that ref_memo deduplicates stream_refs
        result = encode_result({"a": gen, "b": gen}, "test-session", "test", "test")

        # Both should be stream_refs
        assert result["a"].get("__type__") == "stream_ref"
        assert result["b"].get("__type__") == "stream_ref"

        # Both should have the same id (memoization)
        assert result["a"]["id"] == result["b"]["id"], \
            "Same generator should yield same stream_ref id"

    def test_cycle_uses_memoization(self):
        """Cyclic structures should use memoized refs for back-references."""
        lst = [1, 2]
        lst.append(lst)  # Self-reference

        result = encode_result(lst, "test-session", "test", "test")

        # The cyclic reference should be a memoized ref
        cyclic_ref = result[2]
        assert cyclic_ref.get("__type__") == "ref"


class TestAtomEncoding:
    """Test that Atom objects encode as tagged atoms, not refs."""

    def test_atom_encodes_as_tagged_atom(self):
        """Atom should encode as a tagged atom value, not a ref."""
        atom = Atom("ok")
        result = encode_result(atom, "test-session", "test", "test")

        # Result should be a tagged atom, not a ref
        assert isinstance(result, dict)
        assert result.get("__type__") == "atom"
        assert result.get("value") == "ok"
        assert result.get("__schema__") == SCHEMA_VERSION

    def test_atom_in_list_encodes_correctly(self):
        """Atom nested in a list should encode as tagged atom."""
        atom = Atom("error")
        result = encode_result([1, atom, 3], "test-session", "test", "test")

        assert isinstance(result, list)
        assert len(result) == 3
        assert result[0] == 1
        assert result[2] == 3

        # Middle element should be tagged atom
        atom_result = result[1]
        assert isinstance(atom_result, dict)
        assert atom_result.get("__type__") == "atom"
        assert atom_result.get("value") == "error"

    def test_atom_in_dict_encodes_correctly(self):
        """Atom nested in a dict should encode as tagged atom."""
        atom = Atom("nil")
        result = encode_result({"status": atom, "count": 42}, "test", "test", "test")

        assert isinstance(result, dict)
        assert result["count"] == 42

        status_result = result["status"]
        assert isinstance(status_result, dict)
        assert status_result.get("__type__") == "atom"
        assert status_result.get("value") == "nil"

    def test_atom_is_json_safe(self):
        """Encoded atoms should be JSON-safe."""
        atom = Atom("test")
        result = encode_result(atom, "test", "test", "test")
        assert _is_json_safe(result)


# Simple test runner for when pytest is not available
def run_tests():
    """Run all tests and report results."""
    import traceback

    test_classes = [
        TestEncodeResultGraceful,
        TestEncodeResultRefMetadata,
        TestCycleDetection,
        TestRecursionErrorFallback,
        TestAsyncGeneratorHandling,
        TestRefMemoization,
        TestAtomEncoding,
    ]

    total = 0
    passed = 0
    failed = 0

    for test_class in test_classes:
        print(f"\n{test_class.__name__}:")
        instance = test_class()
        for name in dir(instance):
            if name.startswith('test_'):
                total += 1
                try:
                    getattr(instance, name)()
                    passed += 1
                    print(f"  PASS: {name}")
                except AssertionError as e:
                    failed += 1
                    print(f"  FAIL: {name}")
                    print(f"        {e}")
                except Exception as e:
                    failed += 1
                    print(f"  ERROR: {name}")
                    traceback.print_exc()

    print(f"\n{'=' * 60}")
    print(f"Results: {passed}/{total} passed, {failed} failed")
    return failed == 0


if __name__ == "__main__":
    # Try pytest first, fall back to simple runner
    try:
        import pytest
        sys.exit(pytest.main([__file__, "-v"]))
    except ImportError:
        print("pytest not available, using simple test runner\n")
        success = run_tests()
        sys.exit(0 if success else 1)