"""
Tests for Kaizen-enhanced LLMAgentNode with extension points and Kaizen integration.

This test module validates:
1. Extension points (pre/post execute, error handling)
2. Kaizen config compatibility
3. Backward compatibility with Core SDK usage
4. Strategy integration hooks
"""

import pytest
from unittest.mock import Mock, patch
from kaizen.nodes.ai import LLMAgentNode


class TestLLMAgentKaizenExtensions:
    """Test Kaizen-specific enhancements to LLMAgentNode."""

    def test_llm_agent_has_extension_point_method(self):
        """Test that LLMAgentNode exposes extension points."""
        agent = LLMAgentNode()

        # Check for extension point method
        assert hasattr(agent, '_get_extension_points'), \
            "LLMAgentNode should have _get_extension_points method"

    def test_extension_points_structure(self):
        """Test that extension points have expected structure."""
        agent = LLMAgentNode()
        extension_points = agent._get_extension_points()

        # Should be a dictionary
        assert isinstance(extension_points, dict), \
            "Extension points should be a dictionary"

        # Should have expected keys
        expected_keys = ['pre_execute', 'post_execute', 'on_error']
        for key in expected_keys:
            assert key in extension_points, \
                f"Extension points should include '{key}'"
            assert callable(extension_points[key]), \
                f"Extension point '{key}' should be callable"

    def test_pre_execute_hook_called(self):
        """Test that pre_execute hook is called during execution."""
        agent = LLMAgentNode()

        # Track if hook was called
        hook_called = {'called': False, 'inputs': None}

        original_hook = agent._pre_execute_hook
        def tracked_hook(inputs):
            hook_called['called'] = True
            hook_called['inputs'] = inputs
            return original_hook(inputs)

        agent._pre_execute_hook = tracked_hook

        # Execute with mock provider
        result = agent.execute(
            provider="mock",
            model="test-model",
            messages=[{"role": "user", "content": "test"}],
            mock_response="Test response"
        )

        # Verify hook was called
        assert hook_called['called'], "pre_execute_hook should be called"
        assert hook_called['inputs'] is not None, "pre_execute_hook should receive inputs"

    def test_post_execute_hook_called(self):
        """Test that post_execute hook is called after execution."""
        agent = LLMAgentNode()

        # Track if hook was called
        hook_called = {'called': False, 'result': None}

        original_hook = agent._post_execute_hook
        def tracked_hook(result):
            hook_called['called'] = True
            hook_called['result'] = result
            return original_hook(result)

        agent._post_execute_hook = tracked_hook

        # Execute with mock provider
        result = agent.execute(
            provider="mock",
            model="test-model",
            messages=[{"role": "user", "content": "test"}],
            mock_response="Test response"
        )

        # Verify hook was called
        assert hook_called['called'], "post_execute_hook should be called"
        assert hook_called['result'] is not None, "post_execute_hook should receive result"

    def test_on_error_hook_called_on_failure(self):
        """Test that on_error hook is called when errors occur."""
        agent = LLMAgentNode()

        # Track if hook was called
        hook_called = {'called': False, 'error': None}

        original_hook = agent._on_error_hook
        def tracked_hook(error, context):
            hook_called['called'] = True
            hook_called['error'] = error
            return original_hook(error, context)

        agent._on_error_hook = tracked_hook

        # Execute with invalid configuration to trigger error
        with pytest.raises(Exception):
            agent.execute(
                provider="invalid_provider",
                model="test-model",
                messages=[{"role": "user", "content": "test"}]
            )

        # Verify hook was called
        assert hook_called['called'], "on_error_hook should be called on error"
        assert hook_called['error'] is not None, "on_error_hook should receive error"

    def test_extension_point_customization(self):
        """Test that extension points can be customized via subclassing."""

        class CustomLLMAgent(LLMAgentNode):
            """Custom agent with enhanced pre-processing."""

            def _pre_execute_hook(self, inputs):
                """Add custom pre-processing."""
                # Add custom metadata
                if 'metadata' not in inputs:
                    inputs['metadata'] = {}
                inputs['metadata']['custom_processing'] = True
                return inputs

        agent = CustomLLMAgent()

        # Execute with mock provider
        result = agent.execute(
            provider="mock",
            model="test-model",
            messages=[{"role": "user", "content": "test"}],
            mock_response="Test response"
        )

        # Verify custom processing was applied
        assert result.get('success'), "Execution should succeed"

    def test_backward_compatibility_core_sdk_usage(self):
        """Test that LLMAgentNode maintains backward compatibility with Core SDK."""
        agent = LLMAgentNode()

        # Execute using Core SDK pattern (no Kaizen config)
        result = agent.execute(
            provider="mock",
            model="test-model",
            messages=[{"role": "user", "content": "What is 2+2?"}],
            mock_response="4"
        )

        # Should work exactly as before
        assert result.get('success'), "Core SDK usage should still work"
        assert 'response' in result, "Should return response"

    def test_extension_points_dont_break_normal_execution(self):
        """Test that extension points don't interfere with normal execution."""
        agent = LLMAgentNode()

        # Normal execution should work
        result = agent.execute(
            provider="mock",
            model="test-model",
            messages=[{"role": "user", "content": "Hello"}],
            mock_response="Hi there!"
        )

        assert result.get('success'), "Normal execution should succeed"
        assert 'response' in result, "Should have response"
        assert 'content' in result['response'], "Response should have content"


