"""
Test suite for SubstrateAgent - the core agent implementation.

Tests cover:
- Agent initialization and configuration
- Task processing with and without retry
- Custom handlers and presets
- Error handling and classification
- Focus parameter application
- Relay publishing
- Plugin integration
- Token tracking
"""
import pytest
import asyncio
from unittest.mock import Mock, AsyncMock, patch
from typing import Dict, Any

from daita.agents.substrate import SubstrateAgent
from daita.config.base import AgentConfig, AgentType, RetryPolicy, RetryStrategy
from daita.core.exceptions import (
    AgentError, LLMError, ValidationError, InvalidDataError,
    TransientError, RetryableError, PermanentError
)
from daita.llm.mock import MockLLMProvider


class TestSubstrateAgentInitialization:
    """Test agent initialization and configuration."""
    
    def test_minimal_initialization(self):
        """Test creating agent with minimal configuration."""
        agent = SubstrateAgent()
        
        assert agent.name == "Substrate Agent"
        assert agent.agent_type == AgentType.SUBSTRATE
        assert agent.agent_id is not None
        assert len(agent.agent_id) > 0
        assert not agent.config.retry_enabled
        assert agent.llm is None
        assert agent.default_focus is None
        assert agent.relay is None
    
    def test_custom_configuration(self):
        """Test creating agent with custom configuration."""
        config = AgentConfig(
            name="Test Agent",
            enable_retry=True,
            retry_policy=RetryPolicy(max_retries=5, initial_delay=2.0)
        )

        llm = MockLLMProvider()

        agent = SubstrateAgent(
            config=config,
            llm_provider=llm,
            name="Custom Agent",
            focus=["field1", "field2"],
            relay="test_channel",
            api_key="test_key"  # Provide api_key to skip auto-detection
        )

        assert agent.name == "Custom Agent"  # name parameter overrides config
        assert agent.config.retry_enabled
        assert agent.config.retry_policy.max_retries == 5
        assert agent.llm == llm
        assert agent.default_focus == ["field1", "field2"]
        assert agent.relay == "test_channel"
    
    def test_agent_id_generation(self):
        """Test agent ID generation is unique."""
        agent1 = SubstrateAgent(name="Test Agent")
        agent2 = SubstrateAgent(name="Test Agent")
        
        assert agent1.agent_id != agent2.agent_id
        assert "test_agent" in agent1.agent_id.lower()
    
    def test_custom_handlers_initialization(self):
        """Test initialization with custom handlers."""
        async def custom_handler(data, context, agent):
            return {"custom": True, "data": data}
        
        handlers = {"custom_task": custom_handler}
        agent = SubstrateAgent(handlers=handlers)
        
        assert "custom_task" in agent.handlers
        assert agent.handlers["custom_task"] == custom_handler
        # Built-in handlers should still be present
        assert "analyze" in agent.handlers
        assert "transform" in agent.handlers


