import json
import pickle
from unittest.mock import AsyncMock, MagicMock, patch

import pytest

from letta.schemas.pip_requirement import PipRequirement
from letta.schemas.sandbox_config import ModalSandboxConfig, SandboxConfig, SandboxType
from letta.schemas.tool import Tool
from letta.services.tool_sandbox.modal_sandbox_v2 import AsyncToolSandboxModalV2
from letta.services.tool_sandbox.modal_version_manager import ModalVersionManager
from sandbox.modal_executor import ModalFunctionExecutor


class TestModalFunctionExecutor:
    """Test the ModalFunctionExecutor class."""

    def test_execute_tool_dynamic_success(self):
        """Test successful execution of a simple tool."""
        tool_source = """
def add_numbers(a: int, b: int) -> int:
    return a + b
"""

        args = {"a": 5, "b": 3}
        args_pickled = pickle.dumps(args)

        result = ModalFunctionExecutor.execute_tool_dynamic(
            tool_source=tool_source,
            tool_name="add_numbers",
            args_pickled=args_pickled,
            agent_state_pickled=None,
            inject_agent_state=False,
            is_async=False,
            args_schema_code=None,
        )

        assert result["error"] is None
        assert result["result"] == 8  # Actual integer value
        assert result["agent_state"] is None

    def test_execute_tool_dynamic_with_error(self):
        """Test execution with an error."""
        tool_source = """
def divide_numbers(a: int, b: int) -> float:
    return a / b
"""

        args = {"a": 5, "b": 0}
        args_pickled = pickle.dumps(args)

        result = ModalFunctionExecutor.execute_tool_dynamic(
            tool_source=tool_source,
            tool_name="divide_numbers",
            args_pickled=args_pickled,
            agent_state_pickled=None,
            inject_agent_state=False,
            is_async=False,
            args_schema_code=None,
        )

        assert result["error"] is not None
        assert result["error"]["name"] == "ZeroDivisionError"
        assert "division by zero" in result["error"]["value"]
        assert result["result"] is None

    def test_execute_async_tool(self):
        """Test execution of an async tool."""
        tool_source = """
async def async_add(a: int, b: int) -> int:
    import asyncio
    await asyncio.sleep(0.001)
    return a + b
"""

        args = {"a": 10, "b": 20}
        args_pickled = pickle.dumps(args)

        result = ModalFunctionExecutor.execute_tool_dynamic(
            tool_source=tool_source,
            tool_name="async_add",
            args_pickled=args_pickled,
            agent_state_pickled=None,
            inject_agent_state=False,
            is_async=True,
            args_schema_code=None,
        )

        assert result["error"] is None
        assert result["result"] == 30

    def test_execute_with_stdout_capture(self):
        """Test that stdout is properly captured."""
        tool_source = """
def print_and_return(message: str) -> str:
    print(f"Processing: {message}")
    print("Done!")
    return message.upper()
"""

        args = {"message": "hello"}
        args_pickled = pickle.dumps(args)

        result = ModalFunctionExecutor.execute_tool_dynamic(
            tool_source=tool_source,
            tool_name="print_and_return",
            args_pickled=args_pickled,
            agent_state_pickled=None,
            inject_agent_state=False,
            is_async=False,
            args_schema_code=None,
        )

        assert result["error"] is None
        assert result["result"] == "HELLO"
        assert "Processing: hello" in result["stdout"]
        assert "Done!" in result["stdout"]


