"""
Test CrewManager functionality

This module contains both unit tests (with mocks) and integration tests (with real OpenAI API).
Integration tests require OPENAI_API_KEY environment variable.

About resolve_tool's three resolution methods:

Method 1 (registry) - Tool Registry (highest priority):
  - Use case: Custom tools registered via @crew_tool() decorator or register_tool() function
  - Examples: TextProcessorTool, LimitedScrapeWebsiteTool, etc. registered with @crew_tool()
  - Advantages: Centralized management, easy to find and maintain
  - Applicable: Project custom tools, tools that need to be shared globally

Method 2 (crewai_tools) - CrewAI Official Tools Package (second priority):
  - Use case: Standard tools provided by CrewAI, from crewai_tools package
  - Examples: SerperDevTool, ScrapeWebsiteTool, FileReadTool, etc.
  - Advantages: No manual registration needed, can be referenced directly as strings
  - Applicable: Using CrewAI official standard tools, tools from third-party crewai_tools package

Method 3 (globals) - Call Stack Global Variables (fallback):
  - Use case: Functions/classes defined in the global scope of current module or call stack
  - Examples: Tool classes defined in the same file but not registered with @crew_tool()
  - Advantages: High flexibility, suitable for temporary tools or test scenarios
  - Applicable: Temporary tools, test tools, tools that don't want to be registered globally, module-level tool definitions
  - Note: Tools must be defined at module level (not inside functions) to be found
"""
import pytest
import os
from unittest.mock import Mock, patch, MagicMock

# Try to import CrewManager, skip tests if not available
try:
    from aipartnerupflow.extensions.crewai import CrewManager
    from aipartnerupflow.core.tools import register_tool, resolve_tool
    from crewai import LLM, Agent
except ImportError:
    CrewManager = None
    LLM = None
    Agent = None
    register_tool = None
    resolve_tool = None
    pytestmark = pytest.mark.skip("crewai module not available")


