"""
Unit tests for PromptManager

These tests verify the functionality of the PromptManager class,
including template rendering, term merging, JSON validation, and error handling.
"""

import pytest
import json
from unittest.mock import AsyncMock, MagicMock, patch

from trustgraph.template.prompt_manager import PromptManager, PromptConfiguration, Prompt


@pytest.mark.unit
class TestPromptManager:
    """Unit tests for PromptManager template functionality"""

    @pytest.fixture
    def sample_config(self):
        """Sample configuration dict for PromptManager"""
        return {
            "system": json.dumps("You are a helpful assistant."),
            "template-index": json.dumps(["simple_text", "json_response", "complex_template"]),
            "template.simple_text": json.dumps({
                "prompt": "Hello {{ name }}, welcome to {{ system_name }}!",
                "response-type": "text"
            }),
            "template.json_response": json.dumps({
                "prompt": "Generate a user profile for {{ username }}",
                "response-type": "json",
                "schema": {
                    "type": "object",
                    "properties": {
                        "name": {"type": "string"},
                        "age": {"type": "number"}
                    },
                    "required": ["name", "age"]
                }
            }),
            "template.complex_template": json.dumps({
                "prompt": """
                {% for item in items %}
                - {{ item.name }}: {{ item.value }}
                {% endfor %}
                Total: {{ items|length }}
                """,
                "response-type": "text"
            })
        }

    @pytest.fixture
    def prompt_manager(self, sample_config):
        """Create a PromptManager with sample configuration"""
        pm = PromptManager()
        pm.load_config(sample_config)
        # Add global terms manually since load_config doesn't handle them
        pm.terms["system_name"] = "TrustGraph"
        pm.terms["version"] = "1.0"
        return pm

    def test_prompt_manager_initialization(self, prompt_manager, sample_config):
        """Test PromptManager initialization with configuration"""
        assert prompt_manager.config.system_template == "You are a helpful assistant."
        assert len(prompt_manager.prompts) == 3
        assert "simple_text" in prompt_manager.prompts

    def test_simple_text_template_rendering(self, prompt_manager):
        """Test basic template rendering with text response"""
        terms = {"name": "Alice"}
        
        rendered = prompt_manager.render("simple_text", terms)
        
        assert rendered == "Hello Alice, welcome to TrustGraph!"

    def test_global_terms_merging(self, prompt_manager):
        """Test that global terms are properly merged"""
        terms = {"name": "Bob"}
        
        # Global terms should be available in template
        rendered = prompt_manager.render("simple_text", terms)
        
        assert "TrustGraph" in rendered  # From global terms
        assert "Bob" in rendered  # From input terms

    def test_term_override_priority(self, prompt_manager):
        """Test term override priority: input > prompt > global"""
        # Add a test prompt with overlapping terms
        test_config = {
            "template.test": json.dumps({
                "prompt": "Value is: {{ value }}",
                "response-type": "text"
            })
        }
        prompt_manager.load_config({**prompt_manager.config.__dict__, **test_config})
        prompt_manager.terms["value"] = "global"
        prompt_manager.prompts["test"].terms["value"] = "prompt"
        
        # Test with no input override
        rendered = prompt_manager.render("test", {})
        assert rendered == "Value is: prompt"  # Prompt terms override global
        
        # Test with input override
        rendered = prompt_manager.render("test", {"value": "input"})
        assert rendered == "Value is: input"  # Input terms override all

    def test_complex_template_rendering(self, prompt_manager):
        """Test complex template with loops and filters"""
        terms = {
            "items": [
                {"name": "Item1", "value": 10},
                {"name": "Item2", "value": 20},
                {"name": "Item3", "value": 30}
            ]
        }
        
        rendered = prompt_manager.render("complex_template", terms)
        
        assert "Item1: 10" in rendered
        assert "Item2: 20" in rendered
        assert "Item3: 30" in rendered
        assert "Total: 3" in rendered

    @pytest.mark.asyncio
    async def test_invoke_text_response(self, prompt_manager):
        """Test invoking a prompt with text response"""
        mock_llm = AsyncMock()
        mock_llm.return_value = "Welcome Alice to TrustGraph!"
        
        result = await prompt_manager.invoke(
            "simple_text",
            {"name": "Alice"},
            mock_llm
        )
        
        assert result == "Welcome Alice to TrustGraph!"
        
        # Verify LLM was called with correct prompts
        mock_llm.assert_called_once()
        call_args = mock_llm.call_args[1]
        assert call_args["system"] == "You are a helpful assistant."
        assert "Hello Alice, welcome to TrustGraph!" in call_args["prompt"]

    @pytest.mark.asyncio
    async def test_invoke_json_response_valid(self, prompt_manager):
        """Test invoking a prompt with valid JSON response"""
        mock_llm = AsyncMock()
        mock_llm.return_value = '{"name": "John Doe", "age": 30}'
        
        result = await prompt_manager.invoke(
            "json_response",
            {"username": "johndoe"},
            mock_llm
        )
        
        assert isinstance(result, dict)
        assert result["name"] == "John Doe"
        assert result["age"] == 30

    @pytest.mark.asyncio
    async def test_invoke_json_response_with_markdown(self, prompt_manager):
        """Test JSON extraction from markdown code blocks"""
        mock_llm = AsyncMock()
        mock_llm.return_value = """
        Here is the user profile:
        
        ```json
        {
            "name": "Jane Smith",
            "age": 25
        }
        ```
        
        This is a valid profile.
        """
        
        result = await prompt_manager.invoke(
            "json_response",
            {"username": "janesmith"},
            mock_llm
        )
        
        assert isinstance(result, dict)
        assert result["name"] == "Jane Smith"
        assert result["age"] == 25

    @pytest.mark.asyncio
    async def test_invoke_json_validation_failure(self, prompt_manager):
        """Test JSON schema validation failure"""
        mock_llm = AsyncMock()
        # Missing required 'age' field
        mock_llm.return_value = '{"name": "Invalid User"}'
        
        with pytest.raises(ValueError) as exc_info:
            await prompt_manager.invoke(
                "json_response",
                {"username": "invalid"},
                mock_llm
            )
        
        assert "JSON schema validation failed" in str(exc_info.value)

    @pytest.mark.asyncio
    async def test_invoke_json_parse_failure(self, prompt_manager):
        """Test invalid JSON parsing"""
        mock_llm = AsyncMock()
        mock_llm.return_value = "This is not JSON at all"
        
        with pytest.raises(ValueError) as exc_info:
            await prompt_manager.invoke(
                "json_response",
                {"username": "test"},
                mock_llm
            )
        
        assert "No JSON found in response" in str(exc_info.value)

    @pytest.mark.asyncio
    async def test_invoke_unknown_prompt(self, prompt_manager):
        """Test invoking an unknown prompt ID"""
        mock_llm = AsyncMock()
        
        with pytest.raises(KeyError):
            await prompt_manager.invoke(
                "nonexistent_prompt",
                {},
                mock_llm
            )

    def test_template_rendering_with_undefined_variable(self, prompt_manager):
        """Test template rendering with undefined variables"""
        terms = {}  # Missing 'name' variable
        
        # This should raise an error or use empty string depending on config
        with pytest.raises(Exception):  # ibis/Jinja2 will raise on undefined
            prompt_manager.render("simple_text", terms)

    @pytest.mark.asyncio
    async def test_json_response_without_schema(self):
        """Test JSON response without schema validation"""
        pm = PromptManager()
        config = {
            "system": json.dumps("Test"),
            "template-index": json.dumps(["no_schema"]),
            "template.no_schema": json.dumps({
                "prompt": "Generate any JSON",
                "response-type": "json"
                # No schema defined
            })
        }
        pm.load_config(config)
        
        mock_llm = AsyncMock()
        mock_llm.return_value = '{"any": "json", "is": "valid"}'
        
        result = await pm.invoke("no_schema", {}, mock_llm)
        
        assert result == {"any": "json", "is": "valid"}

    def test_prompt_configuration_validation(self):
        """Test PromptConfiguration validation"""
        # Valid configuration
        config = PromptConfiguration(
            system_template="Test system",
            prompts={
                "test": Prompt(
                    template="Hello {{ name }}",
                    response_type="text"
                )
            }
        )
        assert config.system_template == "Test system"
        assert len(config.prompts) == 1

    def test_nested_template_includes(self, prompt_manager):
        """Test templates with nested variable references"""
        # Add global terms
        prompt_manager.terms["company"] = "TrustGraph"
        prompt_manager.terms["year"] = 2024
        
        # Add a nested template
        config = {
            "template.nested": json.dumps({
                "prompt": "{{ greeting }} from {{ company }} in {{ year }}!",
                "response-type": "text"
            })
        }
        prompt_manager.load_config({**prompt_manager.config.__dict__, **config})
        prompt_manager.prompts["nested"].terms = {"greeting": "Welcome"}
        
        rendered = prompt_manager.render("nested", {"user": "Alice"})
        
        # Should contain company and year from global terms
        assert "TrustGraph" in rendered
        assert "2024" in rendered

    @pytest.mark.asyncio
    async def test_concurrent_invocations(self, prompt_manager):
        """Test concurrent prompt invocations"""
        mock_llm = AsyncMock()
        mock_llm.side_effect = [
            "Response for Alice",
            "Response for Bob",
            "Response for Charlie"
        ]
        
        # Simulate concurrent invocations
        import asyncio
        results = await asyncio.gather(
            prompt_manager.invoke("simple_text", {"name": "Alice"}, mock_llm),
            prompt_manager.invoke("simple_text", {"name": "Bob"}, mock_llm),
            prompt_manager.invoke("simple_text", {"name": "Charlie"}, mock_llm)
        )
        
        assert len(results) == 3
        assert "Alice" in results[0]
        assert "Bob" in results[1]
        assert "Charlie" in results[2]

    def test_empty_configuration(self):
        """Test PromptManager with minimal configuration"""
        pm = PromptManager()
        pm.load_config({})  # Empty config
        
        assert pm.config.system_template == "Be helpful."  # Default system
        assert pm.terms == {}  # Default empty terms
        assert len(pm.prompts) == 0