"""
Tests for KivotRunner.

Note: These are unit tests for the runner logic. Integration tests
with actual Temporal workflows are in ../tests/integration/.
"""

import pytest
from unittest.mock import AsyncMock, Mock, patch
from typing import Any, Dict, List

from .runner import KivotRunner, _prepare_messages, _format_tools, _execute_local_tool, _extract_function_tool_func
from .types import AgentResult


# Mock Agent class for testing
class MockAgent:
    def __init__(
        self,
        name: str = "TestAgent",
        instructions: str = "You are helpful.",
        model: str = "gpt-4o-mini",
        tools: List[Any] = None,
    ):
        self.name = name
        self.instructions = instructions
        self.model = model
        self.tools = tools or []


class TestPrepareMessages:
    """Test _prepare_messages helper function."""

    def test_string_input(self):
        agent = MockAgent(instructions="You are helpful.")
        messages = _prepare_messages(agent, "Hello!")

        assert len(messages) == 2
        assert messages[0] == {"role": "system", "content": "You are helpful."}
        assert messages[1] == {"role": "user", "content": "Hello!"}

    def test_dict_with_messages(self):
        agent = MockAgent(instructions="Test")
        input_dict = {
            "messages": [
                {"role": "user", "content": "First message"},
                {"role": "assistant", "content": "Response"},
            ]
        }
        messages = _prepare_messages(agent, input_dict)

        assert len(messages) == 3  # system + 2 messages
        assert messages[0]["role"] == "system"
        assert messages[1]["role"] == "user"
        assert messages[2]["role"] == "assistant"

    def test_dict_with_content(self):
        agent = MockAgent(instructions="Test")
        input_dict = {"content": "Direct content"}
        messages = _prepare_messages(agent, input_dict)

        assert len(messages) == 2
        assert messages[1] == {"role": "user", "content": "Direct content"}

    def test_no_instructions(self):
        agent = MockAgent(instructions="")
        messages = _prepare_messages(agent, "Hello")

        # Should not include system message
        assert len(messages) == 1
        assert messages[0]["role"] == "user"

    def test_complex_dict_fallback(self):
        agent = MockAgent(instructions="")
        input_dict = {"key1": "value1", "key2": 42}
        messages = _prepare_messages(agent, input_dict)

        assert len(messages) == 1
        assert messages[0]["role"] == "user"
        assert "key1" in messages[0]["content"]
        assert "value1" in messages[0]["content"]