class TestLLMAgentKaizenConfig:
    """Test Kaizen configuration integration."""

    def test_accepts_standard_config(self):
        """Test that LLMAgentNode accepts standard configuration."""
        agent = LLMAgentNode()

        # Should work with standard dict config
        result = agent.execute(
            provider="mock",
            model="test-model",
            messages=[{"role": "user", "content": "test"}],
            mock_response="response"
        )

        assert result.get('success'), "Should accept standard config"

    def test_preserves_all_functionality(self):
        """Test that all original LLMAgentNode functionality is preserved."""
        agent = LLMAgentNode()

        # Test various features
        features_to_test = [
            # Basic execution
            {
                'provider': 'mock',
                'model': 'test',
                'messages': [{'role': 'user', 'content': 'test'}],
                'mock_response': 'response'
            },
            # With system prompt
            {
                'provider': 'mock',
                'model': 'test',
                'system_prompt': 'You are helpful',
                'messages': [{'role': 'user', 'content': 'test'}],
                'mock_response': 'response'
            },
            # With generation config
            {
                'provider': 'mock',
                'model': 'test',
                'messages': [{'role': 'user', 'content': 'test'}],
                'generation_config': {'temperature': 0.7},
                'mock_response': 'response'
            },
        ]

        for config in features_to_test:
            result = agent.execute(**config)
            assert result.get('success'), f"Feature test should succeed: {config}"


class TestLLMAgentNodeParameters:
    """Test LLMAgentNode parameter handling."""

    def test_get_parameters_includes_all_params(self):
        """Test that get_parameters returns all expected parameters."""
        agent = LLMAgentNode()
        params = agent.get_parameters()

        # Should have key parameters
        param_names = list(params.keys())
        expected_params = [
            'provider',
            'model',
            'messages',
            'system_prompt',
            'generation_config',
            'tools',
            'mcp_servers',
        ]

        for expected in expected_params:
            assert expected in param_names, \
                f"Parameter '{expected}' should be defined"

    def test_parameters_have_proper_types(self):
        """Test that parameters have proper type annotations."""
        agent = LLMAgentNode()
        params = agent.get_parameters()

        for param_name, param in params.items():
            assert hasattr(param, 'name'), "Parameter should have name"
            assert hasattr(param, 'type'), "Parameter should have type"
            assert hasattr(param, 'required'), "Parameter should have required flag"


class TestLLMAgentNodeIntegration:
    """Test LLMAgentNode integration with Kaizen framework."""

    def test_node_registration(self):
        """Test that LLMAgentNode is properly registered."""
        # Should be importable from kaizen.nodes.ai
        from kaizen.nodes.ai import LLMAgentNode as ImportedNode
        assert ImportedNode is LLMAgentNode, "Should be importable from kaizen.nodes.ai"

    def test_can_be_used_in_workflow(self):
        """Test that LLMAgentNode can be used in a workflow."""
        from kailash.workflow.builder import WorkflowBuilder
        from kailash.runtime.local import LocalRuntime

        # Build workflow
        workflow = WorkflowBuilder()
        workflow.add_node("LLMAgentNode", "agent", {
            "provider": "mock",
            "model": "test-model",
            "messages": [{"role": "user", "content": "test"}],
            "mock_response": "Test response"
        })

        # Execute workflow
        runtime = LocalRuntime()
        results, run_id = runtime.execute(workflow.build())

        # Verify execution
        assert 'agent' in results, "Agent node should be in results"
        assert results['agent'].get('success'), "Agent execution should succeed"


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