class TestSubstrateAgentProcessing:
    """Test core task processing functionality."""
    
    @pytest.fixture
    def agent(self):
        """Create test agent with mock LLM."""
        llm = MockLLMProvider()
        llm.set_response("default", "Mock LLM response for testing")

        return SubstrateAgent(
            name="Test Agent",
            llm_provider=llm,
            api_key="test_key"  # Provide api_key to skip auto-detection
        )
    
    @pytest.fixture
    def retry_agent(self):
        """Create test agent with retry enabled."""
        config = AgentConfig(
            name="Retry Agent",
            enable_retry=True,
            retry_policy=RetryPolicy(max_retries=3, initial_delay=0.1)
        )

        llm = MockLLMProvider()
        return SubstrateAgent(config=config, llm_provider=llm, api_key="test_key")
    
    @pytest.mark.asyncio
    async def test_basic_task_processing(self, agent):
        """Test basic task processing without errors."""
        await agent.start()
        
        result = await agent.process("test_task", {"test": "data"})
        
        assert result["status"] == "success"
        assert "result" in result
        assert result["agent_id"] == agent.agent_id
        assert result["agent_name"] == agent.name
        
        await agent.stop()
    
    @pytest.mark.asyncio
    async def test_llm_query_handler(self, agent):
        """Test built-in LLM query handler."""
        await agent.start()
        
        result = await agent.process("llm_query", "Analyze this text")
        
        assert result["status"] == "success"
        assert "agent_response" in result["result"]
        assert "prompt_used" in result["result"]
        
        await agent.stop()
    
    @pytest.mark.asyncio
    async def test_analyze_handler(self, agent):
        """Test built-in analyze handler."""
        await agent.start()
        
        test_data = {"sales": [100, 200, 300], "region": "North"}
        result = await agent.process("analyze", test_data)
        
        assert result["status"] == "success"
        assert "data_type" in result["result"]
        assert "llm_analysis" in result["result"]
        
        await agent.stop()
    
    @pytest.mark.asyncio
    async def test_transform_handler(self, agent):
        """Test built-in transform handler."""
        await agent.start()
        
        result = await agent.process("transform", {"test": "data"})
        
        assert result["status"] == "success"
        assert "transformation_type" in result["result"]
        assert "transformed" in result["result"]
        
        await agent.stop()
    
    @pytest.mark.asyncio
    async def test_custom_handler(self, agent):
        """Test custom handler execution."""
        async def custom_handler(data, context, agent_instance):
            return {
                "custom_result": True,
                "received_data": data,
                "agent_name": agent_instance.name
            }
        
        agent.add_handler("custom_task", custom_handler)
        await agent.start()
        
        result = await agent.process("custom_task", {"input": "test"})
        
        assert result["status"] == "success"
        assert result["result"]["custom_result"] is True
        assert result["result"]["received_data"] == {"input": "test"}
        assert result["result"]["agent_name"] == agent.name
        
        await agent.stop()
    
    @pytest.mark.asyncio
    async def test_sync_handler(self, agent):
        """Test synchronous handler execution."""
        def sync_handler(data, context, agent_instance):
            return {"sync": True, "data": data}
        
        agent.add_handler("sync_task", sync_handler)
        await agent.start()
        
        result = await agent.process("sync_task", "test")
        
        assert result["status"] == "success"
        assert result["result"]["sync"] is True
        
        await agent.stop()


class TestSubstrateAgentRetry:
    """Test retry behavior and error handling."""
    
    @pytest.mark.asyncio
    async def test_no_retry_on_success(self):
        """Test that successful operations don't trigger retry."""
        config = AgentConfig(
            name="Test Agent",
            enable_retry=True,
            retry_policy=RetryPolicy(max_retries=3)
        )

        agent = SubstrateAgent(config=config)
        await agent.start()

        result = await agent.process("custom", {"test": "data"})

        assert result["status"] == "success"
        assert result["retry_info"]["attempt"] == 1
        # No retry needed for successful operations

        await agent.stop()
    
    @pytest.mark.asyncio
    async def test_retry_on_transient_error(self):
        """Test retry behavior with transient errors."""
        config = AgentConfig(
            name="Test Agent",
            enable_retry=True,
            retry_policy=RetryPolicy(max_retries=2, initial_delay=0.01)
        )

        call_count = 0

        async def failing_handler(data, context, agent):
            nonlocal call_count
            call_count += 1
            if call_count <= 2:
                raise TransientError("Temporary failure")
            return {"success": True, "attempts": call_count}

        agent = SubstrateAgent(config=config, handlers={"failing_task": failing_handler})
        await agent.start()

        result = await agent.process("failing_task", {})

        assert result["status"] == "success"
        assert result["retry_info"]["attempt"] == 3
        assert call_count == 3
        # Retries were attempted and succeeded

        await agent.stop()
    
    @pytest.mark.asyncio
    async def test_no_retry_on_permanent_error(self):
        """Test that permanent errors are not retried."""
        config = AgentConfig(
            name="Test Agent",
            enable_retry=True,
            retry_policy=RetryPolicy(max_retries=3)
        )

        call_count = 0

        # Create a handler that raises PermanentError directly without any wrapping
        async def failing_handler(data, context, agent):
            nonlocal call_count
            call_count += 1
            # Create the error with explicit retry hint
            error = PermanentError("This should not be retried")
            error.retry_hint = "permanent"  # Ensure retry hint is set
            raise error

        agent = SubstrateAgent(config=config, handlers={"failing_task": failing_handler})
        await agent.start()

        result = await agent.process("failing_task", {})

        assert result["status"] == "error"
        assert result["retry_info"]["attempt"] == 1
        assert call_count == 1  # Only called once, no retries on permanent errors

        await agent.stop()
    
    @pytest.mark.asyncio
    async def test_fail_fast_mode(self):
        """Test fail-fast mode (retry disabled)."""
        agent = SubstrateAgent(name="Fail Fast Agent")  # No retry enabled

        async def failing_handler(data, context, agent):
            raise Exception("This will fail immediately")

        agent.add_handler("failing_task", failing_handler)
        await agent.start()

        result = await agent.process("failing_task", {})

        assert result["status"] == "error"
        assert result["retry_info"] is None  # No retry info in fail-fast mode

        await agent.stop()