class TestFormatTools:
    """Test _format_tools helper function."""

    def test_empty_tools(self):
        tools = _format_tools([])
        assert tools == []

    def test_dict_tools(self):
        tool_dict = {
            "type": "function",
            "function": {"name": "test", "description": "Test tool"},
        }
        tools = _format_tools([tool_dict])

        assert len(tools) == 1
        assert tools[0] == tool_dict

    def test_callable_tool(self):
        def get_weather(city: str) -> str:
            """Get weather for a city."""
            return f"Weather in {city}"

        tools = _format_tools([get_weather])

        assert len(tools) == 1
        tool = tools[0]
        assert tool["type"] == "function"
        assert tool["function"]["name"] == "get_weather"
        assert tool["function"]["description"] == "Get weather for a city."
        assert "parameters" in tool["function"]
        assert "city" in tool["function"]["parameters"]["properties"]

    def test_tool_with_to_dict(self):
        mock_tool = Mock()
        mock_tool.to_dict.return_value = {"type": "function", "name": "mock"}

        tools = _format_tools([mock_tool])

        assert len(tools) == 1
        assert tools[0] == {"type": "function", "name": "mock"}
        mock_tool.to_dict.assert_called_once()

    def test_callable_tool_required_params(self):
        """Verify required params only include those without defaults."""
        def search(query: str, limit: int = 10, offset: int = 0) -> str:
            """Search for items."""
            return ""

        tools = _format_tools([search])
        assert len(tools) == 1
        func_def = tools[0]["function"]
        assert func_def["name"] == "search"
        # Only 'query' should be required (limit and offset have defaults)
        assert func_def["parameters"]["required"] == ["query"]
        assert "query" in func_def["parameters"]["properties"]
        assert "limit" in func_def["parameters"]["properties"]
        assert func_def["parameters"]["properties"]["query"]["type"] == "string"
        assert func_def["parameters"]["properties"]["limit"]["type"] == "integer"

    def test_callable_tool_type_mapping(self):
        """Verify Python types map correctly to JSON schema types."""
        def typed_func(name: str, count: int, ratio: float, flag: bool, items: list, meta: dict):
            """Function with many types."""
            pass

        tools = _format_tools([typed_func])
        props = tools[0]["function"]["parameters"]["properties"]
        assert props["name"]["type"] == "string"
        assert props["count"]["type"] == "integer"
        assert props["ratio"]["type"] == "number"
        assert props["flag"]["type"] == "boolean"
        assert props["items"]["type"] == "array"
        assert props["meta"]["type"] == "object"

    def test_mixed_tool_formats(self):
        def test_func():
            """Test function."""
            pass

        tool_dict = {"type": "function", "function": {"name": "dict_tool"}}

        tools = _format_tools([tool_dict, test_func])

        assert len(tools) == 2
        assert tools[0] == tool_dict
        assert tools[1]["function"]["name"] == "test_func"

    def test_function_tool_object(self):
        """FunctionTool from OpenAI Agents SDK (has name + params_json_schema)."""

        class FakeFunctionTool:
            """Mimics agents.tool.FunctionTool — not callable, no to_dict."""

            def __init__(self):
                self.name = "greet"
                self.description = "Greet a person by name"
                self.params_json_schema = {
                    "type": "object",
                    "properties": {"name": {"type": "string"}},
                    "required": ["name"],
                }

        tool = FakeFunctionTool()
        tools = _format_tools([tool])

        assert len(tools) == 1
        assert tools[0]["type"] == "function"
        assert tools[0]["function"]["name"] == "greet"
        assert tools[0]["function"]["description"] == "Greet a person by name"
        assert tools[0]["function"]["parameters"]["type"] == "object"
        assert "name" in tools[0]["function"]["parameters"]["properties"]

    def test_function_tool_no_description(self):
        """FunctionTool with no description defaults to empty string."""

        class FakeFunctionTool:
            def __init__(self):
                self.name = "noop"
                self.params_json_schema = {"type": "object", "properties": {}}

        tool = FakeFunctionTool()
        tools = _format_tools([tool])

        assert len(tools) == 1
        assert tools[0]["function"]["name"] == "noop"
        assert tools[0]["function"]["description"] == ""

    def test_function_tool_mixed_with_callable(self):
        """Mix of FunctionTool objects and plain callables."""

        class FakeFunctionTool:
            def __init__(self):
                self.name = "greet"
                self.description = "Greet"
                self.params_json_schema = {"type": "object", "properties": {}}

        def add(a: int, b: int) -> int:
            """Add two numbers."""
            return a + b

        tools = _format_tools([FakeFunctionTool(), add])
        assert len(tools) == 2
        assert tools[0]["function"]["name"] == "greet"
        assert tools[1]["function"]["name"] == "add"


class TestKivotRunnerOutsideWorkflow:
    """Test KivotRunner behavior outside Temporal workflow."""

    @pytest.mark.asyncio
    @patch("kivot_agents.runner._HAS_TEMPORAL", False)
    @pytest.mark.skip(reason="Requires 'agents' module and real OpenAI API key")
    async def test_fallback_to_official_runner(self):
        """Should use official Runner when outside workflow."""
        # This test requires the actual 'agents' module and API key
        # Skip in unit tests, covered by integration tests
        pass

    @pytest.mark.asyncio
    @patch("kivot_agents.runner._HAS_TEMPORAL", False)
    async def test_error_without_official_runner(self):
        """Should raise error when official Runner not available."""
        agent = MockAgent()

        # Mock the agents import to fail inside _run_standalone
        import builtins
        original_import = builtins.__import__

        def mock_import(name, *args, **kwargs):
            if name == "agents":
                raise ImportError("No module named 'agents'")
            return original_import(name, *args, **kwargs)

        with patch.object(builtins, "__import__", side_effect=mock_import):
            with pytest.raises(RuntimeError, match="OpenAI Agents SDK"):
                await KivotRunner.run(agent, input="Test")


