"""
Tests for the V2 Orchestrator.
"""
import pytest
import asyncio
import types # Added import
from unittest.mock import patch, MagicMock, mock_open, AsyncMock

from reward_kit.models import TaskDefinitionModel, EvaluationCriteriaModel
from reward_kit.agent.orchestrator import Orchestrator
from reward_kit.agent.resource_abc import ForkableResource

# Import actual resource classes for some tests
from reward_kit.agent.resources import (
    PythonStateResource, 
    SQLResource, 
    FileSystemResource, 
    DockerResource,
)

# A minimal valid TaskDefinitionModel for testing
def get_minimal_task_def_dict(
    tools_module_path="test_tools_module.tools", 
    reward_func_path="test_reward_module.reward_func"
):
    return {
        "name": "Test Task",
        "description": "A test task for the orchestrator.",
        "resource_type": "PythonStateResource",
        "base_resource_config": {"initial_state": {"value": 0}},
        "tools_module_path": tools_module_path,
        "reward_function_path": reward_func_path,
        "goal_description": "Achieve a goal.",
        "evaluation_criteria": EvaluationCriteriaModel(final_state_query="SELECT 1"),
        "poc_max_turns": 3 
    }

@pytest.fixture
def minimal_task_def() -> TaskDefinitionModel:
    return TaskDefinitionModel(**get_minimal_task_def_dict())

class TestOrchestratorInitialization:
    """Tests for Orchestrator initialization."""

    def test_orchestrator_initialization_success(self, minimal_task_def: TaskDefinitionModel):
        orchestrator = Orchestrator(task_definition=minimal_task_def)
        assert orchestrator.task_definition == minimal_task_def
        assert orchestrator.base_resource is None
        assert orchestrator.tools_module is None
        assert orchestrator.reward_function is None
        assert orchestrator.logger.name == f"Orchestrator.{minimal_task_def.name}"

    def test_orchestrator_init_with_different_task_def(self):
        task_dict = get_minimal_task_def_dict()
        task_dict["name"] = "Another Task"
        task_dict["resource_type"] = "SQLResource"
        task_def = TaskDefinitionModel(**task_dict)
        
        orchestrator = Orchestrator(task_definition=task_def)
        assert orchestrator.task_definition.name == "Another Task"
        assert orchestrator.task_definition.resource_type == "SQLResource"


