import pytest
import os
from pathlib import Path
import tempfile
import shutil
from multi_swarm.core.agent import Agent
from multi_swarm.core.thread import Thread
from multi_swarm.core.agency import Agency

class TestExperimentalModels:
    @pytest.fixture
    def temp_workspace(self):
        """Create a temporary workspace for testing."""
        temp_dir = tempfile.mkdtemp()
        # Create instructions and tools directories
        os.makedirs(os.path.join(temp_dir, "instructions"))
        os.makedirs(os.path.join(temp_dir, "tools"))
        
        # Create basic instructions file
        with open(os.path.join(temp_dir, "instructions", "basic_instructions.md"), "w") as f:
            f.write("# Agent Instructions\nPerform tasks as requested.")
        
        yield temp_dir
        shutil.rmtree(temp_dir)
    
    @pytest.fixture
    def thread(self):
        """Create a test thread."""
        return Thread()

    def test_world_knowledge_model(self, temp_workspace, thread):
        """Test the gemini-2.0-pro-exp-02-05 model for world knowledge tasks"""
        workspace = Path(temp_workspace)
        instructions_path = str(workspace / "instructions/basic_instructions.md")
        
        # Create agent with world knowledge configuration
        agent = Agent(
            name="WorldKnowledgeAgent",
            description="Agent for testing world knowledge capabilities",
            instructions=instructions_path,
            tools_folder=str(workspace / "tools"),
            llm_provider="gemini",
            provider_config={
                "model": "gemini-2.0-pro-exp-02-05",
                "api_version": "2024-02",
                "max_tokens": 8192,
                "temperature": 0.7
            }
        )
        
        # Verify the agent's configuration
        assert agent.name == "WorldKnowledgeAgent"
        assert agent.config.provider_config["model"] == "gemini-2.0-pro-exp-02-05"
        assert os.path.exists(instructions_path)
        
        # Add message to thread
        thread.add_message(
            role="user",
            content="Explain the impact of quantum computing on cryptography"
        )
        
        # Mock the response for testing
        agent.client = type('MockClient', (), {
            'chat': type('MockChat', (), {
                'completions': type('MockCompletions', (), {
                    'create': lambda **kwargs: type('MockResponse', (), {
                        'choices': [type('MockChoice', (), {
                            'message': type('MockMessage', (), {
                                'content': "Mocked response about quantum computing impact"
                            })
                        })]
                    })
                })
            })
        })
        
        response = agent._process_with_llm(thread)
        assert response == "Mocked response about quantum computing impact"

    def test_complex_reasoning_model(self, temp_workspace, thread):
        """Test the gemini-2.0-flash-thinking-exp-01-21 model for complex reasoning"""
        workspace = Path(temp_workspace)
        instructions_path = str(workspace / "instructions/basic_instructions.md")
        
        agent = Agent(
            name="ReasoningAgent",
            description="Agent for testing complex reasoning capabilities",
            instructions=instructions_path,
            tools_folder=str(workspace / "tools"),
            llm_provider="gemini",
            provider_config={
                "model": "gemini-2.0-flash-thinking-exp-01-21",
                "api_version": "2024-01",
                "max_tokens": 4096,
                "temperature": 0.8
            }
        )
        
        # Verify the agent's configuration
        assert agent.name == "ReasoningAgent"
        assert agent.config.provider_config["model"] == "gemini-2.0-flash-thinking-exp-01-21"
        assert os.path.exists(instructions_path)
        
        # Add message to thread
        thread.add_message(
            role="user",
            content="Analyze the ethical implications of AI decision-making in healthcare"
        )
        
        # Mock the response for testing
        agent.client = type('MockClient', (), {
            'chat': type('MockChat', (), {
                'completions': type('MockCompletions', (), {
                    'create': lambda **kwargs: type('MockResponse', (), {
                        'choices': [type('MockChoice', (), {
                            'message': type('MockMessage', (), {
                                'content': "Mocked response about AI ethics in healthcare"
                            })
                        })]
                    })
                })
            })
        })
        
        response = agent._process_with_llm(thread)
        assert response == "Mocked response about AI ethics in healthcare"

    def test_multimodal_model(self, temp_workspace, thread):
        """Test the learnlm-1.5-pro-experimental model for multimodal tasks"""
        workspace = Path(temp_workspace)
        instructions_path = str(workspace / "instructions/basic_instructions.md")
        
        agent = Agent(
            name="MultimodalAgent",
            description="Agent for testing multimodal capabilities",
            instructions=instructions_path,
            tools_folder=str(workspace / "tools"),
            llm_provider="gemini",
            provider_config={
                "model": "learnlm-1.5-pro-experimental",
                "api_version": "2024-02",
                "max_tokens": 4096,
                "temperature": 0.6
            }
        )
        
        # Verify the agent's configuration
        assert agent.name == "MultimodalAgent"
        assert agent.config.provider_config["model"] == "learnlm-1.5-pro-experimental"
        assert os.path.exists(instructions_path)
        
        # Add message to thread
        thread.add_message(
            role="user",
            content="Describe the capabilities of multimodal processing"
        )
        
        # Mock the response for testing
        agent.client = type('MockClient', (), {
            'chat': type('MockChat', (), {
                'completions': type('MockCompletions', (), {
                    'create': lambda **kwargs: type('MockResponse', (), {
                        'choices': [type('MockChoice', (), {
                            'message': type('MockMessage', (), {
                                'content': "Mocked response about multimodal capabilities"
                            })
                        })]
                    })
                })
            })
        })
        
        response = agent._process_with_llm(thread)
        assert response == "Mocked response about multimodal capabilities"

    def test_automatic_model_selection(self, temp_workspace):
        """Test automatic model selection based on task description"""
        workspace = Path(temp_workspace)
        instructions_path = str(workspace / "instructions/basic_instructions.md")
        
        # Create an agent without specifying model
        agent = Agent(
            name="AutoSelectAgent",
            description="Agent for testing automatic model selection",
            instructions=instructions_path,
            tools_folder=str(workspace / "tools"),
            llm_provider="gemini"
        )
        
        # Test world knowledge task
        world_knowledge_task = "Explain the theory of relativity"
        provider, config = agent._determine_provider_from_description(world_knowledge_task)
        assert provider == "gemini"
        assert isinstance(config, dict)
        assert config.get("model") == "gemini-2.0-pro-exp-02-05"

    def test_model_fallback(self, temp_workspace, thread):
        """Test fallback behavior when experimental models are unavailable"""
        workspace = Path(temp_workspace)
        instructions_path = str(workspace / "instructions/basic_instructions.md")
        
        agent = Agent(
            name="FallbackAgent",
            description="Agent for testing fallback behavior",
            instructions=instructions_path,
            tools_folder=str(workspace / "tools"),
            llm_provider="gemini",
            provider_config={
                "model": "gemini-2.0-pro-exp-02-05",
                "api_version": "2024-02"
            }
        )
        
        # Add message to thread
        thread.add_message(
            role="user",
            content="Test query"
        )
        
        # Simulate model unavailability
        try:
            # Mock an error response
            agent.client = type('MockClient', (), {
                'chat': type('MockChat', (), {
                    'completions': type('MockCompletions', (), {
                        'create': lambda **kwargs: (_ for _ in ()).throw(Exception("Model unavailable"))
                    })
                })
            })
            
            response = agent._process_with_llm(thread)
        except Exception as e:
            # Should fall back to default model
            provider, config = agent._determine_provider_from_description("fallback test")
            assert provider == "gemini"
            assert isinstance(config, dict)
            assert config.get("model") == "gemini-pro"  # Default model

    def test_version_compatibility(self, temp_workspace):
        """Test version compatibility checks"""
        workspace = Path(temp_workspace)
        instructions_path = str(workspace / "instructions/basic_instructions.md")
        
        # Test with compatible version
        agent1 = Agent(
            name="CompatibleAgent",
            description="Agent with compatible version",
            instructions=instructions_path,
            tools_folder=str(workspace / "tools"),
            llm_provider="gemini",
            provider_config={
                "model": "gemini-2.0-pro-exp-02-05",
                "api_version": "2024-02"
            }
        )
        
        # Test with incompatible version
        with pytest.raises(ValueError):
            Agent(
                name="IncompatibleAgent",
                description="Agent with incompatible version",
                instructions=instructions_path,
                tools_folder=str(workspace / "tools"),
                llm_provider="gemini",
                provider_config={
                    "model": "gemini-2.0-pro-exp-02-05",
                    "api_version": "2023-01"  # Old version
                }
            )

    def test_multi_agent_coordination(self, temp_workspace, thread):
        """Test coordination between agents with different experimental models"""
        workspace = Path(temp_workspace)
        instructions_path = str(workspace / "instructions/basic_instructions.md")
        
        # Create agents with different experimental models
        world_agent = Agent(
            name="WorldKnowledge",
            description="World knowledge specialist",
            instructions=instructions_path,
            tools_folder=str(workspace / "tools"),
            llm_provider="gemini",
            provider_config={
                "model": "gemini-2.0-pro-exp-02-05",
                "api_version": "2024-02"
            }
        )
        
        reasoning_agent = Agent(
            name="Reasoning",
            description="Complex reasoning specialist",
            instructions=instructions_path,
            tools_folder=str(workspace / "tools"),
            llm_provider="gemini",
            provider_config={
                "model": "gemini-2.0-flash-thinking-exp-01-21",
                "api_version": "2024-01"
            }
        )
        
        # Create agency
        agency = Agency(
            name="ExperimentalAgency",
            description="Agency with experimental models",
            agents=[world_agent, reasoning_agent],
            flows=[(world_agent, reasoning_agent), (reasoning_agent, world_agent)],
            storage_path=temp_workspace
        )
        
        # Verify agency configuration
        assert len(agency.agents) == 2
        assert agency.can_communicate(world_agent.name, reasoning_agent.name)
        assert agency.can_communicate(reasoning_agent.name, world_agent.name)

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