#
# Copyright 2026 The Dapr Authors
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#     http://www.apache.org/licenses/LICENSE-2.0
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

# TODO(@Sicoyle): this test file is a bit of a mess, and needs to be refactored when we clean the remaining classes up.
# Right now we have to do a bunch of patching at the class-level instead of patching at the instance-level.
# In future, we should do dependency injection instead of patching at the class-level to make it easier to test.
# This applies to all areas in this file where we have with patch.object()...
from datetime import timedelta
import os
from typing import Optional
from unittest.mock import AsyncMock, Mock, patch, MagicMock

import pytest
from dapr.ext.workflow import DaprWorkflowContext

from dapr_agents.agents.durable import DurableAgent
from dapr_agents.agents.configs import (
    AgentPubSubConfig,
    AgentStateConfig,
    AgentRegistryConfig,
    AgentMemoryConfig,
    AgentExecutionConfig,
    WorkflowRetryPolicy,
)
from dapr_agents.agents.schemas import (
    AgentWorkflowMessage,
    AgentWorkflowEntry,
)
from dapr_agents.llm import OpenAIChatClient
from dapr_agents.memory import ConversationDaprStateMemory
from dapr_agents.storage.daprstores.stateservice import StateStoreService
from dapr_agents.tool.base import AgentTool
from dapr_agents.types import AgentError, DaprWorkflowStatus


# We need this otherwise these tests all fail since they require Dapr to be available.
@pytest.fixture(autouse=True)
def patch_dapr_check(monkeypatch):
    """Mock Dapr dependencies to prevent requiring a running Dapr instance."""
    from unittest.mock import Mock
    import dapr.ext.workflow as wf

    # Mock WorkflowRuntime to prevent Dapr checks
    mock_runtime = Mock(spec=wf.WorkflowRuntime)
    monkeypatch.setattr(wf, "WorkflowRuntime", lambda: mock_runtime)

    class MockRetryPolicy:
        def __init__(
            self,
            max_number_of_attempts=1,
            first_retry_interval=timedelta(seconds=1),
            max_retry_interval=timedelta(seconds=60),
            backoff_coefficient=2.0,
            retry_timeout: Optional[timedelta] = None,
        ):
            self.max_number_of_attempts = max_number_of_attempts
            self.first_retry_interval = first_retry_interval
            self.max_retry_interval = max_retry_interval
            self.backoff_coefficient = backoff_coefficient
            self.retry_timeout = retry_timeout

    monkeypatch.setattr(wf, "RetryPolicy", MockRetryPolicy)

    # Return the mock runtime for tests that need it
    yield mock_runtime


class MockDaprClient:
    """Mock DaprClient that supports context manager protocol"""

    def __init__(self, http_timeout_seconds=10):
        self.get_state = MagicMock(return_value=Mock(data=None, json=lambda: {}))
        self.save_state = MagicMock()
        self.delete_state = MagicMock()
        self.query_state = MagicMock()

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        pass

    def __call__(self, *args, **kwargs):
        return self

    def get_metadata(self):
        """Mock get_metadata that returns empty metadata."""
        from unittest.mock import MagicMock

        response = MagicMock()
        response.registered_components = []
        response.application_id = "test-app-id"
        return response