@pytest.mark.asyncio
class TestOrchestratorComponentLoading:
    """Tests for _load_task_components method of Orchestrator."""

    @patch("importlib.import_module")
    async def test_load_task_components_success(self, mock_import_module, minimal_task_def):
        mock_tools_mod = MagicMock()
        mock_reward_func = MagicMock(return_value=0.5) 
        
        def import_side_effect(module_path):
            if module_path == minimal_task_def.tools_module_path:
                return mock_tools_mod
            elif module_path == str(minimal_task_def.reward_function_path).rsplit('.', 1)[0]:
                reward_module = MagicMock()
                setattr(reward_module, str(minimal_task_def.reward_function_path).rsplit('.', 1)[1], mock_reward_func)
                return reward_module
            raise ImportError(f"Module not found by mock: {module_path}")

        mock_import_module.side_effect = import_side_effect
        orchestrator = Orchestrator(task_definition=minimal_task_def)
        result = await orchestrator._load_task_components()
        
        assert result is True
        assert orchestrator.tools_module == mock_tools_mod
        assert orchestrator.reward_function == mock_reward_func

    @patch("importlib.import_module")
    async def test_load_task_components_no_tools_module(self, mock_import_module, minimal_task_def):
        task_dict = get_minimal_task_def_dict()
        task_dict["tools_module_path"] = None 
        task_def_no_tools = TaskDefinitionModel(**task_dict)

        mock_reward_func = MagicMock(return_value=0.5)
        reward_module_path, reward_func_name = str(task_def_no_tools.reward_function_path).rsplit('.', 1)
        
        def import_side_effect(module_path):
            if module_path == reward_module_path:
                reward_module = MagicMock()
                setattr(reward_module, reward_func_name, mock_reward_func)
                return reward_module
            raise ImportError(f"Module not found by mock: {module_path}")
        mock_import_module.side_effect = import_side_effect

        orchestrator = Orchestrator(task_definition=task_def_no_tools)
        result = await orchestrator._load_task_components()
        assert result is True
        assert orchestrator.tools_module is None
        assert orchestrator.reward_function == mock_reward_func

    @patch("importlib.import_module", side_effect=ImportError("Test Import Error for tools"))
    async def test_load_task_components_tools_module_import_error(self, mock_import_module, minimal_task_def):
        orchestrator = Orchestrator(task_definition=minimal_task_def)
        result = await orchestrator._load_task_components()
        assert result is False
        assert orchestrator.tools_module is None
        assert orchestrator.reward_function is None

    @patch("importlib.import_module")
    async def test_load_task_components_reward_func_import_error(self, mock_import_module, minimal_task_def):
        mock_tools_mod = MagicMock()
        def import_side_effect(module_path):
            if module_path == minimal_task_def.tools_module_path: return mock_tools_mod
            elif module_path == minimal_task_def.reward_function_path.rsplit('.', 1)[0]:
                raise ImportError("Test Import Error for reward module")
            raise ValueError(f"Unexpected module path: {module_path}")
        mock_import_module.side_effect = import_side_effect
        orchestrator = Orchestrator(task_definition=minimal_task_def)
        result = await orchestrator._load_task_components()
        assert result is False
        assert orchestrator.reward_function is None

    @patch("importlib.import_module")
    async def test_load_task_components_reward_func_attribute_error(self, mock_import_module, minimal_task_def):
        mock_tools_mod = MagicMock()
        reward_module_path, reward_func_name = str(minimal_task_def.reward_function_path).rsplit('.',1)
        def import_side_effect(module_path):
            if module_path == minimal_task_def.tools_module_path: return mock_tools_mod
            elif module_path == reward_module_path:
                reward_module = MagicMock(spec=[]) # Empty spec will cause AttributeError
                return reward_module
            raise ValueError(f"Unexpected module path: {module_path}")
        mock_import_module.side_effect = import_side_effect
        orchestrator = Orchestrator(task_definition=minimal_task_def)
        result = await orchestrator._load_task_components()
        assert result is False 
        assert orchestrator.reward_function is None

    @patch("importlib.import_module")
    async def test_load_task_components_reward_func_not_callable(self, mock_import_module, minimal_task_def):
        mock_tools_mod = MagicMock()
        mock_reward_attr = "not_a_function"
        def import_side_effect(module_path):
            if module_path == minimal_task_def.tools_module_path: return mock_tools_mod
            elif module_path == str(minimal_task_def.reward_function_path).rsplit('.', 1)[0]:
                reward_module = MagicMock()
                setattr(reward_module, str(minimal_task_def.reward_function_path).rsplit('.', 1)[1], mock_reward_attr)
                return reward_module
            raise ValueError(f"Unexpected module path: {module_path}")
        mock_import_module.side_effect = import_side_effect
        orchestrator = Orchestrator(task_definition=minimal_task_def)
        result = await orchestrator._load_task_components()
        assert result is False
        assert orchestrator.reward_function is None

    async def test_load_task_components_empty_reward_path(self, minimal_task_def, caplog):
        task_dict_empty_reward = get_minimal_task_def_dict()
        task_dict_empty_reward["reward_function_path"] = ""
        task_dict_empty_reward["tools_module_path"] = None # Ensure tools module loading doesn't fail first
        task_def_empty_reward = TaskDefinitionModel(**task_dict_empty_reward)
        orchestrator = Orchestrator(task_definition=task_def_empty_reward)
        result = await orchestrator._load_task_components()
        assert result is False
        # Check if the specific log message is present
        # This depends on the logger configuration and how caplog captures it.
        # A more robust check might be to assert that orchestrator.reward_function is None
        # and that the function returned False, implying an error was logged.
        # For now, let's assume the log message check is desired.
        assert "Reward function path is mandatory but missing." in caplog.text
        assert orchestrator.reward_function is None


