import pytest
from unittest.mock import Mock, patch
from swarms.tools.agent_tools import *
from swarms.boss.boss_node import BossNodeInitializer, BossNode
# For initializing BossNodeInitializer in multiple tests
@pytest.fixture
def mock_boss_node_initializer():
    llm = Mock()
    vectorstore = Mock()
    agent_executor = Mock()
    max_iterations = 5

    boss_node_initializer = BossNodeInitializer(llm, vectorstore, agent_executor, max_iterations)

    return boss_node_initializer


# Test BossNodeInitializer class __init__ method
def test_boss_node_initializer_init(mock_boss_node_initializer):
    with patch('swarms.tools.agent_tools.BabyAGI.from_llm') as mock_from_llm:
        assert isinstance(mock_boss_node_initializer, BossNodeInitializer)
        mock_from_llm.assert_called_once()


# Test initialize_vectorstore method of BossNodeInitializer class
def test_boss_node_initializer_initialize_vectorstore(mock_boss_node_initializer):
    with patch('swarms.tools.agent_tools.OpenAIEmbeddings') as mock_embeddings, \
        patch('swarms.tools.agent_tools.FAISS') as mock_faiss:

        result = mock_boss_node_initializer.initialize_vectorstore()
        mock_embeddings.assert_called_once()
        mock_faiss.assert_called_once()
        assert result is not None


# Test initialize_llm method of BossNodeInitializer class
def test_boss_node_initializer_initialize_llm(mock_boss_node_initializer):
    with patch('swarms.tools.agent_tools.OpenAI') as mock_llm:
        result = mock_boss_node_initializer.initialize_llm(mock_llm)
        mock_llm.assert_called_once()
        assert result is not None


# Test create_task method of BossNodeInitializer class
@pytest.mark.parametrize("objective", ['valid objective', ''])
def test_boss_node_initializer_create_task(objective, mock_boss_node_initializer):
    if objective == '':
        with pytest.raises(ValueError):
            mock_boss_node_initializer.create_task(objective)
    else:
        assert mock_boss_node_initializer.create_task(objective) == {"objective": objective}


# Test run method of BossNodeInitializer class
@pytest.mark.parametrize("task", ['valid task', ''])
def test_boss_node_initializer_run(task, mock_boss_node_initializer):
    with patch.object(mock_boss_node_initializer, 'baby_agi'):
        if task == '':
            with pytest.raises(ValueError):
                mock_boss_node_initializer.run(task)
        else:
            try:
                mock_boss_node_initializer.run(task)
                mock_boss_node_initializer.baby_agi.assert_called_once_with(task)
            except Exception:
                pytest.fail("Unexpected Error!")


# Test BossNode function
@pytest.mark.parametrize("api_key, objective, llm_class, max_iterations",
                         [('valid_key', 'valid_objective', OpenAI, 5), 
                          ('', 'valid_objective', OpenAI, 5), 
                          ('valid_key', '', OpenAI, 5), 
                          ('valid_key', 'valid_objective', '', 5),
                          ('valid_key', 'valid_objective', OpenAI, 0)])
def test_boss_node(api_key, objective, llm_class, max_iterations):
    with patch('os.getenv') as mock_getenv, \
        patch('swarms.tools.agent_tools.PromptTemplate.from_template') as mock_from_template, \
        patch('swarms.tools.agent_tools.LLMChain') as mock_llm_chain, \
        patch('swarms.tools.agent_tools.ZeroShotAgent.create_prompt') as mock_create_prompt, \
        patch('swarms.tools.agent_tools.ZeroShotAgent') as mock_zero_shot_agent, \
        patch('swarms.tools.agent_tools.AgentExecutor.from_agent_and_tools') as mock_from_agent_and_tools, \
        patch('swarms.tools.agent_tools.BossNodeInitializer') as mock_boss_node_initializer, \
        patch.object(mock_boss_node_initializer, 'create_task') as mock_create_task, \
        patch.object(mock_boss_node_initializer, 'run') as mock_run:

        if api_key == '' or objective == '' or llm_class == '' or max_iterations <= 0:
            with pytest.raises(ValueError):
                BossNode(objective, api_key, vectorstore=None, worker_node=None, llm_class=llm_class, max_iterations=max_iterations, verbose=False)
        else:
            mock_getenv.return_value = 'valid_key'
            BossNode(objective, api_key, vectorstore=None, worker_node=None, llm_class=llm_class, max_iterations=max_iterations, verbose=False)
            mock_from_template.assert_called_once()
            mock_llm_chain.assert_called_once()
            mock_create_prompt.assert_called_once()
            mock_zero_shot_agent.assert_called_once()
            mock_from_agent_and_tools.assert_called_once()
            mock_boss_node_initializer.assert_called_once()
            mock_create_task.assert_called_once()
            mock_run.assert_called_once()