class TestSubstrateAgentFocus:
    """Test focus parameter application."""
    
    @pytest.fixture
    def agent(self):
        return SubstrateAgent(name="Focus Test Agent")
    
    @pytest.mark.asyncio
    async def test_dict_focus_list(self, agent):
        """Test focus with list of keys on dictionary data."""
        agent.default_focus = ["name", "age"]
        await agent.start()
        
        data = {"name": "John", "age": 30, "city": "NYC", "job": "Engineer"}
        result = await agent.process("process_data", data)
        
        # The focused data should be passed to the handler
        assert result["status"] == "success"
        
        await agent.stop()
    
    @pytest.mark.asyncio
    async def test_dict_focus_single(self, agent):
        """Test focus with single key on dictionary data."""
        agent.default_focus = "name"
        await agent.start()
        
        data = {"name": "John", "age": 30, "city": "NYC"}
        result = await agent.process("process_data", data)
        
        assert result["status"] == "success"
        
        await agent.stop()
    
    @pytest.mark.asyncio
    async def test_focus_validation_error(self, agent):
        """Test focus validation with invalid configuration."""
        # Test with whitespace-only string - agent logs warning and continues with original data
        agent.default_focus = "   "  # Whitespace-only string
        await agent.start()

        result = await agent.process("process_data", {"test": "data"})

        # Focus application fails gracefully, processing continues with original data
        assert result["status"] == "success"

        await agent.stop()
    
    @pytest.mark.asyncio 
    async def test_no_focus_passthrough(self, agent):
        """Test that data passes through unchanged when no focus is set."""
        await agent.start()
        
        data = {"complex": {"nested": {"data": "value"}}}
        result = await agent.process("process_data", data)
        
        assert result["status"] == "success"
        
        await agent.stop()


class TestSubstrateAgentRelay:
    """Test relay channel publishing."""
    
    @pytest.mark.asyncio
    async def test_relay_publishing(self):
        """Test that successful results are published to relay channel."""
        # Patch the actual import location inside the _publish_to_relay method
        with patch('daita.core.relay.publish') as mock_publish:
            agent = SubstrateAgent(name="Relay Agent", relay="test_channel")
            await agent.start()
            
            result = await agent.process("custom", {"test": "data"})
            
            assert result["status"] == "success"
            
            # Verify publish was called
            mock_publish.assert_called_once()
            
            # Check the call arguments (publish is called with keyword arguments)
            call_args = mock_publish.call_args
            assert call_args[1]["channel"] == "test_channel"
            assert call_args[1]["publisher"] == "Relay Agent"
            assert call_args[1]["agent_response"]["status"] == "success"
            assert call_args[1]["agent_response"]["agent_name"] == "Relay Agent"
            
            await agent.stop()
    
    @pytest.mark.asyncio
    async def test_relay_error_handling(self):
        """Test that relay publishing errors don't break main processing."""
        with patch('daita.core.relay.publish', side_effect=Exception("Relay failed")):
            agent = SubstrateAgent(name="Relay Agent", relay="test_channel")
            await agent.start()
            
            # Main processing should still succeed despite relay failure
            result = await agent.process("custom", {"test": "data"})
            
            assert result["status"] == "success"
            
            await agent.stop()