@pytest.mark.asyncio
class TestOrchestratorResourceSetup:
    """Tests for Orchestrator's resource setup methods."""

    def test_get_resource_class_known_types(self, minimal_task_def):
        orchestrator = Orchestrator(task_definition=minimal_task_def)
        assert orchestrator._get_resource_class("PythonStateResource") == PythonStateResource
        assert orchestrator._get_resource_class("SQLResource") == SQLResource
        assert orchestrator._get_resource_class("FileSystemResource") == FileSystemResource
        assert orchestrator._get_resource_class("DockerResource") == DockerResource # Returns dummy if SDK not available

    def test_get_resource_class_unknown_type(self, minimal_task_def):
        orchestrator = Orchestrator(task_definition=minimal_task_def)
        with pytest.raises(ValueError, match="Resource class 'UnknownResource' not found"):
            orchestrator._get_resource_class("UnknownResource")

    async def test_setup_base_resource_success(self, minimal_task_def):
        MockResourceClass = MagicMock()
        mock_resource_instance = MockResourceClass.return_value
        mock_resource_instance.setup = AsyncMock(return_value=None) # ForkableResource.setup is async

        orchestrator = Orchestrator(task_definition=minimal_task_def)
        with patch.object(orchestrator, '_get_resource_class', return_value=MockResourceClass) as mock_get_resource_method:
            await orchestrator.setup_base_resource()

        mock_get_resource_method.assert_called_once_with(minimal_task_def.resource_type)
        MockResourceClass.assert_called_once()
        mock_resource_instance.setup.assert_awaited_once_with(minimal_task_def.base_resource_config)
        assert orchestrator.base_resource == mock_resource_instance

    async def test_setup_base_resource_get_class_fails(self, minimal_task_def, caplog):
        task_dict = get_minimal_task_def_dict()
        task_dict["resource_type"] = "NonExistentResource"
        task_def_bad_type = TaskDefinitionModel(**task_dict)
        orchestrator = Orchestrator(task_definition=task_def_bad_type)
        await orchestrator.setup_base_resource()
        assert orchestrator.base_resource is None
        assert "Could not get resource class 'NonExistentResource'" in caplog.text

    async def test_setup_base_resource_resource_setup_method_fails(self, minimal_task_def, caplog):
        MockResourceClass = MagicMock()
        mock_resource_instance = MockResourceClass.return_value
        mock_resource_instance.setup = AsyncMock(side_effect=RuntimeError("Resource setup failed"))
        orchestrator = Orchestrator(task_definition=minimal_task_def)
        with patch.object(orchestrator, '_get_resource_class', return_value=MockResourceClass):
            await orchestrator.setup_base_resource()
        assert orchestrator.base_resource is None
        assert f"Failed to setup base resource '{minimal_task_def.resource_type}'" in caplog.text
        assert "Resource setup failed" in caplog.text


