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

class TestGeminiAPIConnections:
    @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_connection(self, temp_workspace, thread):
        """Test actual connection to gemini-2.0-pro-exp-02-05 model."""
        if not os.getenv("GOOGLE_API_KEY"):
            pytest.skip("GOOGLE_API_KEY not found in environment variables")
            
        workspace = Path(temp_workspace)
        instructions_path = str(workspace / "instructions/basic_instructions.md")
        
        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
            }
        )
        
        # Add message to thread
        thread.add_message(
            role="user",
            content="What is quantum computing and how does it work?"
        )
        
        try:
            response = agent._process_with_llm(thread)
            assert response is not None
            assert len(response) > 0
            print(f"\nWorld Knowledge Model Response:\n{response}")
        except Exception as e:
            pytest.fail(f"World Knowledge Model API call failed: {str(e)}")

    def test_complex_reasoning_model_connection(self, temp_workspace, thread):
        """Test actual connection to gemini-2.0-flash-thinking-exp-01-21 model."""
        if not os.getenv("GOOGLE_API_KEY"):
            pytest.skip("GOOGLE_API_KEY not found in environment variables")
            
        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
            }
        )
        
        # Add message to thread
        thread.add_message(
            role="user",
            content="Analyze the ethical implications of AI in autonomous vehicles."
        )
        
        try:
            response = agent._process_with_llm(thread)
            assert response is not None
            assert len(response) > 0
            print(f"\nComplex Reasoning Model Response:\n{response}")
        except Exception as e:
            pytest.fail(f"Complex Reasoning Model API call failed: {str(e)}")

    def test_multimodal_model_connection(self, temp_workspace, thread):
        """Test actual connection to learnlm-1.5-pro-experimental model."""
        if not os.getenv("GOOGLE_API_KEY"):
            pytest.skip("GOOGLE_API_KEY not found in environment variables")
            
        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
            }
        )
        
        # Add message to thread
        thread.add_message(
            role="user",
            content="Explain how multimodal AI processes different types of input data."
        )
        
        try:
            response = agent._process_with_llm(thread)
            assert response is not None
            assert len(response) > 0
            print(f"\nMultimodal Model Response:\n{response}")
        except Exception as e:
            pytest.fail(f"Multimodal Model API call failed: {str(e)}")

    def test_fallback_behavior(self, temp_workspace, thread):
        """Test fallback to standard model when experimental model fails."""
        if not os.getenv("GOOGLE_API_KEY"):
            pytest.skip("GOOGLE_API_KEY not found in environment variables")
            
        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 fallback behavior with a simple query."
        )
        
        try:
            response = agent._process_with_llm(thread)
            assert response is not None
            assert len(response) > 0
            print(f"\nFallback Model Response:\n{response}")
        except Exception as e:
            # If experimental model fails, it should fall back to standard model
            provider, config = agent._determine_provider_from_description("fallback test")
            assert provider == "gemini"
            assert config["model"] == "gemini-pro" 