class TestModalVersionManager:
    """Test the Modal Version Manager."""

    @pytest.mark.asyncio
    async def test_register_and_get_deployment(self):
        """Test registering and retrieving deployments."""
        from unittest.mock import AsyncMock

        from letta.schemas.user import User

        manager = ModalVersionManager()

        # Mock the tool manager
        mock_tool = MagicMock()
        mock_tool.id = "tool-abc12345"
        mock_tool.metadata_ = {}

        manager.tool_manager.get_tool_by_id = MagicMock(return_value=mock_tool)
        manager.tool_manager.update_tool_by_id_async = AsyncMock(return_value=mock_tool)

        # Create a mock actor
        mock_actor = MagicMock(spec=User)
        mock_actor.id = "user-123"

        # Register a deployment
        mock_app = MagicMock(spec=["deploy", "stop"])
        info = await manager.register_deployment(
            tool_id="tool-abc12345",
            app_name="test-app",
            version_hash="abc123",
            app=mock_app,
            dependencies={"pandas", "numpy"},
            sandbox_config_id="config-123",
            actor=mock_actor,
        )

        assert info.app_name == "test-app"
        assert info.version_hash == "abc123"
        assert info.dependencies == {"pandas", "numpy"}

        # Retrieve the deployment
        retrieved = await manager.get_deployment("tool-abc12345", "config-123", actor=mock_actor)
        assert retrieved.app_name == info.app_name
        assert retrieved.version_hash == info.version_hash

    @pytest.mark.asyncio
    async def test_needs_redeployment(self):
        """Test checking if redeployment is needed."""
        from unittest.mock import AsyncMock

        from letta.schemas.user import User

        manager = ModalVersionManager()

        # Mock the tool manager
        mock_tool = MagicMock()
        mock_tool.id = "tool-def45678"
        mock_tool.metadata_ = {}

        manager.tool_manager.get_tool_by_id = MagicMock(return_value=mock_tool)
        manager.tool_manager.update_tool_by_id_async = AsyncMock(return_value=mock_tool)

        # Create a mock actor
        mock_actor = MagicMock(spec=User)

        # No deployment exists yet
        assert await manager.needs_redeployment("tool-def45678", "v1", "config-123", actor=mock_actor) is True

        # Register a deployment
        mock_app = MagicMock()
        await manager.register_deployment(
            tool_id="tool-def45678",
            app_name="test-app",
            version_hash="v1",
            app=mock_app,
            sandbox_config_id="config-123",
            actor=mock_actor,
        )

        # Update mock to return the registered deployment
        mock_tool.metadata_ = {
            "modal_deployments": {
                "config-123": {
                    "app_name": "test-app",
                    "version_hash": "v1",
                    "deployed_at": "2024-01-01T00:00:00",
                    "dependencies": [],
                }
            }
        }

        # Same version - no redeployment needed
        assert await manager.needs_redeployment("tool-def45678", "v1", "config-123", actor=mock_actor) is False

        # Different version - redeployment needed
        assert await manager.needs_redeployment("tool-def45678", "v2", "config-123", actor=mock_actor) is True

    @pytest.mark.skip(reason="get_deployment_stats method not implemented in ModalVersionManager")
    @pytest.mark.asyncio
    async def test_deployment_stats(self):
        """Test getting deployment statistics."""
        from unittest.mock import AsyncMock

        from letta.schemas.user import User

        manager = ModalVersionManager()

        # Mock the tool manager
        mock_tools = {}
        for i in range(3):
            tool_id = f"tool-{i:08x}"
            mock_tool = MagicMock()
            mock_tool.id = tool_id
            mock_tool.metadata_ = {}
            mock_tools[tool_id] = mock_tool

        def get_tool_by_id(tool_id, actor=None):
            return mock_tools.get(tool_id)

        manager.tool_manager.get_tool_by_id = MagicMock(side_effect=get_tool_by_id)
        manager.tool_manager.update_tool_by_id_async = AsyncMock()

        # Create a mock actor
        mock_actor = MagicMock(spec=User)

        # Register multiple deployments
        for i in range(3):
            tool_id = f"tool-{i:08x}"
            mock_app = MagicMock()
            await manager.register_deployment(
                tool_id=tool_id,
                app_name=f"app-{i}",
                version_hash=f"v{i}",
                app=mock_app,
                sandbox_config_id="config-123",
                actor=mock_actor,
            )

        stats = await manager.get_deployment_stats()

        # Note: The actual implementation may store deployments differently
        # This test assumes the stats method exists and returns expected format
        assert stats["total_deployments"] >= 0  # Adjust based on actual implementation
        assert "deployments" in stats

    @pytest.mark.skip(reason="export_state and import_state methods not implemented in ModalVersionManager")
    @pytest.mark.asyncio
    async def test_export_import_state(self):
        """Test exporting and importing deployment state."""
        from unittest.mock import AsyncMock

        from letta.schemas.user import User

        manager1 = ModalVersionManager()

        # Mock the tool manager for manager1
        mock_tools = {
            "tool-11111111": MagicMock(id="tool-11111111", metadata_={}),
            "tool-22222222": MagicMock(id="tool-22222222", metadata_={}),
        }

        def get_tool_by_id(tool_id, actor=None):
            return mock_tools.get(tool_id)

        manager1.tool_manager.get_tool_by_id = MagicMock(side_effect=get_tool_by_id)
        manager1.tool_manager.update_tool_by_id_async = AsyncMock()

        # Create a mock actor
        mock_actor = MagicMock(spec=User)

        # Register deployments
        mock_app = MagicMock()
        await manager1.register_deployment(
            tool_id="tool-11111111",
            app_name="app1",
            version_hash="v1",
            app=mock_app,
            dependencies={"dep1"},
            sandbox_config_id="config-123",
            actor=mock_actor,
        )
        await manager1.register_deployment(
            tool_id="tool-22222222",
            app_name="app2",
            version_hash="v2",
            app=mock_app,
            dependencies={"dep2", "dep3"},
            sandbox_config_id="config-123",
            actor=mock_actor,
        )

        # Export state
        state_json = await manager1.export_state()
        state = json.loads(state_json)

        # Verify exported state structure
        assert "tool-11111111" in state or "deployments" in state  # Depends on implementation

        # Import into new manager
        manager2 = ModalVersionManager()
        manager2.tool_manager.get_tool_by_id = MagicMock(side_effect=get_tool_by_id)

        await manager2.import_state(state_json)

        # Note: The actual implementation may not have export/import methods
        # This test assumes they exist or should be modified based on actual API