class TestKivotRunnerInsideWorkflow:
    """Test KivotRunner behavior inside Temporal workflow."""

    @pytest.mark.asyncio
    @patch("kivot_agents.runner.workflow.in_workflow", return_value=True)
    @patch("kivot_agents.runner.workflow.info")
    @patch("kivot_agents.runner.workflow.execute_activity")
    @patch("kivot_agents.runner.workflow.logger")
    async def test_simple_completion_without_tools(
        self, mock_logger, mock_execute, mock_info, mock_in_workflow
    ):
        """Test agent completes without tool calls."""
        # Setup
        mock_info.return_value.workflow_id = "test-run-123"

        # Mock ModelStepActivity response (no tools) - must be awaitable
        async def mock_activity(*args, **kwargs):
            return {
                "content": "Hello! How can I help you?",
                "tool_calls": None,
                "usage": {"prompt_tokens": 10, "completion_tokens": 8, "total_tokens": 18},
            }

        mock_execute.side_effect = mock_activity

        agent = MockAgent()
        result = await KivotRunner.run(agent, input="Hello!")

        # Assertions
        assert result.final_output == "Hello! How can I help you?"
        assert len(result.messages) == 3  # system + user + assistant
        assert result.usage["total_tokens"] == 18

        # Verify ModelStepActivity was called
        mock_execute.assert_called_once()
        call_args = mock_execute.call_args[0]
        assert call_args[0] == "ModelStepActivity"

    @pytest.mark.asyncio
    @patch("kivot_agents.runner.workflow.in_workflow", return_value=True)
    @patch("kivot_agents.runner.workflow.info")
    @patch("kivot_agents.runner.workflow.execute_activity")
    @patch("kivot_agents.runner.workflow.logger")
    async def test_agent_with_tool_calls(
        self, mock_logger, mock_execute, mock_info, mock_in_workflow
    ):
        """Test agent makes tool calls and continues."""
        # Setup
        mock_info.return_value.workflow_id = "test-run-123"

        # Mock responses: first with tool call, then final response
        responses = [
            # First call - agent requests tool
            {
                "content": "",
                "tool_calls": [
                    {
                        "id": "call_123",
                        "name": "get_weather",
                        "arguments": '{"city": "San Francisco"}',
                    }
                ],
                "usage": {"total_tokens": 25},
            },
            # Second call - agent responds with tool result
            {
                "content": "The weather in San Francisco is sunny, 72°F.",
                "tool_calls": None,
                "usage": {"total_tokens": 15},
            },
        ]

        # Track calls to distinguish ModelStepActivity vs ToolStepActivity
        call_count = {"model": 0, "tool": 0}

        async def mock_activity_call(activity_name, *args, **kwargs):
            if activity_name == "ModelStepActivity":
                idx = call_count["model"]
                call_count["model"] += 1
                return responses[idx]
            elif activity_name == "ToolStepActivity":
                call_count["tool"] += 1
                return {"output": "Sunny, 72°F", "status": "completed"}

        mock_execute.side_effect = mock_activity_call

        agent = MockAgent(tools=[{"type": "function", "function": {"name": "get_weather"}}])
        result = await KivotRunner.run(agent, input="What's the weather in SF?")

        # Assertions
        assert "sunny" in result.final_output.lower()
        assert call_count["model"] == 2  # Two model calls
        assert call_count["tool"] == 1  # One tool call
        assert len(result.messages) >= 4  # system + user + assistant + tool + assistant

    @pytest.mark.asyncio
    @patch("kivot_agents.runner.workflow.in_workflow", return_value=True)
    @patch("kivot_agents.runner.workflow.info")
    @patch("kivot_agents.runner.workflow.execute_activity")
    @patch("kivot_agents.runner.workflow.logger")
    async def test_max_iterations_limit(
        self, mock_logger, mock_execute, mock_info, mock_in_workflow
    ):
        """Test agent stops after max iterations."""
        # Setup
        mock_info.return_value.workflow_id = "test-run-123"

        # Mock response that keeps requesting tools (infinite loop scenario)
        async def mock_activity_call(activity_name, *args, **kwargs):
            if activity_name == "ModelStepActivity":
                return {
                    "content": "",
                    "tool_calls": [{"id": "call_123", "name": "test", "arguments": "{}"}],
                    "usage": None,
                }
            else:  # ToolStepActivity
                return {"output": "result", "status": "completed"}

        mock_execute.side_effect = mock_activity_call

        agent = MockAgent()

        # Should raise after max_iterations
        with pytest.raises(RuntimeError, match="exceeded max iterations"):
            await KivotRunner.run(agent, input="Loop forever", max_iterations=3)