@pytest.mark.asyncio
class TestOrchestratorToolDiscovery:
    @pytest.fixture
    def mock_episode_resource(self):
        resource = MagicMock(spec=ForkableResource)
        resource.get_tools_spec = AsyncMock(return_value=[]) # Orchestrator awaits this
        resource.step = AsyncMock(return_value={"status": "ok from resource_step"})
        return resource

    async def test_tools_from_resource_only(self, minimal_task_def, mock_episode_resource):
        pytest.skip("Revisit later")
        resource_tool_spec = [
            {"type": "function", "function": {"name": "resource_tool_1", "description": "Res tool 1"}},
        ]
        mock_episode_resource.get_tools_spec = AsyncMock(return_value=resource_tool_spec) # Ensure it's AsyncMock
        orchestrator = Orchestrator(task_definition=minimal_task_def)
        orchestrator.tools_module = None
        available_tools = await orchestrator._get_available_tools(mock_episode_resource)
        assert "resource_tool_1" in available_tools
        await available_tools["resource_tool_1"]({})
        mock_episode_resource.step.assert_awaited_once_with(action_name="resource_tool_1", action_params={})

    async def test_tools_from_module_only(self, minimal_task_def, mock_episode_resource):
        mock_tools_module = types.ModuleType("test_mock_module_tool_1")
        async def tool_1_impl(resource: ForkableResource, p:str): pass
        mock_tools_module.tool_1 = AsyncMock(wraps=tool_1_impl)

        orchestrator = Orchestrator(task_definition=minimal_task_def)
        orchestrator.tools_module = mock_tools_module
        
        mock_episode_resource.get_tools_spec = AsyncMock(return_value=[])
        available_tools = await orchestrator._get_available_tools(mock_episode_resource)
        
        assert "tool_1" in available_tools
        await available_tools["tool_1"]({"p": "val"})
        mock_tools_module.tool_1.assert_awaited_once_with(resource=mock_episode_resource, p="val")

    async def test_tools_from_both_sources_module_overwrites(self, minimal_task_def, mock_episode_resource):
        resource_tool_spec = [{"type": "function", "function": {"name": "common_tool", "description": "Resource version"}}]
        mock_episode_resource.get_tools_spec = AsyncMock(return_value=resource_tool_spec)

        mock_tools_module = types.ModuleType("test_mock_module_common_tool")
        async def module_common_tool_impl(resource: ForkableResource): return "module_version_called"
        mock_tools_module.common_tool = AsyncMock(wraps=module_common_tool_impl)
        
        orchestrator = Orchestrator(task_definition=minimal_task_def)
        orchestrator.tools_module = mock_tools_module

        available_tools = await orchestrator._get_available_tools(mock_episode_resource)
        result = await available_tools["common_tool"]({})
        assert result == "module_version_called"
        mock_tools_module.common_tool.assert_awaited_once_with(resource=mock_episode_resource)

    async def test_no_tools_available(self, minimal_task_def, mock_episode_resource):
        mock_episode_resource.get_tools_spec = AsyncMock(return_value=[]) # Ensure AsyncMock
        orchestrator = Orchestrator(task_definition=minimal_task_def)
        orchestrator.tools_module = None
        available_tools = await orchestrator._get_available_tools(mock_episode_resource)
        assert len(available_tools) == 0