class TestSubstrateAgentUtilities:
    """Test utility methods and properties."""
    
    def test_add_remove_handlers(self):
        """Test adding and removing custom handlers."""
        agent = SubstrateAgent()
        
        async def test_handler(data, context, agent):
            return {"test": True}
        
        # Add handler
        agent.add_handler("test_task", test_handler)
        assert "test_task" in agent.handlers
        
        # Remove handler
        agent.remove_handler("test_task")
        assert "test_task" not in agent.handlers
        
        # Remove non-existent handler should not error
        agent.remove_handler("non_existent")
    
    def test_handler_validation(self):
        """Test handler validation."""
        agent = SubstrateAgent()

        # Empty task name should raise error
        with pytest.raises(ValidationError):
            agent.add_handler("", lambda: None)

        # Non-callable should raise error
        with pytest.raises(ValidationError):
            agent.add_handler("test", "not callable")

    @pytest.mark.asyncio
    async def test_health_property(self):
        """Test health property provides comprehensive status."""
        llm = MockLLMProvider()
        agent = SubstrateAgent(
            name="Health Test",
            llm_provider=llm,
            relay="test_channel",
            api_key="test_key"
        )
        
        health = agent.health

        assert health["name"] == "Health Test"
        assert health["id"] == agent.agent_id
        assert health["type"] == "substrate"
        assert "handlers" in health
        assert "plugins" in health
        assert "relay" in health
        assert "llm" in health

        # Check relay and LLM status
        assert health["relay"]["enabled"] is True
        assert health["relay"]["channel"] == "test_channel"
        assert health["llm"]["available"] is True
    
    def test_token_usage_tracking(self):
        """Test token usage tracking integration."""
        llm = MockLLMProvider()
        agent = SubstrateAgent(llm_provider=llm, agent_id="test_agent_123", api_key="test_key")

        # Should return token usage data structure
        usage = agent.get_token_usage()

        assert isinstance(usage, dict)
        assert "total_tokens" in usage
        assert "prompt_tokens" in usage
        assert "completion_tokens" in usage
        assert "total_calls" in usage  # MockLLMProvider uses "total_calls" instead of "requests"


class TestSubstrateAgentPlugins:
    """Test plugin integration."""
    
    def test_plugin_access(self):
        """Test plugin access object."""
        agent = SubstrateAgent()
        
        # Should have plugin access
        assert hasattr(agent, 'plugins')
        assert hasattr(agent.plugins, 'postgresql')
        assert hasattr(agent.plugins, 'mysql') 
        assert hasattr(agent.plugins, 'mongodb')
        assert hasattr(agent.plugins, 'rest')
    
    def test_add_plugin(self):
        """Test adding plugins to agent."""
        agent = SubstrateAgent()
        
        mock_plugin = Mock()
        mock_plugin.__class__.__name__ = "MockPlugin"
        
        agent.add_plugin(mock_plugin)
        
        assert mock_plugin in agent.plugin_instances
        
        # Should appear in health info
        health = agent.health
        assert health["plugins"]["count"] == 1
        assert "MockPlugin" in health["plugins"]["types"]


class TestSubstrateAgentEdgeCases:
    """Test edge cases and error conditions."""
    
    @pytest.mark.asyncio
    async def test_empty_task_name(self):
        """Test processing with empty task name."""
        agent = SubstrateAgent()
        await agent.start()

        result = await agent.process("", {"test": "data"})

        # Empty task uses default handler
        assert result["status"] == "success"

        await agent.stop()
    
    @pytest.mark.asyncio
    async def test_none_data(self):
        """Test processing with None data."""
        agent = SubstrateAgent()
        await agent.start()
        
        result = await agent.process("custom", None)
        
        # Should not error with None data
        assert result["status"] == "success"
        
        await agent.stop()
    
    @pytest.mark.asyncio
    async def test_unknown_task_fallback(self):
        """Test fallback to LLM for unknown tasks."""
        llm = MockLLMProvider()
        agent = SubstrateAgent(llm_provider=llm, api_key="test_key")
        await agent.start()
        
        result = await agent.process("unknown_task", {"test": "data"})
        
        # Should fall back to LLM query
        assert result["status"] == "success"
        
        await agent.stop()
    
    @pytest.mark.asyncio
    async def test_no_llm_fallback(self):
        """Test behavior when no LLM is available for unknown tasks."""
        agent = SubstrateAgent()  # No LLM provider
        await agent.start()

        result = await agent.process("unknown_task", {"test": "data"})

        # Unknown task with no LLM uses default handler
        assert result["status"] == "success"

        await agent.stop()


if __name__ == "__main__":
    # Run tests with pytest
    pytest.main([__file__, "-v"])