class TestKivotRunnerWithThread:
    """Test KivotRunner behavior with thread_id for multi-turn conversations."""

    @pytest.mark.asyncio
    @patch("kivot_agents.runner.workflow.in_workflow", return_value=True)
    @patch("kivot_agents.runner.workflow.info")
    @patch("kivot_agents.runner.workflow.execute_activity")
    @patch("kivot_agents.runner.workflow.logger")
    async def test_run_with_thread_loads_history(
        self, mock_logger, mock_execute, mock_info, mock_in_workflow
    ):
        """When thread_id is provided, GetThreadMessagesActivity should be called first."""
        mock_info.return_value.workflow_id = "test-run-123"

        call_order = []

        async def mock_activity_call(activity_name, *args, **kwargs):
            call_order.append(activity_name)
            if activity_name == "GetThreadMessagesActivity":
                return {
                    "thread_id": "thread-1",
                    "messages": [
                        {"role": "user", "content": "Previous question"},
                        {"role": "assistant", "content": "Previous answer"},
                    ],
                    "count": 2,
                }
            elif activity_name == "ModelStepActivity":
                return {
                    "content": "Based on our conversation, here is my answer.",
                    "tool_calls": None,
                    "usage": {"total_tokens": 20},
                }
            elif activity_name == "AppendThreadMessagesActivity":
                return {"count": 2}

        mock_execute.side_effect = mock_activity_call

        agent = MockAgent()
        result = await KivotRunner.run(agent, input="Follow up question", thread_id="thread-1")

        # Verify call order: GetThreadMessages -> ModelStep -> AppendThreadMessages
        assert call_order[0] == "GetThreadMessagesActivity"
        assert call_order[1] == "ModelStepActivity"
        assert call_order[2] == "AppendThreadMessagesActivity"
        assert result.final_output == "Based on our conversation, here is my answer."

    @pytest.mark.asyncio
    @patch("kivot_agents.runner.workflow.in_workflow", return_value=True)
    @patch("kivot_agents.runner.workflow.info")
    @patch("kivot_agents.runner.workflow.execute_activity")
    @patch("kivot_agents.runner.workflow.logger")
    async def test_run_with_thread_saves_messages(
        self, mock_logger, mock_execute, mock_info, mock_in_workflow
    ):
        """After completion, new messages should be saved to thread."""
        mock_info.return_value.workflow_id = "test-run-456"

        saved_messages = []

        async def mock_activity_call(activity_name, *args, **kwargs):
            if activity_name == "GetThreadMessagesActivity":
                return {"thread_id": "thread-1", "messages": [], "count": 0}
            elif activity_name == "ModelStepActivity":
                return {
                    "content": "Hello!",
                    "tool_calls": None,
                    "usage": None,
                }
            elif activity_name == "AppendThreadMessagesActivity":
                saved_messages.extend(args[0]["messages"] if args else kwargs.get("arg", {}).get("messages", []))
                # Capture from the positional arg (activity input dict)
                input_dict = args[0] if args else {}
                saved_messages.clear()
                saved_messages.extend(input_dict.get("messages", []))
                return {"count": len(input_dict.get("messages", []))}

        mock_execute.side_effect = mock_activity_call

        agent = MockAgent()
        result = await KivotRunner.run(agent, input="Hi", thread_id="thread-1")

        assert result.final_output == "Hello!"
        # New messages should include user input + assistant response
        assert len(saved_messages) == 2
        assert saved_messages[0]["role"] == "user"
        assert saved_messages[0]["content"] == "Hi"
        assert saved_messages[1]["role"] == "assistant"
        assert saved_messages[1]["content"] == "Hello!"

    @pytest.mark.asyncio
    @patch("kivot_agents.runner.workflow.in_workflow", return_value=True)
    @patch("kivot_agents.runner.workflow.info")
    @patch("kivot_agents.runner.workflow.execute_activity")
    @patch("kivot_agents.runner.workflow.logger")
    async def test_run_without_thread_no_session_calls(
        self, mock_logger, mock_execute, mock_info, mock_in_workflow
    ):
        """When thread_id is None, no thread activities should be called."""
        mock_info.return_value.workflow_id = "test-run-789"

        called_activities = []

        async def mock_activity_call(activity_name, *args, **kwargs):
            called_activities.append(activity_name)
            if activity_name == "ModelStepActivity":
                return {"content": "Hi!", "tool_calls": None, "usage": None}
            return {}

        mock_execute.side_effect = mock_activity_call

        agent = MockAgent()
        await KivotRunner.run(agent, input="Hello")

        # Only ModelStepActivity should be called, no thread activities
        assert "GetThreadMessagesActivity" not in called_activities
        assert "AppendThreadMessagesActivity" not in called_activities
        assert "CreateThreadActivity" not in called_activities
        assert "ModelStepActivity" in called_activities

    @pytest.mark.asyncio
    @patch("kivot_agents.runner.workflow.in_workflow", return_value=True)
    @patch("kivot_agents.runner.workflow.info")
    @patch("kivot_agents.runner.workflow.execute_activity")
    @patch("kivot_agents.runner.workflow.logger")
    async def test_history_inserted_after_system_before_input(
        self, mock_logger, mock_execute, mock_info, mock_in_workflow
    ):
        """History messages should be inserted between system message and user input."""
        mock_info.return_value.workflow_id = "test-run-ordering"

        captured_messages = []

        async def mock_activity_call(activity_name, *args, **kwargs):
            if activity_name == "GetThreadMessagesActivity":
                return {
                    "thread_id": "thread-1",
                    "messages": [
                        {"role": "user", "content": "Old question"},
                        {"role": "assistant", "content": "Old answer"},
                    ],
                    "count": 2,
                }
            elif activity_name == "ModelStepActivity":
                # Capture messages sent to model to verify ordering
                input_dict = args[0] if args else {}
                captured_messages.extend(input_dict.get("messages", []))
                return {"content": "New answer", "tool_calls": None, "usage": None}
            elif activity_name == "AppendThreadMessagesActivity":
                return {"count": 2}

        mock_execute.side_effect = mock_activity_call

        agent = MockAgent(instructions="You are helpful.")
        await KivotRunner.run(agent, input="New question", thread_id="thread-1")

        # Expected order: [system, old_user, old_assistant, new_user]
        assert len(captured_messages) == 4
        assert captured_messages[0]["role"] == "system"
        assert captured_messages[0]["content"] == "You are helpful."
        assert captured_messages[1]["role"] == "user"
        assert captured_messages[1]["content"] == "Old question"
        assert captured_messages[2]["role"] == "assistant"
        assert captured_messages[2]["content"] == "Old answer"
        assert captured_messages[3]["role"] == "user"
        assert captured_messages[3]["content"] == "New question"

    @pytest.mark.asyncio
    @patch("kivot_agents.runner.workflow.in_workflow", return_value=True)
    @patch("kivot_agents.runner.workflow.info")
    @patch("kivot_agents.runner.workflow.execute_activity")
    @patch("kivot_agents.runner.workflow.logger")
    async def test_only_new_messages_saved(
        self, mock_logger, mock_execute, mock_info, mock_in_workflow
    ):
        """Only new messages (not history) should be saved to thread."""
        mock_info.return_value.workflow_id = "test-run-new-only"

        saved_input = {}

        async def mock_activity_call(activity_name, *args, **kwargs):
            if activity_name == "GetThreadMessagesActivity":
                return {
                    "thread_id": "thread-1",
                    "messages": [
                        {"role": "user", "content": "Old question"},
                        {"role": "assistant", "content": "Old answer"},
                    ],
                    "count": 2,
                }
            elif activity_name == "ModelStepActivity":
                return {"content": "New response", "tool_calls": None, "usage": None}
            elif activity_name == "AppendThreadMessagesActivity":
                saved_input.update(args[0] if args else {})
                return {"count": len(saved_input.get("messages", []))}

        mock_execute.side_effect = mock_activity_call

        agent = MockAgent(instructions="Be helpful.")
        await KivotRunner.run(agent, input="New question", thread_id="thread-1")

        # Only new messages should be saved (user input + assistant response)
        # History messages (Old question, Old answer) should NOT be re-saved
        saved_msgs = saved_input.get("messages", [])
        assert len(saved_msgs) == 2
        roles = [m["role"] for m in saved_msgs]
        contents = [m["content"] for m in saved_msgs]
        assert "user" in roles
        assert "assistant" in roles
        assert "New question" in contents
        assert "New response" in contents
        # Old messages must NOT be present
        assert "Old question" not in contents
        assert "Old answer" not in contents