@pytest.mark.skipif(CrewManager is None, reason="CrewManager not available")
class TestCrewManager:
    """Test cases for CrewManager"""
    
    @pytest.mark.asyncio
    async def test_execute_with_mock(self):
        """Test crew execution with mocked CrewAI"""
        # Create mock crew result
        mock_result = Mock()
        mock_result.raw = "Test execution result"
        mock_result.token_usage = {
            "total_tokens": 100,
            "prompt_tokens": 60,
            "completion_tokens": 40,
            "cached_prompt_tokens": 0,
            "successful_requests": 1
        }
        
        # Create mock crew
        mock_crew = Mock()
        mock_crew.kickoff = Mock(return_value=mock_result)
        
        # Create mock Agent and Task to avoid LLM initialization
        mock_agent = Mock()
        mock_task = Mock()
        
        # Create CrewManager with all mocks in place
        with patch('aipartnerupflow.extensions.crewai.crew_manager.CrewAI') as mock_crew_class, \
             patch('aipartnerupflow.extensions.crewai.crew_manager.Agent') as mock_agent_class, \
             patch('aipartnerupflow.extensions.crewai.crew_manager.Task') as mock_task_class:
            mock_crew_class.return_value = mock_crew
            mock_agent_class.return_value = mock_agent
            mock_task_class.return_value = mock_task
            
            # Initialize CrewManager with works format
            crew_manager = CrewManager(
                name="Test Crew",
                works={
                    "agents": {
                        "researcher": {
                            "role": "Researcher",
                            "goal": "Research and gather information",
                            "backstory": "You are a research assistant"
                        }
                    },
                    "tasks": {
                        "research_task": {
                            "description": "Research a topic",
                            "expected_output": "A summary of the research findings",
                            "agent": "researcher"
                        }
                    }
                }
            )
            
            # Replace the crew instance with our mock
            crew_manager.crew = mock_crew
            
            # Execute crew
            result = await crew_manager.execute()
            
            # Verify result structure
            assert result["status"] == "success"
            assert result["result"] == "Test execution result"
            assert "token_usage" in result
            assert result["token_usage"]["total_tokens"] == 100
            assert result["token_usage"]["prompt_tokens"] == 60
            assert result["token_usage"]["completion_tokens"] == 40
            
            # Verify crew.kickoff was called
            mock_crew.kickoff.assert_called_once()
    
    @pytest.mark.asyncio
    async def test_execute_with_error_mock(self):
        """Test crew execution error handling with mocked CrewAI"""
        # Create mock crew that raises an error
        mock_crew = Mock()
        mock_crew.kickoff = Mock(side_effect=Exception("Test error"))
        
        # Mock Task class
        mock_task = Mock()
        
        # Create CrewManager with mock crew
        with patch('aipartnerupflow.extensions.crewai.crew_manager.CrewAI') as mock_crew_class, \
             patch('aipartnerupflow.extensions.crewai.crew_manager.Task') as mock_task_class, \
             patch('aipartnerupflow.extensions.crewai.crew_manager.Agent') as mock_agent_class:
            mock_crew_class.return_value = mock_crew
            mock_task_class.return_value = mock_task
            mock_agent_class.return_value = Mock()
            
            # Initialize CrewManager with works format
            crew_manager = CrewManager(
                name="Test Crew",
                works={
                    "agents": {
                        "researcher": {
                            "role": "Researcher",
                            "goal": "Research and gather information",
                            "backstory": "You are a research assistant"
                        }
                    },
                    "tasks": {
                        "research_task": {
                            "description": "Research a topic",
                            "expected_output": "A summary of the research findings",
                            "agent": "researcher"
                        }
                    }
                }
            )
            
            # Replace the crew instance with our mock
            crew_manager.crew = mock_crew
            
            # Execute crew (should handle error gracefully)
            result = await crew_manager.execute()
            
            # Verify error result structure
            assert result["status"] == "failed"
            assert "error" in result
            assert "Test error" in result["error"]
            assert result["result"] is None

    @pytest.mark.skipif(CrewManager is None or LLM is None, reason="CrewManager or LLM not available")
    def test_llm_string_conversion(self):
        """Test that string LLM names are converted to LLM objects"""
        # Mock LLM class
        mock_llm_instance = Mock()
        mock_llm_instance.model = "gpt-4"
        
        with patch('aipartnerupflow.extensions.crewai.crew_manager.LLM') as mock_llm_class:
            mock_llm_class.return_value = mock_llm_instance
            
            # Mock Agent, Task and CrewAI
            mock_agent = Mock()
            mock_task = Mock()
            mock_crew = Mock()
            
            with patch('aipartnerupflow.extensions.crewai.crew_manager.Agent') as mock_agent_class, \
                 patch('aipartnerupflow.extensions.crewai.crew_manager.Task') as mock_task_class, \
                 patch('aipartnerupflow.extensions.crewai.crew_manager.CrewAI') as mock_crew_class:
                mock_agent_class.return_value = mock_agent
                mock_task_class.return_value = mock_task
                mock_crew_class.return_value = mock_crew
                
                # Create CrewManager with string LLM in works format
                crew_manager = CrewManager(
                    name="Test Crew",
                    works={
                        "agents": {
                            "researcher": {
                                "role": "Researcher",
                                "goal": "Research",
                                "backstory": "You are a researcher",
                                "llm": "gpt-4"  # String LLM name
                            }
                        },
                        "tasks": {
                            "research_task": {
                                "description": "Research a topic",
                                "expected_output": "A summary of the research findings",
                                "agent": "researcher"
                            }
                        }
                    }
                )
                
                # Verify LLM was called with the model name
                mock_llm_class.assert_called_with(model="gpt-4")
                
                # Verify Agent was called with LLM object
                call_args = mock_agent_class.call_args
                assert call_args is not None
                assert "llm" in call_args.kwargs
                assert call_args.kwargs["llm"] == mock_llm_instance
    
    @pytest.mark.skipif(CrewManager is None or LLM is None, reason="CrewManager or LLM not available")
    def test_model_from_kwargs(self):
        """Test CrewManager receives model from kwargs (from schemas)"""
        with patch('aipartnerupflow.extensions.crewai.crew_manager.Agent') as mock_agent_class, \
             patch('aipartnerupflow.extensions.crewai.crew_manager.Task') as mock_task_class, \
             patch('aipartnerupflow.extensions.crewai.crew_manager.CrewAI') as mock_crew_class, \
             patch('aipartnerupflow.extensions.crewai.crew_manager.LLM') as mock_llm_class:
            
            mock_agent = Mock()
            mock_task = Mock()
            mock_crew = Mock()
            mock_llm_instance = Mock()
            
            mock_agent_class.return_value = mock_agent
            mock_task_class.return_value = mock_task
            mock_crew_class.return_value = mock_crew
            mock_llm_class.return_value = mock_llm_instance
            
            # Create CrewManager with model from kwargs (simulating schemas["model"])
            crew_manager = CrewManager(
                name="Test Crew",
                works={
                    "agents": {
                        "researcher": {
                            "role": "Researcher",
                            "goal": "Research",
                            "backstory": "You are a researcher"
                        }
                    },
                    "tasks": {
                        "research_task": {
                            "description": "Research a topic",
                            "expected_output": "A summary",
                            "agent": "researcher"
                        }
                    }
                },
                model="openai/gpt-4o"  # Model from schemas
            )
            
            # Verify model was set to self.llm
            assert crew_manager.llm == "openai/gpt-4o"
            
            # Verify LLM was called with the model name in _initialize_crew
            # (This happens during __init__, so we check the call was made)
            mock_llm_class.assert_called_with(model="openai/gpt-4o")
    
    @pytest.mark.skipif(CrewManager is None or resolve_tool is None, reason="CrewManager or resolve_tool not available")
    def test_tools_string_conversion(self):
        """Test that string tool names are converted to callable objects"""
        # Create mock tool
        mock_tool = Mock()
        mock_tool.run = Mock(return_value="tool result")
        
        # Mock resolve_tool to return our mock tool
        with patch('aipartnerupflow.extensions.crewai.crew_manager.resolve_tool') as mock_resolve_tool:
            mock_resolve_tool.return_value = mock_tool
            
            # Mock Agent, Task and CrewAI
            mock_agent = Mock()
            mock_task = Mock()
            mock_crew = Mock()
            
            with patch('aipartnerupflow.extensions.crewai.crew_manager.Agent') as mock_agent_class, \
                 patch('aipartnerupflow.extensions.crewai.crew_manager.Task') as mock_task_class, \
                 patch('aipartnerupflow.extensions.crewai.crew_manager.CrewAI') as mock_crew_class:
                mock_agent_class.return_value = mock_agent
                mock_task_class.return_value = mock_task
                mock_crew_class.return_value = mock_crew
                
                # Create CrewManager with string tool names in works format
                crew_manager = CrewManager(
                    name="Test Crew",
                    works={
                        "agents": {
                            "researcher": {
                                "role": "Researcher",
                                "goal": "Research",
                                "backstory": "You are a researcher",
                                "tools": ["SerperDevTool()", "ScrapeWebsiteTool()"]  # String tool names
                            }
                        },
                        "tasks": {
                            "research_task": {
                                "description": "Research a topic",
                                "expected_output": "A summary of the research findings",
                                "agent": "researcher"
                            }
                        }
                    }
                )
                
                # Verify resolve_tool was called for each tool
                assert mock_resolve_tool.call_count == 2
                mock_resolve_tool.assert_any_call("SerperDevTool()")
                mock_resolve_tool.assert_any_call("ScrapeWebsiteTool()")
                
                # Verify Agent was called with converted tools
                call_args = mock_agent_class.call_args
                assert call_args is not None
                assert "tools" in call_args.kwargs
                assert len(call_args.kwargs["tools"]) == 2
                assert call_args.kwargs["tools"][0] == mock_tool
                assert call_args.kwargs["tools"][1] == mock_tool
    
    @pytest.mark.skipif(CrewManager is None or LLM is None or resolve_tool is None, 
                        reason="Required modules not available")
    def test_llm_and_tools_together(self):
        """Test that both LLM and tools string conversion work together"""
        # Mock LLM
        mock_llm_instance = Mock()
        mock_llm_instance.model = "gpt-4"
        
        # Mock tool
        mock_tool = Mock()
        mock_tool.run = Mock(return_value="tool result")
        
        with patch('aipartnerupflow.extensions.crewai.crew_manager.LLM') as mock_llm_class, \
             patch('aipartnerupflow.extensions.crewai.crew_manager.resolve_tool') as mock_resolve_tool:
            mock_llm_class.return_value = mock_llm_instance
            mock_resolve_tool.return_value = mock_tool
            
            # Mock Agent, Task and CrewAI
            mock_agent = Mock()
            mock_task = Mock()
            mock_crew = Mock()
            
            with patch('aipartnerupflow.extensions.crewai.crew_manager.Agent') as mock_agent_class, \
                 patch('aipartnerupflow.extensions.crewai.crew_manager.Task') as mock_task_class, \
                 patch('aipartnerupflow.extensions.crewai.crew_manager.CrewAI') as mock_crew_class:
                mock_agent_class.return_value = mock_agent
                mock_task_class.return_value = mock_task
                mock_crew_class.return_value = mock_crew

                # Create CrewManager with both string LLM and tools in works format
                crew_manager = CrewManager(
                    name="Test Crew",
                    works={
                        "agents": {
                            "researcher": {
                                "role": "Researcher",
                                "goal": "Research",
                                "backstory": "You are a researcher",
                                "llm": "gpt-4",  # String LLM
                                "tools": ["SerperDevTool()"]  # String tool
                            }
                        },
                        "tasks": {
                            "research_task": {
                                "description": "Research a topic",
                                "expected_output": "A summary of the research findings",
                                "agent": "researcher"
                            }
                        }
                    }
                )
                
                # Verify LLM conversion
                mock_llm_class.assert_called_with(model="gpt-4")
                
                # Verify tools conversion
                mock_resolve_tool.assert_called_with("SerperDevTool()")
                
                # Verify Agent was called with both LLM and tools
                call_args = mock_agent_class.call_args
                assert call_args is not None
                assert "llm" in call_args.kwargs
                assert call_args.kwargs["llm"] == mock_llm_instance
                assert "tools" in call_args.kwargs
                assert len(call_args.kwargs["tools"]) == 1
                assert call_args.kwargs["tools"][0] == mock_tool
    

    
    @pytest.mark.asyncio
    @pytest.mark.skipif(
        not os.getenv("OPENAI_API_KEY"),
        reason="OPENAI_API_KEY is not set - skipping integration test"
    )
    async def test_execute_with_real_openai(self):
        """Test crew execution with real OpenAI API (requires OPENAI_API_KEY)"""
        # Check if OpenAI API key is available
        openai_key = os.getenv("OPENAI_API_KEY")
        if not openai_key:
            pytest.skip("OPENAI_API_KEY is not set")
        
        # Create a simple crew for testing
        # LLM is now set at agent level, not crew level
        crew_manager = CrewManager(
            name="Test Research Crew",
            works={
                "agents": {
                    "researcher": {
                        "role": "Researcher",
                        "goal": "Research and provide a brief summary",
                        "backstory": "You are a helpful research assistant",
                        "verbose": False,
                        "allow_delegation": False,
                        "llm": "openai/gpt-3.5-turbo"  # LLM set at agent level
                    }
                },
                "tasks": {
                    "research_task": {
                        "description": "Research and summarize what Python is in one sentence",
                        "expected_output": "A one-sentence summary of Python",
                        "agent": "researcher"
                    }
                }
            }
        )
        
        # Execute crew
        result = await crew_manager.execute()
        print("=== result: ===")
        import json
        print(json.dumps(result, indent=2, default=str))
        # Verify result structure
        assert result["status"] in ["success", "failed"]
        
        if result["status"] == "success":
            # Verify success result
            assert "result" in result
            assert result["result"] is not None
            
            # Verify token usage is present (if available)
            if "token_usage" in result:
                token_usage = result["token_usage"]
                assert "total_tokens" in token_usage or "status" in token_usage
        else:
            # If failed, verify error message
            assert "error" in result
            # Log the error for debugging
            print(f"Crew execution failed: {result.get('error')}")
    
    @pytest.mark.asyncio
    @pytest.mark.skipif(
        not os.getenv("OPENAI_API_KEY"),
        reason="OPENAI_API_KEY is not set - skipping integration test"
    )
    async def test_resolve_tool_method_2_crewai_tools(self):
        """
        Test Method 2 (crewai_tools) usage scenario
        
        Use case description:
        Method 2 is used to find tools provided by CrewAI (from crewai_tools package),
        such as SerperDevTool, ScrapeWebsiteTool, etc. These tools don't need manual registration,
        they can be used directly by importing from crewai_tools package.
        
        Applicable scenarios:
        - Using standard tools provided by CrewAI
        - Don't want to manually register tools, use string references directly
        - Tools from third-party crewai_tools package
        """
        # Check if OpenAI API key is available
        openai_key = os.getenv("OPENAI_API_KEY")
        if not openai_key:
            pytest.skip("OPENAI_API_KEY is not set")
        
        try:
            import crewai_tools
            from aipartnerupflow.core.tools import resolve_tool, get_tool_registry
        except ImportError:
            pytest.skip("crewai_tools or CrewAI module not available")
        
        # Test: Use tools from crewai_tools (if available)
        # Note: This is just a demonstration that resolve_tool will use Method 2
        # In actual usage, if a tool is not found in the registry, it will automatically try to find it from crewai_tools
        
        # First ensure the tool is not in the registry (if already in registry, will use Method 1)
        registry = get_tool_registry()
        
        # Try to resolve a tool that might come from crewai_tools
        # Note: This is just a demonstration, actual tool may not exist or require configuration
        try:
            # If SerperDevTool is not in the registry, resolve_tool will try to find it from crewai_tools
            if "SerperDevTool" not in registry.list_tools():
                # Try to resolve, should use Method 2
                tool = resolve_tool("SerperDevTool()")
                print(f"Resolved tool from crewai_tools: {type(tool).__name__}")
        except (NameError, AttributeError):
            # If tool doesn't exist, this is normal, we're just demonstrating the flow
            pass
        
        # The main purpose of this test is to demonstrate Method 2 usage scenario
        # In actual projects, when you use string references like "SerperDevTool()",
        # if the tool is not registered in the registry, resolve_tool will automatically search from crewai_tools package
        assert True  # Test passes, understanding is correct
    
    @pytest.mark.asyncio
    @pytest.mark.skipif(
        not os.getenv("OPENAI_API_KEY"),
        reason="OPENAI_API_KEY is not set - skipping integration test"
    )
    async def test_resolve_tool_method_3_globals(self):
        """
        Test Method 3 (globals) usage scenario
        
        Use case description:
        Method 3 is used to find functions/classes defined in the global variables of the current module or call stack.
        This is a fallback when:
        1. Tool is not registered with @crew_tool() (not using Method 1)
        2. Tool is not in crewai_tools package (not using Method 2)
        3. Tool is defined directly in the current scope
        
        Applicable scenarios:
        - Tool classes defined in the same file but not registered with @crew_tool()
        - Temporarily defined tools that don't want to be registered in the global registry
        - Quickly defined tools in test scenarios
        - Tools defined in the same scope where resolve_tool is called
        
        Note: Method 3 searches for tools by backtracking the call stack using inspect.currentframe(),
        so tools must be defined at module-level global scope, or in the global scope of some module
        in the call stack of resolve_tool.
        """
        # Check if OpenAI API key is available
        openai_key = os.getenv("OPENAI_API_KEY")
        if not openai_key:
            pytest.skip("OPENAI_API_KEY is not set")
        
        try:
            from aipartnerupflow.core.tools import resolve_tool, get_tool_registry, BaseTool
            from pydantic import BaseModel, Field
            from typing import Type
            import inspect
        except ImportError:
            pytest.skip("CrewAI tools module not available")
        
        # Ensure tool is not in the registry
        registry = get_tool_registry()
        
        # Define a tool class but don't register it with @crew_tool()
        # This way it won't appear in the registry (not using Method 1)
        class QuickCalculatorInputSchema(BaseModel):
            a: float = Field(..., description="First number")
            b: float = Field(..., description="Second number")
            operation: str = Field(default="add", description="Operation: add, subtract, multiply, divide")
        
        class QuickCalculatorTool(BaseTool):
            """A quick calculator tool defined in local scope"""
            name: str = "Quick Calculator"
            description: str = "Perform basic arithmetic operations"
            args_schema: Type[BaseModel] = QuickCalculatorInputSchema
            
            def _run(self, a: float, b: float, operation: str = "add") -> str:
                """Perform calculation"""
                if operation == "add":
                    return str(a + b)
                elif operation == "subtract":
                    return str(a - b)
                elif operation == "multiply":
                    return str(a * b)
                elif operation == "divide":
                    if b == 0:
                        return "Error: Division by zero"
                    return str(a / b)
                else:
                    return f"Unknown operation: {operation}"
        
        # Ensure tool is not in the registry
        assert "QuickCalculatorTool" not in registry.list_tools(), "Tool should not be in registry"
        
        # Add tool to current module's global scope so Method 3 can find it
        # Get current module
        current_module = inspect.getmodule(inspect.currentframe())
        if current_module:
            # Add tool to module's global dictionary
            current_module.__dict__['QuickCalculatorTool'] = QuickCalculatorTool
        
        # Now resolve_tool should be able to find it via Method 3
        try:
            tool = resolve_tool("QuickCalculatorTool()")
            print(f"Resolved tool from globals (Method 3): {type(tool).__name__}")
            # Verify tool works correctly
            result = tool._run(10, 5, "add")
            assert result == "15", f"Expected '15', got '{result}'"
        except NameError as e:
            # If not found, Method 3 is not applicable in this context
            # This is normal because the tool is defined in function scope
            print(f"Method 3 not applicable in this context: {e}")
            # The main purpose of this test is to demonstrate Method 3 usage scenario and limitations
            pass
        
        # Cleanup: Remove from module global dictionary
        if current_module and 'QuickCalculatorTool' in current_module.__dict__:
            del current_module.__dict__['QuickCalculatorTool']
    