class TestDurableAgent:
    """Test cases for the DurableAgent class."""

    @pytest.fixture(autouse=True)
    def setup_env(self, monkeypatch):
        """Set up environment variables and mocks for testing."""
        os.environ["OPENAI_API_KEY"] = "test-api-key"

        # Mock DaprClient to use our context manager supporting mock
        mock_client = MockDaprClient()
        mock_client.get_state.return_value = Mock(data=None)  # Default empty state

        # Patch both the client import locations
        monkeypatch.setattr("dapr.clients.DaprClient", lambda: mock_client)
        monkeypatch.setattr(
            "dapr_agents.storage.daprstores.statestore.DaprClient", lambda: mock_client
        )

        yield
        if "OPENAI_API_KEY" in os.environ:
            del os.environ["OPENAI_API_KEY"]

    @pytest.fixture
    def mock_llm(self):
        """Create a mock LLM client."""
        mock = Mock(spec=OpenAIChatClient)
        mock.generate = AsyncMock()
        mock.prompt_template = None
        # Set the class name to avoid OpenAI validation
        mock.__class__.__name__ = "MockLLMClient"
        mock.provider = "MockOpenAIProvider"
        mock.api = "MockOpenAIAPI"
        mock.model = "gpt-4o-mock"
        return mock

    @pytest.fixture
    def mock_tool(self):
        """Create a mock tool."""
        tool = Mock(spec=AgentTool)
        tool.name = "test_tool"
        tool.description = "A test tool"
        tool.run = AsyncMock(return_value="test_result")
        tool._is_async = True
        return tool

    @pytest.fixture
    def mock_workflow_context(self):
        """Create a mock Dapr workflow context."""
        context = DaprWorkflowContext()
        context.instance_id = "test-instance-123"
        context.is_replaying = False
        context.call_activity = Mock()
        context.wait_for_external_event = Mock()
        context.set_custom_status = Mock()
        context.current_utc_datetime = Mock()
        context.current_utc_datetime.isoformat = Mock(
            return_value="2024-01-01T00:00:00.000000"
        )
        return context

    @pytest.fixture
    def basic_durable_agent(self, mock_llm):
        """Create a basic durable agent instance for testing."""
        return DurableAgent(
            name="TestDurableAgent",
            role="Test Durable Assistant",
            goal="Help with testing",
            instructions=["Be helpful", "Test things"],
            llm=mock_llm,
            pubsub=AgentPubSubConfig(
                pubsub_name="testpubsub",
                agent_topic="TestDurableAgent",
            ),
            state=AgentStateConfig(
                store=StateStoreService(store_name="teststatestore")
            ),
            registry=AgentRegistryConfig(
                store=StateStoreService(store_name="testregistry")
            ),
            memory=AgentMemoryConfig(
                store=ConversationDaprStateMemory(
                    store_name="teststatestore",
                    workflow_instance_id="test_session",
                )
            ),
            execution=AgentExecutionConfig(max_iterations=5),
        )

    @pytest.fixture
    def durable_agent_with_tools(self, mock_llm, mock_tool):
        """Create a durable agent with tools for testing."""
        return DurableAgent(
            name="ToolDurableAgent",
            role="Tool Durable Assistant",
            goal="Execute tools",
            instructions=["Use tools when needed"],
            llm=mock_llm,
            pubsub=AgentPubSubConfig(
                pubsub_name="testpubsub",
                agent_topic="ToolDurableAgent",
            ),
            state=AgentStateConfig(
                store=StateStoreService(store_name="teststatestore")
            ),
            registry=AgentRegistryConfig(
                store=StateStoreService(store_name="testregistry")
            ),
            memory=AgentMemoryConfig(
                store=ConversationDaprStateMemory(
                    store_name="teststatestore",
                    workflow_instance_id="test_session",
                )
            ),
            tools=[mock_tool],
            execution=AgentExecutionConfig(max_iterations=5),
        )

    def test_durable_agent_initialization(self, mock_llm):
        """Test durable agent initialization with basic parameters."""
        agent = DurableAgent(
            name="TestDurableAgent",
            role="Test Durable Assistant",
            goal="Help with testing",
            instructions=["Be helpful"],
            llm=mock_llm,
            pubsub=AgentPubSubConfig(
                pubsub_name="testpubsub",
                agent_topic="TestDurableAgent",
            ),
            state=AgentStateConfig(
                store=StateStoreService(store_name="teststatestore")
            ),
            registry=AgentRegistryConfig(
                store=StateStoreService(store_name="testregistry")
            ),
        )

        assert agent.name == "TestDurableAgent"
        assert agent.prompting_helper.role == "Test Durable Assistant"
        assert agent.prompting_helper.goal == "Help with testing"
        assert agent.prompting_helper.instructions == ["Be helpful"]
        assert agent.execution.max_iterations == 10  # default value
        assert agent.tool_history == []
        assert agent.pubsub.pubsub_name == "testpubsub"
        assert agent.pubsub.agent_topic == "TestDurableAgent"

    def test_durable_agent_initialization_with_custom_topic(self, mock_llm):
        """Test durable agent initialization with custom topic name."""
        agent = DurableAgent(
            name="TestDurableAgent",
            role="Test Durable Assistant",
            goal="Help with testing",
            llm=mock_llm,
            pubsub=AgentPubSubConfig(
                pubsub_name="testpubsub",
                agent_topic="custom-topic",
            ),
            state=AgentStateConfig(
                store=StateStoreService(store_name="teststatestore")
            ),
            registry=AgentRegistryConfig(
                store=StateStoreService(store_name="testregistry")
            ),
        )

        assert agent.pubsub.agent_topic == "custom-topic"

    def test_durable_agent_initialization_name_from_role(self, mock_llm):
        """Test durable agent initialization with name derived from role."""
        agent = DurableAgent(
            name="Test Durable Assistant",
            role="Test Durable Assistant",
            goal="Help with testing",
            llm=mock_llm,
            pubsub=AgentPubSubConfig(
                pubsub_name="testpubsub",
                agent_topic="Test Durable Assistant",
            ),
            state=AgentStateConfig(
                store=StateStoreService(store_name="teststatestore")
            ),
            registry=AgentRegistryConfig(
                store=StateStoreService(store_name="testregistry")
            ),
        )

        assert agent.name == "Test Durable Assistant"
        assert agent.pubsub.agent_topic == "Test Durable Assistant"

    def test_durable_agent_metadata(self, basic_durable_agent):
        """Test durable agent metadata creation."""
        metadata = basic_durable_agent.agent_metadata

        assert metadata is not None
        assert metadata.name == "TestDurableAgent"
        assert metadata.agent.role == "Test Durable Assistant"
        assert metadata.agent.goal == "Help with testing"
        assert metadata.pubsub.agent_topic == "TestDurableAgent"
        assert metadata.pubsub.resource_name == "testpubsub"
        assert metadata.agent.orchestrator is False

    def test_tool_calling_workflow_initialization(
        self, basic_durable_agent, mock_workflow_context
    ):
        """Test workflow initialization on first iteration."""
        message = {
            "task": "Test task",
            "iteration": 0,
            "workflow_instance_id": "parent-instance-123",
        }

        mock_workflow_context.instance_id = "test-instance-123"
        mock_workflow_context.call_activity.side_effect = [
            {"content": "Test response"},
            {"message": "Test response"},
            "stop",
        ]

        # Use AgentWorkflowEntry for state setup (single entry per instance)
        entry = AgentWorkflowEntry(
            source=None,
            triggering_workflow_instance_id="parent-instance-123",
            messages=[],
            tool_history=[],
        )
        basic_durable_agent._infra._state_model = entry
        with patch.object(
            basic_durable_agent._infra,
            "get_state",
            side_effect=lambda wid: basic_durable_agent._infra._state_model,
        ):
            workflow_gen = basic_durable_agent.agent_workflow(
                mock_workflow_context, message
            )
            try:
                next(workflow_gen)  # agent_workflow is a generator, not async
            except StopIteration:
                pass

        # State is the current workflow entry (single entry model)
        instance_data = basic_durable_agent._state_model
        assert instance_data.source is None
        assert instance_data.triggering_workflow_instance_id == "parent-instance-123"

    @pytest.mark.asyncio
    async def test_broadcast_to_team_activity(self, basic_durable_agent):
        """Test broadcasting message to agents activity."""
        message = {
            "type": "broadcast",
            "content": "Test broadcast message",
            "sender": "TestDurableAgent",
        }

        # Mock the activity context
        mock_ctx = Mock()

        # The basic_durable_agent fixture doesn't have a broadcast_topic configured,
        # so this should execute without error but skip the actual broadcast
        basic_durable_agent.broadcast_to_team(mock_ctx, {"message": message})
        # Test passes if no exception is raised

    @pytest.mark.asyncio
    async def test_finish_workflow_activity(self, basic_durable_agent):
        """Test finishing workflow activity."""

        instance_id = "test-instance-123"
        final_output = "Final response"
        entry = AgentWorkflowEntry(
            source="test_source",
            triggering_workflow_instance_id=None,
            messages=[],
            tool_history=[],
        )
        basic_durable_agent._infra._state_model = entry

        # Mock the activity context and save_state
        mock_ctx = Mock()

        with (
            patch.object(basic_durable_agent, "save_state"),
            patch.object(
                basic_durable_agent._infra,
                "load_state",
            ),
            patch.object(
                basic_durable_agent._infra,
                "get_state",
                side_effect=lambda wid: basic_durable_agent._infra._state_model,
            ),
        ):
            basic_durable_agent.finalize_workflow(
                mock_ctx,
                {
                    "instance_id": instance_id,
                    "final_output": final_output,
                    "end_time": "2024-01-01T00:00:00Z",
                    "triggering_workflow_instance_id": None,
                },
            )
        assert basic_durable_agent._state_model.triggering_workflow_instance_id is None

    def test_run_tool(self, basic_durable_agent, mock_tool):
        """Test that run_tool executes a tool and returns the result without persisting state."""
        from datetime import datetime

        instance_id = "test-instance-123"
        tool_call = {
            "id": "call_123",
            "function": {"name": "test_tool", "arguments": '{"arg1": "value1"}'},
        }

        # Set up state: single entry in _state_model
        entry = AgentWorkflowEntry(
            source="test_source",
            triggering_workflow_instance_id=None,
            messages=[],
            tool_history=[],
        )
        basic_durable_agent._infra._state_model = entry

        # Mock the tool executor
        with (
            patch.object(
                type(basic_durable_agent.tool_executor),
                "run_tool",
                new_callable=AsyncMock,
            ) as mock_run_tool,
            patch.object(
                basic_durable_agent._infra,
                "get_state",
                side_effect=lambda wid: basic_durable_agent._infra._state_model,
            ),
        ):
            mock_run_tool.return_value = "tool_result"

            test_time = datetime.fromisoformat(
                "2024-01-01T00:00:00Z".replace("Z", "+00:00")
            )

            # Mock the activity context
            mock_ctx = Mock()

            result = basic_durable_agent.run_tool(
                mock_ctx,
                {
                    "tool_call": tool_call,
                    "instance_id": instance_id,
                    "time": test_time.isoformat(),
                    "order": 1,
                },
            )

            # Verify tool was executed and result was returned
            assert result["tool_call_id"] == "call_123"
            assert result["name"] == "test_tool"
            assert result["content"] == "tool_result"

            # Verify that instance state was NOT modified by run_tool
            entry = basic_durable_agent._state_model
            assert len(entry.messages) == 0  # No tool message added by run_tool
            assert len(entry.tool_history) == 0  # No tool history added by run_tool

    def test_run_tool_unwraps_kwargs_for_mcp_tools(
        self, basic_durable_agent, mock_tool
    ):
        """When tool_call arguments are wrapped as {'kwargs': {...}}, run_tool unwraps so executor gets flat kwargs."""
        instance_id = "test-instance-123"
        # Simulate LLM returning MCP-style wrapped arguments
        tool_call = {
            "id": "call_456",
            "function": {
                "name": "test_tool",
                "arguments": '{"kwargs": {"arg1": "value1"}}',
            },
        }
        entry = AgentWorkflowEntry(
            source="test_source",
            triggering_workflow_instance_id=None,
            messages=[],
            tool_history=[],
        )
        basic_durable_agent._infra._state_model = entry
        mock_ctx = Mock()
        with (
            patch.object(
                type(basic_durable_agent.tool_executor),
                "run_tool",
                new_callable=AsyncMock,
            ) as mock_run_tool,
            patch.object(
                basic_durable_agent._infra,
                "get_state",
                side_effect=lambda wid: basic_durable_agent._state_model,
            ),
        ):
            mock_run_tool.return_value = "ok"
            basic_durable_agent.run_tool(
                mock_ctx,
                {
                    "tool_call": tool_call,
                    "instance_id": instance_id,
                    "time": "2024-01-01T00:00:00Z",
                    "order": 1,
                },
            )
            mock_run_tool.assert_called_once()
            call_kwargs = mock_run_tool.call_args[1]
            assert call_kwargs == {"arg1": "value1"}, (
                "run_tool should unwrap {'kwargs': {...}} so executor receives flat kwargs"
            )

    def test_record_initial_entry(self, basic_durable_agent):
        """Test record_initial_entry helper method."""

        instance_id = "test-instance-123"
        source = "test_source"
        triggering_workflow_instance_id = "parent-instance-123"

        entry = AgentWorkflowEntry(
            source=None,
            triggering_workflow_instance_id=None,
            messages=[],
            tool_history=[],
        )
        basic_durable_agent._infra._state_model = entry

        # Mock the activity context (record_initial_entry uses ctx.workflow_id)
        mock_ctx = Mock()
        mock_ctx.workflow_id = instance_id

        with (
            patch.object(basic_durable_agent, "save_state"),
            patch.object(
                basic_durable_agent._infra,
                "get_state",
                side_effect=lambda wid: basic_durable_agent._infra._state_model,
            ),
        ):
            basic_durable_agent.record_initial_entry(
                mock_ctx,
                {
                    "instance_id": instance_id,
                    "source": source,
                    "triggering_workflow_instance_id": triggering_workflow_instance_id,
                    "trace_context": None,
                },
            )

        # Verify instance was updated
        assert basic_durable_agent._state_model.source == source
        assert (
            basic_durable_agent._state_model.triggering_workflow_instance_id
            == triggering_workflow_instance_id
        )

    def test_process_user_message(self, basic_durable_agent):
        """Test _process_user_message helper method."""

        instance_id = "test-instance-123"
        task = "Hello, world!"
        user_message_copy = {"role": "user", "content": "Hello, world!"}

        entry = AgentWorkflowEntry(
            source="test_source",
            triggering_workflow_instance_id=None,
            messages=[],
            tool_history=[],
        )
        basic_durable_agent._infra._state_model = entry

        # Mock memory.add_message and save_state
        with (
            patch.object(type(basic_durable_agent.memory), "add_message"),
            patch.object(basic_durable_agent, "save_state"),
            patch.object(
                basic_durable_agent._infra,
                "get_state",
                side_effect=lambda wid: basic_durable_agent._infra._state_model,
            ),
        ):
            basic_durable_agent._process_user_message(
                instance_id, task, user_message_copy
            )

        # Verify message was added to instance
        entry = basic_durable_agent._state_model
        assert len(entry.messages) == 1
        assert entry.messages[0].role == "user"
        assert entry.messages[0].content == "Hello, world!"
        assert entry.last_message.role == "user"

    def test_save_assistant_message(self, basic_durable_agent):
        """Test _save_assistant_message helper method."""

        instance_id = "test-instance-123"
        assistant_message = {"role": "assistant", "content": "Hello back!"}

        entry = AgentWorkflowEntry(
            source="test_source",
            triggering_workflow_instance_id=None,
            messages=[],
            tool_history=[],
        )
        basic_durable_agent._infra._state_model = entry

        # Mock memory.add_message and save_state
        with (
            patch.object(type(basic_durable_agent.memory), "add_message"),
            patch.object(basic_durable_agent, "save_state"),
            patch.object(
                basic_durable_agent._infra,
                "get_state",
                side_effect=lambda wid: basic_durable_agent._infra._state_model,
            ),
        ):
            basic_durable_agent._save_assistant_message(instance_id, assistant_message)

        # Verify message was added to instance
        entry = basic_durable_agent._state_model
        assert len(entry.messages) == 1
        assert entry.messages[0].role == "assistant"
        assert entry.messages[0].content == "Hello back!"
        assert entry.last_message.role == "assistant"

    def test_get_last_message_from_state(self, basic_durable_agent):
        """Test accessing last_message from instance state."""

        last_msg = AgentWorkflowMessage(role="assistant", content="Last message")
        entry = AgentWorkflowEntry(
            source="test_source",
            triggering_workflow_instance_id=None,
            messages=[],
            tool_history=[],
            last_message=last_msg,
        )
        basic_durable_agent._infra._state_model = entry

        # Access last_message directly from the current state entry
        assert basic_durable_agent._state_model.last_message is not None
        assert basic_durable_agent._state_model.last_message.role == "assistant"
        assert basic_durable_agent._state_model.last_message.content == "Last message"

    def test_create_tool_message_objects(self, basic_durable_agent):
        """Test that tool message dicts are created correctly by run_tool and persisted by save_tool_results."""
        from datetime import datetime, timezone

        instance_id = "test-instance-123"
        tool_call = {
            "id": "call_123",
            "function": {"name": "test_tool", "arguments": '{"arg1": "value1"}'},
        }

        # Create entry with an assistant message with tool_calls
        # Tool messages must follow an assistant message with tool_calls (OpenAI API requirement)
        assistant_message = AgentWorkflowMessage(
            role="assistant",
            content="I'll help you",
            tool_calls=[
                {
                    "id": "call_123",
                    "type": "function",
                    "function": {"name": "test_tool", "arguments": "{}"},
                }
            ],
        )
        entry = AgentWorkflowEntry(
            source="test_source",
            triggering_workflow_instance_id=None,
            messages=[assistant_message],
            tool_history=[],
            last_message=assistant_message,  # Set last_message to the assistant message
        )
        basic_durable_agent._infra._state_model = entry

        # Mock tool executor
        with patch.object(
            type(basic_durable_agent.tool_executor), "run_tool", new_callable=AsyncMock
        ) as mock_run:
            mock_run.return_value = "tool_result"

            mock_ctx = Mock()

            result = basic_durable_agent.run_tool(
                mock_ctx,
                {
                    "tool_call": tool_call,
                    "instance_id": instance_id,
                    "time": datetime.now(timezone.utc).isoformat(),
                    "order": 1,
                },
            )

        # Verify the tool result structure
        assert result["tool_call_id"] == "call_123"
        assert result["name"] == "test_tool"
        assert result["content"] == "tool_result"

        # Now call save_tool_results to persist the results
        with (
            patch.object(basic_durable_agent, "save_state"),
            patch.object(basic_durable_agent._infra, "load_state"),
            patch.object(
                basic_durable_agent._infra,
                "get_state",
                side_effect=lambda wid: basic_durable_agent._infra._state_model,
            ),
        ):
            basic_durable_agent.save_tool_results(
                mock_ctx,
                {
                    "tool_results": [result],
                    "instance_id": instance_id,
                },
            )

        # Verify messages were added to instance by save_tool_results
        entry = basic_durable_agent._state_model
        assert len(entry.messages) == 2  # Assistant message + tool message
        # Find the tool message (last one should be the tool message)
        tool_messages = [m for m in entry.messages if m.role == "tool"]
        assert len(tool_messages) == 1
        assert (
            tool_messages[0].tool_call_id == "call_123"
        )  # Check tool_call_id, not the message UUID id
        assert tool_messages[0].name == "test_tool"

    def test_append_tool_message_to_instance(self, basic_durable_agent):
        """Test that tool messages are appended to instance via save_tool_results activity."""
        instance_id = "test-instance-123"

        # Create entry with an assistant message with tool_calls
        # Tool messages must follow an assistant message with tool_calls (OpenAI API requirement)
        assistant_message = AgentWorkflowMessage(
            role="assistant",
            content="I'll help you",
            tool_calls=[
                {
                    "id": "call_123",
                    "type": "function",
                    "function": {"name": "TestToolFunc", "arguments": "{}"},
                }
            ],
        )
        entry = AgentWorkflowEntry(
            source="test_source",
            triggering_workflow_instance_id=None,
            messages=[assistant_message],
            tool_history=[],
            last_message=assistant_message,  # Set last_message to the assistant message
        )
        basic_durable_agent._infra._state_model = entry

        # Create a simple test tool
        from dapr_agents.tool.base import AgentTool

        def test_tool_func(x):
            """Test tool for verification."""
            return "tool_result"

        test_tool = AgentTool.from_func(test_tool_func)
        basic_durable_agent.tools.append(test_tool)
        # Recreate tool executor with the new tool
        from dapr_agents.tool.executor import AgentToolExecutor

        basic_durable_agent.tool_executor = AgentToolExecutor(
            tools=list(basic_durable_agent.tools)
        )

        mock_ctx = Mock()

        # Call run_tool activity which executes the tool and returns result dict
        result = basic_durable_agent.run_tool(
            mock_ctx,
            {
                "instance_id": instance_id,
                "tool_call": {
                    "id": "call_123",
                    "type": "function",
                    "function": {
                        "name": "TestToolFunc",  # Tool name is CamelCase version of function name
                        "arguments": '{"x": "test"}',  # Pass string to match type hint default
                    },
                },
            },
        )

        # run_tool does NOT persist state, so entry should still have only the assistant message
        assert len(entry.messages) == 1
        assert entry.messages[0].role == "assistant"
        assert len(entry.tool_history) == 0

        with (
            patch.object(basic_durable_agent, "save_state"),
            patch.object(basic_durable_agent._infra, "load_state"),
            patch.object(
                basic_durable_agent._infra,
                "get_state",
                side_effect=lambda wid: basic_durable_agent._infra._state_model,
            ),
        ):
            basic_durable_agent.save_tool_results(
                mock_ctx,
                {
                    "tool_results": [result],
                    "instance_id": instance_id,
                },
            )

        # Verify entry was updated with message by save_tool_results
        assert len(entry.messages) == 2  # Assistant message + tool message
        # Find the tool message (last one should be the tool message)
        tool_messages = [m for m in entry.messages if m.role == "tool"]
        assert len(tool_messages) == 1
        assert (
            tool_messages[0].tool_call_id == "call_123"
        )  # Check tool_call_id, not the message UUID id

    def test_update_agent_memory_and_history(self, basic_durable_agent):
        """Test that memory is updated via save_tool_results activity."""
        instance_id = "test-instance-123"

        # Create entry with an assistant message with tool_calls
        # Tool messages must follow an assistant message with tool_calls (OpenAI API requirement)
        assistant_message = AgentWorkflowMessage(
            role="assistant",
            content="I'll help you",
            tool_calls=[
                {
                    "id": "call_123",
                    "type": "function",
                    "function": {"name": "TestToolFunc", "arguments": "{}"},
                }
            ],
        )
        entry = AgentWorkflowEntry(
            source="test_source",
            triggering_workflow_instance_id=None,
            messages=[assistant_message],
            tool_history=[],
            last_message=assistant_message,  # Set last_message to the assistant message
        )
        basic_durable_agent._infra._state_model = entry

        # Create a simple test tool
        from dapr_agents.tool.base import AgentTool

        def test_tool_func(x: str) -> str:
            """Test tool for verification."""
            return "tool_result"

        test_tool = AgentTool.from_func(test_tool_func)
        basic_durable_agent.tools.append(test_tool)
        # Recreate tool executor with the new tool
        from dapr_agents.tool.executor import AgentToolExecutor

        basic_durable_agent.tool_executor = AgentToolExecutor(
            tools=list(basic_durable_agent.tools)
        )

        mock_ctx = Mock()

        # Call run_tool activity which executes the tool and returns result
        result = basic_durable_agent.run_tool(
            mock_ctx,
            {
                "instance_id": instance_id,
                "tool_call": {
                    "id": "call_123",
                    "type": "function",
                    "function": {
                        "name": "TestToolFunc",
                        "arguments": '{"x": "test"}',
                    },
                },
            },
        )

        # Verify run_tool returns proper result dict
        assert result["tool_call_id"] == "call_123"
        assert result["name"] == "TestToolFunc"
        assert result["content"] == "tool_result"

        # Mock save_state; save_tool_results persists tool results to state entry
        with (
            patch.object(basic_durable_agent, "save_state"),
            patch.object(basic_durable_agent._infra, "load_state"),
            patch.object(
                basic_durable_agent._infra,
                "get_state",
                side_effect=lambda wid: basic_durable_agent._infra._state_model,
            ),
        ):
            basic_durable_agent.save_tool_results(
                mock_ctx,
                {
                    "tool_results": [result],
                    "instance_id": instance_id,
                },
            )

        # Verify tool message was persisted to state entry
        entry = basic_durable_agent._state_model
        assert len(entry.messages) == 2  # Assistant message + tool message
        # Find the tool message (last one should be the tool message)
        tool_messages = [m for m in entry.messages if m.role == "tool"]
        assert len(tool_messages) == 1
        assert tool_messages[0].tool_call_id == "call_123"
        assert tool_messages[0].name == "TestToolFunc"

    def test_reconstruct_conversation_history(self, basic_durable_agent):
        """Test test_reconstruct_conversation_history helper method."""

        instance_id = "test-instance-123"

        entry = AgentWorkflowEntry(
            source="test_source",
            triggering_workflow_instance_id=None,
            messages=[
                AgentWorkflowMessage(role="user", content="Hello"),
                AgentWorkflowMessage(role="assistant", content="Hi there!"),
            ],
            tool_history=[],
        )
        basic_durable_agent._infra._state_model = entry

        with patch.object(
            basic_durable_agent._infra,
            "get_state",
            side_effect=lambda wid: basic_durable_agent._infra._state_model,
        ):
            messages = basic_durable_agent._reconstruct_conversation_history(
                instance_id
            )

        # Should include messages from instance history (system messages excluded from instance timeline)
        assert len(messages) >= 2  # At least the 2 instance messages
        # Messages may be Pydantic models or dicts
        user_messages = [
            m
            for m in messages
            if getattr(m, "role", m.get("role") if isinstance(m, dict) else None)
            == "user"
        ]
        assistant_messages = [
            m
            for m in messages
            if getattr(m, "role", m.get("role") if isinstance(m, dict) else None)
            == "assistant"
        ]
        assert len(user_messages) >= 1
        assert len(assistant_messages) >= 1

    def test_register_agentic_system(self, basic_durable_agent):
        """Test registering agentic system with per-agent key + index."""
        saved = {}
        metadata = basic_durable_agent.agent_metadata

        def fake_save(*, key, value, **kwargs):
            saved[key] = value

        def fake_load_with_etag(*, key, default=None, **kwargs):
            return (saved.get(key, default), "etag-1")

        def fake_load(*, key, default=None, **kwargs):
            return saved.get(key, default)

        with (
            patch.object(
                basic_durable_agent.registry_state, "save", side_effect=fake_save
            ),
            patch.object(
                basic_durable_agent.registry_state,
                "load_with_etag",
                side_effect=fake_load_with_etag,
            ),
            patch.object(
                basic_durable_agent.registry_state, "load", side_effect=fake_load
            ),
        ):
            basic_durable_agent.register_agentic_system(metadata=metadata)

        # Verify per-agent key was saved
        agent_key = f"agents:default:{basic_durable_agent.name}"
        assert agent_key in saved
        assert saved[agent_key] is not None

        # Verify index was updated
        index_key = "agents:default:_index"
        assert index_key in saved
        assert basic_durable_agent.name in saved[index_key]["agents"]

    def test_register_agentic_system_idempotent(self, basic_durable_agent):
        """Test that re-registering the same agent does not duplicate in the index."""
        saved = {}
        metadata = basic_durable_agent.agent_metadata

        def fake_save(*, key, value, **kwargs):
            saved[key] = value

        def fake_load_with_etag(*, key, default=None, **kwargs):
            return (saved.get(key, default), "etag-1")

        def fake_load(*, key, default=None, **kwargs):
            return saved.get(key, default)

        with (
            patch.object(
                basic_durable_agent.registry_state, "save", side_effect=fake_save
            ),
            patch.object(
                basic_durable_agent.registry_state,
                "load_with_etag",
                side_effect=fake_load_with_etag,
            ),
            patch.object(
                basic_durable_agent.registry_state, "load", side_effect=fake_load
            ),
        ):
            basic_durable_agent.register_agentic_system(metadata=metadata)
            basic_durable_agent.register_agentic_system(metadata=metadata)

        index_key = "agents:default:_index"
        agents_list = saved[index_key]["agents"]
        assert agents_list.count(basic_durable_agent.name) == 1

    def test_deregister_agentic_system(self, basic_durable_agent):
        """Test deregistering removes per-agent key and updates index."""
        saved = {
            f"agents:default:{basic_durable_agent.name}": {
                "name": basic_durable_agent.name
            },
            "agents:default:_index": {
                "agents": [basic_durable_agent.name, "other-agent"]
            },
        }
        deleted_keys = []

        def fake_save(*, key, value, **kwargs):
            saved[key] = value

        def fake_load_with_etag(*, key, default=None, **kwargs):
            return (saved.get(key, default), "etag-1")

        def fake_delete(*, key, **kwargs):
            deleted_keys.append(key)
            saved.pop(key, None)

        with (
            patch.object(
                basic_durable_agent.registry_state, "save", side_effect=fake_save
            ),
            patch.object(
                basic_durable_agent.registry_state,
                "load_with_etag",
                side_effect=fake_load_with_etag,
            ),
            patch.object(
                basic_durable_agent.registry_state, "delete", side_effect=fake_delete
            ),
        ):
            basic_durable_agent.deregister_agentic_system()

        # Per-agent key was deleted
        agent_key = f"agents:default:{basic_durable_agent.name}"
        assert agent_key in deleted_keys

        # Index no longer contains this agent but retains the other
        index_key = "agents:default:_index"
        assert basic_durable_agent.name not in saved[index_key]["agents"]
        assert "other-agent" in saved[index_key]["agents"]

    def test_get_agents_metadata_stale_index(self, basic_durable_agent):
        """Test that get_agents_metadata skips stale index entries (missing per-agent keys)."""
        stored = {
            "agents:default:_index": {"agents": ["agent-a", "agent-b", "agent-stale"]},
        }

        def fake_load(*, key, default=None, **kwargs):
            return stored.get(key, default)

        def fake_load_many(*, keys, **kwargs):
            # Only agent-a and agent-b exist; agent-stale is missing
            # Metadata structure must have "agent" key with nested metadata
            return {
                "agents:default:agent-a": {
                    "name": "agent-a",
                    "agent": {"name": "agent-a", "orchestrator": False},
                },
                "agents:default:agent-b": {
                    "name": "agent-b",
                    "agent": {"name": "agent-b", "orchestrator": False},
                },
            }

        with (
            patch.object(
                basic_durable_agent.registry_state, "load", side_effect=fake_load
            ),
            patch.object(
                basic_durable_agent.registry_state,
                "load_many",
                side_effect=fake_load_many,
            ),
        ):
            result = basic_durable_agent.get_agents_metadata(exclude_self=False)

        assert "agent-a" in result
        assert "agent-b" in result
        assert "agent-stale" not in result

    def test_get_agents_metadata_excludes_self_and_orchestrator(
        self, basic_durable_agent
    ):
        """Test that get_agents_metadata applies exclude_self and exclude_orchestrator filters."""
        stored = {
            "agents:default:_index": {
                "agents": [basic_durable_agent.name, "agent-b", "orch-agent"]
            },
        }

        def fake_load(*, key, default=None, **kwargs):
            return stored.get(key, default)

        def fake_load_many(*, keys, **kwargs):
            return {
                f"agents:default:{basic_durable_agent.name}": {
                    "name": basic_durable_agent.name,
                    "agent": {"orchestrator": False},
                },
                "agents:default:agent-b": {
                    "name": "agent-b",
                    "agent": {"orchestrator": False},
                },
                "agents:default:orch-agent": {
                    "name": "orch-agent",
                    "agent": {"orchestrator": True},
                },
            }

        with (
            patch.object(
                basic_durable_agent.registry_state, "load", side_effect=fake_load
            ),
            patch.object(
                basic_durable_agent.registry_state,
                "load_many",
                side_effect=fake_load_many,
            ),
        ):
            result = basic_durable_agent.get_agents_metadata(
                exclude_self=True, exclude_orchestrator=True
            )

        assert basic_durable_agent.name not in result
        assert "agent-b" in result
        assert "orch-agent" not in result

    def test_durable_agent_properties(self, basic_durable_agent):
        """Test durable agent properties."""
        assert basic_durable_agent.tool_executor is not None
        assert basic_durable_agent.text_formatter is not None
        assert basic_durable_agent.state is not None

    def test_durable_agent_state_initialization(self, basic_durable_agent):
        """Test that the agent state is properly initialized."""
        assert basic_durable_agent.state is not None
        # State is the current workflow entry dict (messages, tool_history, etc.)
        assert (
            "messages" in basic_durable_agent.state
            or basic_durable_agent.state is not None
        )

    def test_durable_agent_retry_policy_initialization(self, mock_llm):
        """Test that DurableAgent correctly initializes with retry policy parameters."""
        agent = DurableAgent(
            name="RetryTestAgent",
            role="Retry Test Assistant",
            llm=mock_llm,
            pubsub=AgentPubSubConfig(
                pubsub_name="testpubsub",
                agent_topic="RetryTestAgent",
            ),
            retry_policy=WorkflowRetryPolicy(
                max_attempts=5,
                initial_backoff_seconds=10,
                max_backoff_seconds=60,
                backoff_multiplier=2.0,
                retry_timeout=300,
            ),
        )

        assert agent._retry_policy is not None
        assert agent._retry_policy.max_number_of_attempts == 5
        assert agent._retry_policy.first_retry_interval.total_seconds() == 10
        assert agent._retry_policy.max_retry_interval.total_seconds() == 60
        assert agent._retry_policy.backoff_coefficient == 2.0
        assert agent._retry_policy.retry_timeout.total_seconds() == 300

    def test_durable_agent_retry_policy_defaults(self, mock_llm):
        """Test that DurableAgent uses correct default retry values."""
        agent = DurableAgent(
            name="RetryDefaultAgent",
            role="Retry Default Assistant",
            llm=mock_llm,
            pubsub=AgentPubSubConfig(
                pubsub_name="testpubsub",
                agent_topic="RetryDefaultAgent",
            ),
        )

        assert agent._retry_policy is not None
        assert agent._retry_policy.max_number_of_attempts == 3
        assert agent._retry_policy.first_retry_interval.total_seconds() == 5
        assert agent._retry_policy.max_retry_interval.total_seconds() == 30
        assert agent._retry_policy.backoff_coefficient == 1.5
        assert agent._retry_policy.retry_timeout is None

    def test_durable_agent_retry_policy_env_override(self, mock_llm, monkeypatch):
        """Test that DAPR_API_MAX_RETRIES environment variable overrides max_attempts."""
        monkeypatch.setenv("DAPR_API_MAX_RETRIES", "10")

        agent = DurableAgent(
            name="RetryEnvAgent",
            role="Retry Env Assistant",
            llm=mock_llm,
            pubsub=AgentPubSubConfig(
                pubsub_name="testpubsub",
                agent_topic="RetryEnvAgent",
            ),
            retry_policy=WorkflowRetryPolicy(max_attempts=3),
        )

        # Should use env var value over max_attempts
        assert agent._retry_policy.max_number_of_attempts == 10

    def test_durable_agent_retry_policy_invalid_env(self, mock_llm, monkeypatch):
        """Test that invalid DAPR_API_MAX_RETRIES falls back to max_attempts."""
        monkeypatch.setenv("DAPR_API_MAX_RETRIES", "invalid")

        agent = DurableAgent(
            name="RetryInvalidEnvAgent",
            role="Retry Invalid Env Assistant",
            llm=mock_llm,
            pubsub=AgentPubSubConfig(
                pubsub_name="testpubsub",
                agent_topic="RetryInvalidEnvAgent",
            ),
            retry_policy=WorkflowRetryPolicy(max_attempts=3),
        )

        # Should fall back to max_attempts since env var is invalid
        assert agent._retry_policy.max_number_of_attempts == 3

    def test_durable_agent_retry_policy_min_attempts_validation(self, mock_llm):
        """Test that max_attempts cannot be less than 1."""
        with pytest.raises(
            ValueError, match="max_attempts or DAPR_API_MAX_RETRIES must be at least 1."
        ):
            DurableAgent(
                name="RetryZeroAgent",
                role="Retry Zero Assistant",
                llm=mock_llm,
                pubsub=AgentPubSubConfig(
                    pubsub_name="testpubsub",
                    agent_topic="RetryZeroAgent",
                ),
                retry_policy=WorkflowRetryPolicy(max_attempts=0),
            )

    def test_agent_workflow_applies_retry_policy(
        self, basic_durable_agent, mock_workflow_context
    ):
        """Test that agent_workflow applies retry policy to activity calls."""
        message = {
            "task": "Test task with retries",
            "workflow_instance_id": "parent-instance-123",
        }

        call_activity_calls = []

        def track_call_activity(activity, **kwargs):
            call_activity_calls.append(
                {
                    "activity": activity,
                    "input": kwargs.get("input"),
                    "retry_policy": kwargs.get("retry_policy"),
                }
            )

            if hasattr(activity, "__name__"):
                activity_name = activity.__name__
            elif hasattr(activity, "__func__"):
                activity_name = activity.__func__.__name__
            else:
                activity_name = str(activity)

            if activity_name == "call_llm":
                return {
                    "content": "Test response",
                    "tool_calls": [
                        {
                            "id": "call_test_123",
                            "type": "function",
                            "function": {
                                "name": "test_tool",
                                "arguments": '{"arg": "value"}',
                            },
                        }
                    ],
                    "role": "assistant",
                }
            elif activity_name == "run_tool":
                return {
                    "tool_call_id": "call_test_123",
                    "content": "tool result",
                    "role": "tool",
                    "name": "test_tool",
                }
            elif activity_name in [
                "record_initial_entry",
                "finalize_workflow",
                "save_tool_results",
            ]:
                return None

        mock_workflow_context.instance_id = "test-instance-123"
        mock_workflow_context.call_activity = Mock(side_effect=track_call_activity)

        # Set up minimal state (single entry model)
        entry = AgentWorkflowEntry(
            source=None,
            triggering_workflow_instance_id="parent-instance-123",
            messages=[],
            tool_history=[],
        )
        basic_durable_agent._infra._state_model = entry

        # Run the workflow generator
        with patch.object(
            basic_durable_agent._infra,
            "get_state",
            side_effect=lambda wid: basic_durable_agent._infra._state_model,
        ):
            workflow_gen = basic_durable_agent.agent_workflow(
                mock_workflow_context, message
            )

            # Step through the generator, sending results back
            result = None
            try:
                while True:
                    result = workflow_gen.send(result)
            except StopIteration as e:
                result = e.value

        # Verify that retry_policy was passed to critical activities
        assert len(call_activity_calls) >= 5, (
            f"Expected at least 3 activity calls, got {len(call_activity_calls)}"
        )

        # All activities should have retry_policy parameter
        for call in call_activity_calls:
            assert "retry_policy" in call, f"Missing retry_policy in call: {call}"
            assert call["retry_policy"] == basic_durable_agent._retry_policy, (
                f"Expected retry_policy {basic_durable_agent._retry_policy}, got {call['retry_policy']}"
            )

        # Verify the key activities were called
        activity_names = [
            getattr(call["activity"], "__name__", str(call["activity"]))
            for call in call_activity_calls
        ]
        assert "record_initial_entry" in activity_names, (
            f"Missing record_initial_entry in {activity_names}"
        )
        assert "call_llm" in activity_names, f"Missing call_llm in {activity_names}"
        assert "run_tool" in activity_names, f"Missing run_tool in {activity_names}"
        assert "save_tool_results" in activity_names, (
            f"Missing save_tool_results in {activity_names}"
        )
        assert "finalize_workflow" in activity_names, (
            f"Missing finalize_workflow in {activity_names}"
        )

    def test_agent_workflow_max_iterations_sets_custom_status(
        self, basic_durable_agent, mock_workflow_context, mock_tool
    ):
        """Test that set_custom_status('max_iterations_reached') is called when max iterations is reached."""
        message = {"task": "Test task"}

        # Set up agent with max_iterations=2 and register a tool
        basic_durable_agent.execution.max_iterations = 2
        basic_durable_agent.tool_executor.register_tool(mock_tool)

        # Track call count to ensure we always return tool calls for call_llm
        call_llm_count = 0

        # Mock call_activity to return responses with tool_calls to force iterations
        def mock_call_activity(activity, **kwargs):
            nonlocal call_llm_count
            # Get activity name - handle both bound methods and functions
            if hasattr(activity, "__name__"):
                activity_name = activity.__name__
            elif hasattr(activity, "__func__"):
                activity_name = activity.__func__.__name__
            else:
                activity_name = str(activity)

            if activity_name == "call_llm":
                call_llm_count += 1
                # Always return tool calls to force continuation (never return final response)
                return {
                    "content": "I'll use a tool",
                    "tool_calls": [
                        {
                            "id": f"call_{call_llm_count}",
                            "type": "function",
                            "function": {
                                "name": "test_tool",
                                "arguments": "{}",
                            },
                        }
                    ],
                    "role": "assistant",
                }
            elif activity_name == "run_tool":
                return {
                    "tool_call_id": f"call_{call_llm_count}",
                    "content": "tool result",
                    "role": "tool",
                    "name": "test_tool",
                }
            elif activity_name == "load_tools":
                return []  # Empty list of registered tools
            elif activity_name in [
                "record_initial_entry",
                "finalize_workflow",
                "save_tool_results",
                "broadcast_to_team",
                "return_response",
                "summarize",
            ]:
                return None
            return None

        mock_workflow_context.call_activity = Mock(side_effect=mock_call_activity)

        # Set up state
        entry = AgentWorkflowEntry(
            source="test_source",
            triggering_workflow_instance_id=None,
            messages=[],
            tool_history=[],
        )
        basic_durable_agent._infra._state_model = entry

        with (
            patch.object(
                basic_durable_agent._infra,
                "get_state",
                side_effect=lambda wid: basic_durable_agent._infra._state_model,
            ),
            patch.object(basic_durable_agent, "save_state"),
        ):
            workflow_gen = basic_durable_agent.agent_workflow(
                mock_workflow_context, message
            )
            # Consume the generator until it completes
            result = None
            try:
                while True:
                    result = workflow_gen.send(result)
            except StopIteration:
                pass

        # Verify set_custom_status was called with "max_iterations_reached"
        mock_workflow_context.set_custom_status.assert_called_once_with(
            "max_iterations_reached"
        )

    def test_agent_workflow_exception_raises_agent_error(
        self, basic_durable_agent, mock_workflow_context
    ):
        """Test that an exception in the workflow results in AgentError being raised."""
        message = {"task": "Test task"}

        # Mock call_activity to raise an exception
        test_exception = RuntimeError("Test workflow failure")

        def mock_call_activity(activity, **kwargs):
            # Get activity name - handle both bound methods and functions
            if hasattr(activity, "__name__"):
                activity_name = activity.__name__
            elif hasattr(activity, "__func__"):
                activity_name = activity.__func__.__name__
            else:
                activity_name = str(activity)

            if activity_name == "record_initial_entry":
                return None
            elif activity_name == "load_tools":
                return []  # Empty list of registered tools
            elif activity_name == "call_llm":
                # Raise exception on LLM call
                raise test_exception
            elif activity_name in [
                "broadcast_to_team",
                "return_response",
                "summarize",
                "finalize_workflow",
            ]:
                return None
            return None

        mock_workflow_context.call_activity = Mock(side_effect=mock_call_activity)

        # Set up state
        entry = AgentWorkflowEntry(
            source="test_source",
            triggering_workflow_instance_id=None,
            messages=[],
            tool_history=[],
        )
        basic_durable_agent._infra._state_model = entry

        with (
            patch.object(
                basic_durable_agent._infra,
                "get_state",
                side_effect=lambda wid: basic_durable_agent._infra._state_model,
            ),
            patch.object(basic_durable_agent, "save_state"),
        ):
            workflow_gen = basic_durable_agent.agent_workflow(
                mock_workflow_context, message
            )
            # Consume the generator until it raises AgentError
            # The exception is caught internally, workflow continues to finalize,
            # then AgentError is raised at the end
            with pytest.raises(AgentError) as exc_info:
                try:
                    while True:
                        next(workflow_gen)
                except StopIteration:
                    # Should not reach here - AgentError should be raised first
                    pass

            # Verify the exception message contains the original error and agent name
            assert "workflow failed" in str(exc_info.value).lower()
            assert "TestDurableAgent" in str(exc_info.value)
            # Verify it's chained to the original exception
            assert exc_info.value.__cause__ == test_exception