class TestExecuteLocalTool:
    """Test _execute_local_tool helper function."""

    @pytest.mark.asyncio
    async def test_plain_callable(self):
        """Test calling a plain Python function."""
        def greet(name: str) -> str:
            return f"Hello, {name}!"

        result = await _execute_local_tool(greet, '{"name": "Alice"}')
        assert result == "Hello, Alice!"

    @pytest.mark.asyncio
    async def test_async_callable(self):
        """Test calling an async function."""
        async def greet(name: str) -> str:
            return f"Hi, {name}!"

        result = await _execute_local_tool(greet, '{"name": "Bob"}')
        assert result == "Hi, Bob!"

    @pytest.mark.asyncio
    async def test_callable_with_multiple_args(self):
        """Test function with multiple keyword arguments."""
        def add(a: int, b: int) -> int:
            return a + b

        result = await _execute_local_tool(add, '{"a": 3, "b": 5}')
        assert result == "8"

    @pytest.mark.asyncio
    async def test_empty_arguments(self):
        """Test function with no arguments."""
        def hello() -> str:
            return "hello"

        result = await _execute_local_tool(hello, "{}")
        assert result == "hello"

    @pytest.mark.asyncio
    async def test_function_tool_with_closure(self):
        """Test FunctionTool-like object: extracts underlying function from closure."""
        def _real_greet(name: str) -> str:
            return f"Hello, {name}!"

        # Mimic OpenAI SDK closure chain: on_invoke_tool → impl → original func
        # The impl closure captures the original function alongside other objects.
        def _make_fake_function_tool(func):
            class FT:
                pass
            tool = FT()
            tool.name = func.__name__

            # Two-level closure matching Agents SDK structure
            sentinel = object()
            schema = object()

            def _impl():
                # Captures: some_bool, schema, func (the original function)
                return (sentinel, schema, func)

            def _on_invoke(ctx, input_json):
                # Captures: _impl (+ other things in real SDK)
                return _impl

            tool.on_invoke_tool = _on_invoke
            return tool

        tool = _make_fake_function_tool(_real_greet)
        # _extract_function_tool_func should find _real_greet by matching __name__
        extracted = _extract_function_tool_func(tool)
        assert extracted is _real_greet

        result = await _execute_local_tool(tool, '{"name": "Charlie"}')
        assert result == "Hello, Charlie!"

    @pytest.mark.asyncio
    async def test_function_tool_extraction_fails_gracefully(self):
        """When extraction fails, raises ValueError for non-callable tool."""
        class BrokenFunctionTool:
            name = "broken"
            def on_invoke_tool(self, ctx, input_json):
                pass
        # on_invoke_tool is a bound method, no __closure__ chain → extraction fails
        tool = BrokenFunctionTool()
        extracted = _extract_function_tool_func(tool)
        assert extracted is None
        # Should raise because extraction fails and tool is not callable
        with pytest.raises(ValueError, match="Cannot execute local tool"):
            await _execute_local_tool(tool, '{}')

    @pytest.mark.asyncio
    async def test_non_callable_raises(self):
        """Test that non-callable raises ValueError."""
        with pytest.raises(ValueError, match="Cannot execute local tool"):
            await _execute_local_tool("not a callable", '{}')