class TestAsyncToolSandboxModalV2:
    """Test the AsyncToolSandboxModalV2 class."""

    @pytest.fixture
    def mock_tool(self):
        """Create a mock tool for testing."""
        return Tool(
            id="tool-12345678",  # Valid tool ID format
            name="test_function",
            source_code="""
def test_function(x: int, y: int) -> int:
    '''Add two numbers together.'''
    return x + y
""",
            json_schema={
                "parameters": {
                    "properties": {
                        "x": {"type": "integer"},
                        "y": {"type": "integer"},
                    }
                }
            },
            pip_requirements=[PipRequirement(name="requests")],
        )

    @pytest.fixture
    def mock_user(self):
        """Create a mock user for testing."""
        user = MagicMock()
        user.organization_id = "test-org"
        return user

    @pytest.fixture
    def mock_sandbox_config(self):
        """Create a mock sandbox configuration."""
        modal_config = ModalSandboxConfig(
            timeout=60,
            pip_requirements=["pandas"],
        )
        config = SandboxConfig(
            id="sandbox-12345678",  # Valid sandbox ID format
            type=SandboxType.MODAL,  # Changed from sandbox_type to type
            config=modal_config.model_dump(),
        )
        return config

    def test_version_hash_calculation(self, mock_tool, mock_user, mock_sandbox_config):
        """Test that version hash is calculated correctly."""
        sandbox = AsyncToolSandboxModalV2(
            tool_name="test_function",
            args={"x": 1, "y": 2},
            user=mock_user,
            tool_id=mock_tool.id,
            tool_object=mock_tool,
            sandbox_config=mock_sandbox_config,
        )

        # Access through deployment manager
        version1 = sandbox._deployment_manager.calculate_version_hash(mock_sandbox_config)
        assert version1  # Should not be empty
        assert len(version1) == 12  # We take first 12 chars of hash

        # Same inputs should produce same hash
        version2 = sandbox._deployment_manager.calculate_version_hash(mock_sandbox_config)
        assert version1 == version2

        # Changing tool code should change hash
        mock_tool.source_code = "def test_function(x, y): return x * y"
        sandbox2 = AsyncToolSandboxModalV2(
            tool_name="test_function",
            args={"x": 1, "y": 2},
            user=mock_user,
            tool_id=mock_tool.id,
            tool_object=mock_tool,
            sandbox_config=mock_sandbox_config,
        )
        version3 = sandbox2._deployment_manager.calculate_version_hash(mock_sandbox_config)
        assert version3 != version1

        # Changing dependencies should also change hash
        mock_tool.source_code = "def test_function(x, y): return x + y"  # Reset
        mock_tool.pip_requirements = [PipRequirement(name="numpy")]
        sandbox3 = AsyncToolSandboxModalV2(
            tool_name="test_function",
            args={"x": 1, "y": 2},
            user=mock_user,
            tool_id=mock_tool.id,
            tool_object=mock_tool,
            sandbox_config=mock_sandbox_config,
        )
        version4 = sandbox3._deployment_manager.calculate_version_hash(mock_sandbox_config)
        assert version4 != version1

        # Changing sandbox config should change hash
        modal_config2 = ModalSandboxConfig(
            timeout=120,  # Different timeout
            pip_requirements=["pandas"],
        )
        config2 = SandboxConfig(
            id="sandbox-87654321",
            type=SandboxType.MODAL,
            config=modal_config2.model_dump(),
        )
        version5 = sandbox3._deployment_manager.calculate_version_hash(config2)
        assert version5 != version4

    def test_app_name_generation(self, mock_tool, mock_user):
        """Test app name generation."""
        sandbox = AsyncToolSandboxModalV2(
            tool_name="test_function",
            args={"x": 1, "y": 2},
            user=mock_user,
            tool_id=mock_tool.id,
            tool_object=mock_tool,
        )

        # App name generation is now in deployment manager and uses tool ID
        app_name = sandbox._deployment_manager._generate_app_name()
        # App name is based on tool ID truncated to 40 chars
        assert app_name == mock_tool.id[:40]

    @pytest.mark.asyncio
    async def test_run_with_mocked_modal(self, mock_tool, mock_user, mock_sandbox_config):
        """Test the run method with mocked Modal components."""
        with (
            patch("letta.services.tool_sandbox.modal_sandbox_v2.modal") as mock_modal,
            patch("letta.services.tool_sandbox.modal_deployment_manager.modal") as mock_modal2,
        ):
            # Mock Modal app
            mock_app = MagicMock()  # Use MagicMock for the app itself
            mock_app.run = MagicMock()

            # Mock the function decorator
            def mock_function_decorator(*args, **kwargs):
                def decorator(func):
                    # Create a mock that has a remote attribute
                    mock_func = MagicMock()
                    mock_func.remote = mock_remote
                    # Store the mocked function as tool_executor on the app
                    mock_app.tool_executor = mock_func
                    return mock_func

                return decorator

            mock_app.function = mock_function_decorator

            # Mock deployment
            mock_app.deploy = MagicMock()
            mock_app.deploy.aio = AsyncMock()

            # Mock the remote execution
            mock_remote = MagicMock()
            mock_remote.aio = AsyncMock(
                return_value={
                    "result": 3,  # Return actual integer, not string
                    "agent_state": None,
                    "stdout": "Executing...",
                    "stderr": "",
                    "error": None,
                }
            )

            mock_modal.App.return_value = mock_app
            mock_modal2.App.return_value = mock_app

            # Mock App.lookup.aio to handle app lookup attempts
            mock_modal.App.lookup = MagicMock()
            mock_modal.App.lookup.aio = AsyncMock(side_effect=Exception("App not found"))
            mock_modal2.App.lookup = MagicMock()
            mock_modal2.App.lookup.aio = AsyncMock(side_effect=Exception("App not found"))

            # Mock enable_output context manager
            mock_modal.enable_output = MagicMock()
            mock_modal.enable_output.return_value.__enter__ = MagicMock()
            mock_modal.enable_output.return_value.__exit__ = MagicMock()
            mock_modal2.enable_output = MagicMock()
            mock_modal2.enable_output.return_value.__enter__ = MagicMock()
            mock_modal2.enable_output.return_value.__exit__ = MagicMock()

            # Mock the SandboxConfigManager to avoid type checking issues
            with patch("letta.services.tool_sandbox.base.SandboxConfigManager") as MockSCM:
                mock_scm = MockSCM.return_value
                mock_scm.get_sandbox_env_vars_as_dict_async = AsyncMock(return_value={})

                # Create sandbox
                sandbox = AsyncToolSandboxModalV2(
                    tool_name="test_function",
                    args={"x": 1, "y": 2},
                    user=mock_user,
                    tool_id=mock_tool.id,
                    tool_object=mock_tool,
                    sandbox_config=mock_sandbox_config,
                )

                # Mock the version manager through deployment manager
                version_manager = sandbox._deployment_manager.version_manager
                if version_manager:
                    with patch.object(version_manager, "get_deployment", return_value=None):
                        with patch.object(version_manager, "register_deployment", return_value=None):
                            # Run the tool
                            result = await sandbox.run()
                else:
                    # If no version manager (use_version_tracking=False), just run
                    result = await sandbox.run()

            assert result.func_return == 3  # Check for actual integer
            assert result.status == "success"
            assert "Executing..." in result.stdout[0]

    def test_detect_async_function(self, mock_user):
        """Test detection of async functions."""
        # Test with sync function
        sync_tool = Tool(
            id="tool-abcdef12",  # Valid tool ID format
            name="sync_func",
            source_code="def sync_func(x): return x",
            json_schema={"parameters": {"properties": {}}},
        )

        sandbox_sync = AsyncToolSandboxModalV2(
            tool_name="sync_func",
            args={},
            user=mock_user,
            tool_id=sync_tool.id,
            tool_object=sync_tool,
        )

        assert sandbox_sync._detect_async_function() is False

        # Test with async function
        async_tool = Tool(
            id="tool-fedcba21",  # Valid tool ID format
            name="async_func",
            source_code="async def async_func(x): return x",
            json_schema={"parameters": {"properties": {}}},
        )

        sandbox_async = AsyncToolSandboxModalV2(
            tool_name="async_func",
            args={},
            user=mock_user,
            tool_id=async_tool.id,
            tool_object=async_tool,
        )

        assert sandbox_async._detect_async_function() is True


if __name__ == "__main__":
    pytest.main([__file__, "-v"])