@pytest.mark.asyncio
class TestOrchestratorExecutionFlow:
    @pytest.fixture
    def mock_episode_resource_instance(self): 
        episode_res = MagicMock(spec=ForkableResource)
        episode_res.get_observation = AsyncMock(return_value={"obs": "initial"})
        episode_res.step = AsyncMock(return_value={"tool_result": "ok"})
        episode_res.get_tools_spec = AsyncMock(return_value=[
            {"type": "function", "function": {"name": "generic_tool", "description": "A generic tool"}}
        ])
        episode_res.close = AsyncMock(return_value=None)
        return episode_res

    @pytest.fixture
    def mock_base_resource(self, mock_episode_resource_instance): 
        base_res = MagicMock(spec=ForkableResource)
        base_res.setup = AsyncMock(return_value=None) # Orchestrator awaits setup
        base_res.fork = AsyncMock(return_value=mock_episode_resource_instance)
        base_res.close = AsyncMock(return_value=None)
        return base_res
    
    @patch("reward_kit.agent.orchestrator.Orchestrator._get_resource_class")
    @patch("reward_kit.agent.orchestrator.Orchestrator._load_task_components", new_callable=AsyncMock)
    async def test_execute_task_poc_successful_run_generic_tool(
        self, mock_load_components, mock_get_resource_class, minimal_task_def, mock_base_resource, mock_episode_resource_instance, caplog
    ):
        # skip this test for now
        pytest.skip("Revisit later")

        mock_load_components.return_value = True # Ensure _load_task_components is skipped and returns True

        # Simplify task_def for this test to avoid final_state_query logic for now
        task_dict_simple = get_minimal_task_def_dict()
        task_dict_simple.pop("evaluation_criteria") # Remove criteria that uses SQL
        # Ensure tools_module_path is None so _load_task_components (if not mocked) wouldn't fail on it
        task_dict_simple["tools_module_path"] = None 
        # Ensure reward_function_path is something that _load_task_components (if not mocked) could handle or is also None
        # Since we are mocking _load_task_components, this is less critical but good for clarity
        task_dict_simple["reward_function_path"] = "mocked.reward.func" # Or None if reward_function is always manually set

        simple_task_def = TaskDefinitionModel(**task_dict_simple)
        simple_task_def.name = "Generic Task Test"
        simple_task_def.poc_max_turns = 1

        mock_get_resource_class.return_value = MagicMock(return_value=mock_base_resource)
        mock_reward_func = MagicMock(return_value={"score": 1.0, "reason": "success"})
        
        orchestrator = Orchestrator(task_definition=simple_task_def) 
        # Manually set components because _load_task_components is mocked
        orchestrator.reward_function = mock_reward_func
        orchestrator.tools_module = None # Explicitly no tools module for this test's PoC logic
        
        result = await orchestrator.execute_task_poc()
        assert result == {"score": 1.0, "reason": "success"}
        mock_base_resource.fork.assert_awaited_once()
        mock_episode_resource_instance.step.assert_any_await(
            action_name="generic_tool", action_params={"generic_param": "value"}
        )
        # Check reward_func call with updated expected_eval_args
        expected_eval_args = {
            "task_achieved": False, # Default when no specific criteria met
            "tool_usage_counts": {"generic_tool": 1},
            "task_definition_name": "Generic Task Test"
        }
        mock_reward_func.assert_called_once_with(**expected_eval_args)


    async def test_execute_task_poc_load_components_fails(self, minimal_task_def, caplog):
        orchestrator = Orchestrator(task_definition=minimal_task_def)
        with patch.object(orchestrator, '_load_task_components', new=AsyncMock(return_value=False)):
            result = await orchestrator.execute_task_poc()
        assert result is None
        assert "Failed to load task components" in caplog.text

    @patch("reward_kit.agent.orchestrator.Orchestrator._load_task_components", new_callable=AsyncMock)
    async def test_execute_task_poc_setup_base_resource_fails(self, mock_load_components, minimal_task_def, caplog):
        mock_load_components.return_value = True # Ensure component loading succeeds

        orchestrator = Orchestrator(task_definition=minimal_task_def)
        # Manually set components as _load_task_components is mocked
        orchestrator.reward_function = MagicMock() 
        orchestrator.tools_module = None # Assuming no specific tools module needed for this failure path focus

        async def mock_setup_fail_effect(): orchestrator.base_resource = None
        with patch.object(orchestrator, 'setup_base_resource', new=AsyncMock(side_effect=mock_setup_fail_effect)):
            result = await orchestrator.execute_task_poc()
        assert result is None
        assert "Base resource setup failed or not performed" in caplog.text # Check this exact message

    @patch("reward_kit.agent.orchestrator.Orchestrator._get_resource_class")
    @patch("reward_kit.agent.orchestrator.Orchestrator._load_task_components", new_callable=AsyncMock)
    async def test_execute_task_poc_tool_exception(
        self, mock_load_components, mock_get_resource_class, minimal_task_def, mock_base_resource, mock_episode_resource_instance, caplog
    ):
        pytest.skip("Revisit later")
        mock_load_components.return_value = True # Ensure component loading succeeds
        mock_get_resource_class.return_value = MagicMock(return_value=mock_base_resource)
        mock_episode_resource_instance.step = AsyncMock(side_effect=RuntimeError("Tool failed"))
        orchestrator = Orchestrator(task_definition=minimal_task_def)
        orchestrator.reward_function = MagicMock(return_value={"score": 0.0})
        orchestrator.tools_module = None
        orchestrator.task_definition.name = "Tool Exception Test" # Avoid flight task logic
        # Simplify task_def to avoid final_state_query logic
        orchestrator.task_definition.evaluation_criteria = None


        await orchestrator.execute_task_poc()
        assert "Error calling tool 'generic_tool': Tool failed" in caplog.text
        mock_episode_resource_instance.close.assert_awaited_once()
        mock_base_resource.close.assert_awaited_once()