class TestLocalToolRouting:
    """Test that local tools are routed to LocalToolActivity (not ToolStepActivity)."""

    @pytest.mark.asyncio
    @patch("kivot_agents.runner.workflow.in_workflow", return_value=True)
    @patch("kivot_agents.runner.workflow.info")
    @patch("kivot_agents.runner.workflow.execute_activity")
    @patch("kivot_agents.runner.workflow.logger")
    async def test_local_tool_uses_local_tool_activity(
        self, mock_logger, mock_execute, mock_info, mock_in_workflow
    ):
        """Local tool calls should go to LocalToolActivity, not ToolStepActivity."""
        mock_info.return_value.workflow_id = "test-run-local"

        called_activities = []

        async def mock_activity_call(activity_name, *args, **kwargs):
            called_activities.append(activity_name)
            if activity_name == "ModelStepActivity":
                if len([a for a in called_activities if a == "ModelStepActivity"]) == 1:
                    # First call: return tool call
                    return {
                        "content": "",
                        "tool_calls": [
                            {"id": "call_1", "name": "greet", "arguments": '{"name": "Alice"}'}
                        ],
                        "usage": {"total_tokens": 20},
                    }
                else:
                    # Second call: final response
                    return {
                        "content": "I greeted Alice for you!",
                        "tool_calls": None,
                        "usage": {"total_tokens": 15},
                    }
            elif activity_name == "LocalToolActivity":
                return {"events_emitted": 2}

        mock_execute.side_effect = mock_activity_call

        # Agent with a local greet tool (callable)
        def greet(name: str) -> str:
            """Greet a person."""
            return f"Hello, {name}!"

        agent = MockAgent(tools=[greet])
        result = await KivotRunner.run(agent, input="Greet Alice")

        # LocalToolActivity should be called, NOT ToolStepActivity
        assert "LocalToolActivity" in called_activities
        assert "ToolStepActivity" not in called_activities
        assert result.final_output == "I greeted Alice for you!"

    @pytest.mark.asyncio
    @patch("kivot_agents.runner.workflow.in_workflow", return_value=True)
    @patch("kivot_agents.runner.workflow.info")
    @patch("kivot_agents.runner.workflow.execute_activity")
    @patch("kivot_agents.runner.workflow.logger")
    async def test_mcp_tool_uses_tool_step_activity(
        self, mock_logger, mock_execute, mock_info, mock_in_workflow
    ):
        """MCP tools (dict format, not callable) should go to ToolStepActivity."""
        mock_info.return_value.workflow_id = "test-run-mcp"

        called_activities = []

        async def mock_activity_call(activity_name, *args, **kwargs):
            called_activities.append(activity_name)
            if activity_name == "ModelStepActivity":
                if len([a for a in called_activities if a == "ModelStepActivity"]) == 1:
                    return {
                        "content": "",
                        "tool_calls": [
                            {"id": "call_1", "name": "search_logs", "arguments": '{"query": "error"}'}
                        ],
                        "usage": {"total_tokens": 20},
                    }
                else:
                    return {
                        "content": "Found 3 errors.",
                        "tool_calls": None,
                        "usage": {"total_tokens": 15},
                    }
            elif activity_name == "ToolStepActivity":
                return {"output": '[{"log": "error 1"}]', "status": "completed"}

        mock_execute.side_effect = mock_activity_call

        # Agent with an MCP tool (dict format, no callable — not in local_tools)
        mcp_tool = {"type": "function", "function": {"name": "search_logs"}}
        agent = MockAgent(tools=[mcp_tool])
        result = await KivotRunner.run(agent, input="Find errors")

        # ToolStepActivity should be called, NOT LocalToolActivity
        assert "ToolStepActivity" in called_activities
        assert "LocalToolActivity" not in called_activities
        assert result.final_output == "Found 3 errors."

    @pytest.mark.asyncio
    @patch("kivot_agents.runner.workflow.in_workflow", return_value=True)
    @patch("kivot_agents.runner.workflow.info")
    @patch("kivot_agents.runner.workflow.execute_activity")
    @patch("kivot_agents.runner.workflow.logger")
    async def test_mixed_tools_routing(
        self, mock_logger, mock_execute, mock_info, mock_in_workflow
    ):
        """Agent with both local and MCP tools routes correctly."""
        mock_info.return_value.workflow_id = "test-run-mixed"

        called_activities = []

        async def mock_activity_call(activity_name, *args, **kwargs):
            called_activities.append(activity_name)
            if activity_name == "ModelStepActivity":
                model_calls = len([a for a in called_activities if a == "ModelStepActivity"])
                if model_calls == 1:
                    # First: both local and MCP tool calls
                    return {
                        "content": "",
                        "tool_calls": [
                            {"id": "call_1", "name": "greet", "arguments": '{"name": "Alice"}'},
                            {"id": "call_2", "name": "search_logs", "arguments": '{"query": "error"}'},
                        ],
                        "usage": {"total_tokens": 30},
                    }
                else:
                    return {
                        "content": "Done!",
                        "tool_calls": None,
                        "usage": {"total_tokens": 15},
                    }
            elif activity_name == "LocalToolActivity":
                return {"events_emitted": 2}
            elif activity_name == "ToolStepActivity":
                return {"output": "results", "status": "completed"}

        mock_execute.side_effect = mock_activity_call

        def greet(name: str) -> str:
            """Greet."""
            return f"Hello, {name}!"

        mcp_tool = {"type": "function", "function": {"name": "search_logs"}}
        agent = MockAgent(tools=[greet, mcp_tool])
        result = await KivotRunner.run(agent, input="Greet Alice and find errors")

        # Both activity types should be called
        assert "LocalToolActivity" in called_activities
        assert "ToolStepActivity" in called_activities
        assert result.final_output == "Done!"


class TestLocalToolsMap:
    """Test local tools map building from agent.tools."""

    def test_callable_tool_detected(self):
        """Callable tools should be in local_tools map."""
        def greet(name: str) -> str:
            return f"Hello, {name}!"

        agent = MockAgent(tools=[greet])
        local_tools = {}
        for tool in getattr(agent, "tools", []):
            tool_name = getattr(tool, "name", None) or getattr(tool, "__name__", None)
            if tool_name:
                local_tools[tool_name] = tool

        assert "greet" in local_tools
        assert local_tools["greet"] is greet

    def test_dict_tool_not_in_local(self):
        """Dict tools (MCP) should NOT be in local_tools map."""
        mcp_tool = {"type": "function", "function": {"name": "search_logs"}}

        agent = MockAgent(tools=[mcp_tool])
        local_tools = {}
        for tool in getattr(agent, "tools", []):
            tool_name = getattr(tool, "name", None) or getattr(tool, "__name__", None)
            if tool_name:
                local_tools[tool_name] = tool

        # Dict doesn't have __name__ or name attribute
        assert "search_logs" not in local_tools
        assert len(local_tools) == 0

    def test_named_object_tool(self):
        """Objects with 'name' attribute should be detected."""
        class FakeTool:
            name = "custom_tool"

        tool = FakeTool()
        agent = MockAgent(tools=[tool])
        local_tools = {}
        for t in getattr(agent, "tools", []):
            tool_name = getattr(t, "name", None) or getattr(t, "__name__", None)
            if tool_name:
                local_tools[tool_name] = t

        assert "custom_tool" in local_tools

    def test_mixed_tools(self):
        """Mix of callable and dict tools."""
        def greet(name: str) -> str:
            return ""

        def add(a: int, b: int) -> int:
            return a + b

        mcp_tool = {"type": "function", "function": {"name": "search"}}

        agent = MockAgent(tools=[greet, mcp_tool, add])
        local_tools = {}
        for t in getattr(agent, "tools", []):
            tool_name = getattr(t, "name", None) or getattr(t, "__name__", None)
            if tool_name:
                local_tools[tool_name] = t

        assert "greet" in local_tools
        assert "add" in local_tools
        assert "search" not in local_tools
        assert len(local_tools) == 2


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