"""
Tests for the middleware system.
"""

import asyncio
import time
from typing import Any, Dict, List, Optional, Type, Union
from unittest.mock import AsyncMock, MagicMock

import pytest
from pydantic import BaseModel

from picoagents._middleware import (
    BaseMiddleware,
    GuardrailMiddleware,
    LoggingMiddleware,
    MetricsMiddleware,
    MiddlewareChain,
    MiddlewareContext,
    PIIRedactionMiddleware,
    RateLimitMiddleware,
)
from picoagents.agents import Agent
from picoagents.context import AgentContext
from picoagents.llm import BaseChatCompletionClient
from picoagents.messages import (
    AssistantMessage,
    SystemMessage,
    ToolCallRequest,
    UserMessage,
)
from picoagents.tools import BaseTool
from picoagents.types import ChatCompletionResult, ToolResult, Usage


class MockChatCompletionClient(BaseChatCompletionClient):
    """Mock chat completion client for testing."""

    def __init__(self, model: str = "test-model"):
        super().__init__(model=model)
        self.responses: List[AssistantMessage] = []
        self.call_count = 0

    def set_response(self, content: str, tool_calls: Optional[List] = None):
        """Set the response the mock client will return."""
        message = AssistantMessage(
            content=content, source="mock", tool_calls=tool_calls
        )
        self.responses = [message]

    def set_responses(self, responses: List[AssistantMessage]):
        """Set multiple responses to be returned in sequence."""
        self.responses = responses

    async def create(
        self,
        messages: List[Any],
        tools: Optional[List[Dict[str, Any]]] = None,
        output_format: Optional[Type[BaseModel]] = None,
        **kwargs: Any,
    ) -> ChatCompletionResult:
        """Mock implementation of create method."""
        self.call_count += 1

        if not self.responses:
            response = AssistantMessage(content="Test response", source="mock")
        else:
            # Pop from responses list so each call gets a different response
            # If only one response, keep returning it (backwards compatible)
            if len(self.responses) > 1:
                response = self.responses.pop(0)
            else:
                response = self.responses[0]

        return ChatCompletionResult(
            message=response,
            model="test-model",
            finish_reason="stop",
            usage=Usage(
                duration_ms=100,
                llm_calls=1,
                tokens_input=50,
                tokens_output=25,
                tool_calls=0,
                memory_operations=0,
            ),
        )

    async def create_stream(self, messages, tools=None, output_format=None, **kwargs):
        """Mock stream method - required by abstract base."""
        from picoagents.types import ChatCompletionChunk

        # Yield chunks that simulate streaming
        yield ChatCompletionChunk(
            content="Test", is_complete=False, tool_call_chunk=None
        )
        yield ChatCompletionChunk(
            content=" response", is_complete=True, tool_call_chunk=None
        )


class MockTool(BaseTool):
    """Mock tool for testing."""

    def __init__(self, name: str = "mock_tool", result: str = "mock result"):
        super().__init__(name=name, description="Mock tool for testing")
        self.result_value = result
        self._parameters_schema = {"type": "object", "properties": {}}

    @property
    def parameters(self) -> dict:
        return self._parameters_schema

    async def execute(self, parameters: Dict[str, Any]) -> ToolResult:
        """Mock tool execution."""
        return ToolResult(
            success=True,
            result=self.result_value,
            error=None,
            metadata={"tool_name": self.name},
        )


# Test Middleware Implementations


class MockMiddleware(BaseMiddleware):
    """Test middleware that tracks calls."""

    def __init__(self):
        self.request_calls = []
        self.response_calls = []
        self.error_calls = []

    async def process_request(self, context: MiddlewareContext):
        self.request_calls.append(context.operation)
        # Modify data to test request processing
        if context.operation == "model_call" and isinstance(context.data, list):
            # Add a marker to the last message
            context.metadata["test_marker"] = "request_processed"
        yield context

    async def process_response(self, context: MiddlewareContext, result: Any):
        self.response_calls.append(context.operation)
        # Modify result to test response processing
        if hasattr(result, "message") and hasattr(result.message, "content"):
            # Create new message with modified content (Pydantic models are immutable)
            modified_message = result.message.model_copy(
                update={"content": result.message.content + " [test_processed]"}
            )
            # Create new result with modified message
            result = result.model_copy(update={"message": modified_message})
        yield result

    async def process_error(
        self, context: MiddlewareContext, error: Exception
    ):
        self.error_calls.append((context.operation, type(error).__name__))
        raise error
        yield  # pragma: no cover


class BlockingMiddleware(BaseMiddleware):
    """Middleware that blocks certain operations."""

    def __init__(self, block_operations: Optional[List[str]] = None):
        self.block_operations = block_operations or []

    async def process_request(self, context: MiddlewareContext):
        if context.operation in self.block_operations:
            raise ValueError(f"Operation {context.operation} blocked by middleware")
        yield context

    async def process_response(self, context: MiddlewareContext, result: Any):
        yield result

    async def process_error(
        self, context: MiddlewareContext, error: Exception
    ):
        raise error
        yield  # pragma: no cover


# Tests


@pytest.fixture
def mock_agent():
    """Create a mock agent for testing."""
    client = MockChatCompletionClient()
    context = AgentContext()
    agent = Agent(
        name="test_agent",
        description="Test agent",
        instructions="Test instructions",
        model_client=client,
        context=context,
    )
    return agent, client, context


class TestMiddlewareChain:
    """Test the middleware chain execution."""

    @pytest.mark.asyncio
    async def test_empty_chain(self):
        """Test middleware chain with no middleware."""
        chain = MiddlewareChain([])
        context = AgentContext()

        async def mock_func(data):
            return "test_result"

        result = None
        async for item in chain.execute_stream(
            operation="test_op",
            agent_name="test_agent",
            agent_context=context,
            data="test_data",
            func=mock_func,
        ):
            result = item

        assert result == "test_result"

    @pytest.mark.asyncio
    async def test_single_middleware(self):
        """Test middleware chain with single middleware."""
        test_middleware = MockMiddleware()
        chain = MiddlewareChain([test_middleware])
        context = AgentContext()

        async def mock_func(data):
            return "test_result"

        result = None
        async for item in chain.execute_stream(
            operation="test_op",
            agent_name="test_agent",
            agent_context=context,
            data="test_data",
            func=mock_func,
        ):
            result = item

        assert result == "test_result"
        assert test_middleware.request_calls == ["test_op"]
        assert test_middleware.response_calls == ["test_op"]
        assert test_middleware.error_calls == []

    @pytest.mark.asyncio
    async def test_multiple_middleware_order(self):
        """Test that middleware executes in correct order."""
        middleware1 = MockMiddleware()
        middleware2 = MockMiddleware()
        chain = MiddlewareChain([middleware1, middleware2])
        context = AgentContext()

        async def mock_func(data):
            return "test_result"

        async for item in chain.execute_stream(
            operation="test_op",
            agent_name="test_agent",
            agent_context=context,
            data="test_data",
            func=mock_func,
        ):
            pass  # Just consume the stream

        # Request processing: forward order
        assert middleware1.request_calls == ["test_op"]
        assert middleware2.request_calls == ["test_op"]

        # Response processing: reverse order
        assert middleware1.response_calls == ["test_op"]
        assert middleware2.response_calls == ["test_op"]

    @pytest.mark.asyncio
    async def test_middleware_blocking(self):
        """Test middleware can block operations."""
        blocking_middleware = BlockingMiddleware(["blocked_op"])
        chain = MiddlewareChain([blocking_middleware])
        context = AgentContext()

        async def mock_func(data):
            return "should_not_execute"

        with pytest.raises(ValueError, match="Operation blocked_op blocked"):
            async for item in chain.execute_stream(
                operation="blocked_op",
                agent_name="test_agent",
                agent_context=context,
                data="test_data",
                func=mock_func,
            ):
                pass

    @pytest.mark.asyncio
    async def test_error_handling(self):
        """Test error handling through middleware."""
        test_middleware = MockMiddleware()
        chain = MiddlewareChain([test_middleware])
        context = AgentContext()

        async def failing_func(data):
            raise RuntimeError("Test error")

        with pytest.raises(RuntimeError, match="Test error"):
            async for item in chain.execute_stream(
                operation="test_op",
                agent_name="test_agent",
                agent_context=context,
                data="test_data",
                func=failing_func,
            ):
                pass

        assert test_middleware.error_calls == [("test_op", "RuntimeError")]


class TestAgentMiddlewareIntegration:
    """Test middleware integration with Agent."""

    @pytest.mark.asyncio
    async def test_agent_with_middleware(self, mock_agent):
        """Test agent executes middleware during operations."""
        agent, client, context = mock_agent
        test_middleware = MockMiddleware()
        agent.middleware_chain = MiddlewareChain([test_middleware])

        client.set_response("Test response")

        result = await agent.run("Test query")

        # Should have processed model call
        assert "model_call" in test_middleware.request_calls
        assert "model_call" in test_middleware.response_calls
        # The middleware modifies the response from the model, but that modified response
        # is used to create the AssistantMessage that gets added to context.
        # Check that the middleware was called (we already verified above)
        # The actual modification happens in the agent's internal processing,
        # and the AssistantMessage is created from the modified result.
        # Since messages are immutable and the agent creates new AssistantMessage from
        # the result, we just verify middleware was called correctly.
        assert len(test_middleware.request_calls) > 0
        assert len(test_middleware.response_calls) > 0

    @pytest.mark.asyncio
    async def test_agent_with_tool_middleware(self, mock_agent):
        """Test middleware processes tool calls."""
        agent, client, context = mock_agent
        test_middleware = MockMiddleware()
        agent.middleware_chain = MiddlewareChain([test_middleware])

        # Add a tool
        mock_tool = MockTool()
        agent.tools = [mock_tool]

        # Set response sequence: first call returns tool call, second returns final answer
        tool_call = ToolCallRequest(
            tool_name="mock_tool", parameters={}, call_id="test_call"
        )
        client.set_responses(
            [
                AssistantMessage(
                    content="I'll use the tool", source="mock", tool_calls=[tool_call]
                ),
                AssistantMessage(
                    content="Final answer after using tool",
                    source="mock",
                    tool_calls=None,
                ),
            ]
        )

        result = await agent.run("Test query")

        # Should have processed model call (tool call integration will be tested separately)
        assert "model_call" in test_middleware.request_calls
        assert len(test_middleware.request_calls) >= 1


class TestBuiltInMiddleware:
    """Test the built-in middleware implementations."""

    @pytest.mark.asyncio
    async def test_logging_middleware(self, caplog):
        """Test logging middleware logs operations."""
        import logging

        logger = logging.getLogger("test_middleware")
        logging_middleware = LoggingMiddleware(logger)

        context = MiddlewareContext(
            operation="test_op",
            agent_name="test_agent",
            agent_context=AgentContext(),
            data="test_data",
        )

        # Test request logging
        with caplog.at_level(logging.INFO):
            result_context = None
            async for item in logging_middleware.process_request(context):
                result_context = item
            assert "Starting test_op" in caplog.text

        # Test response logging
        with caplog.at_level(logging.INFO):
            async for item in logging_middleware.process_response(result_context, "test_result"):
                pass
            assert "Completed test_op" in caplog.text

    @pytest.mark.asyncio
    async def test_rate_limit_middleware(self):
        """Test rate limiting middleware."""
        # Set very low limit for testing
        rate_limiter = RateLimitMiddleware(max_calls_per_minute=2)
        context = MiddlewareContext(
            operation="test_op",
            agent_name="test_agent",
            agent_context=AgentContext(),
            data="test_data",
        )

        # First two calls should be fast
        start_time = time.time()
        async for item in rate_limiter.process_request(context):
            pass
        async for item in rate_limiter.process_request(context):
            pass
        _fast_duration = time.time() - start_time

        # Third call should be rate limited (but we'll use a small timeout for testing)
        rate_limiter.max_calls = 1  # Lower the limit to trigger rate limiting faster
        rate_limiter.call_times = [time.time()]  # Set one recent call

        # The next call should wait, but we'll set a very short time to avoid long tests
        # In real usage, this would wait up to 60 seconds
        rate_limiter.call_times = [time.time() - 59]  # Call 59 seconds ago

        start_time = time.time()
        async for item in rate_limiter.process_request(context):
            pass
        limited_duration = time.time() - start_time

        # Should be nearly instant since the call was 59 seconds ago
        assert limited_duration < 2

    @pytest.mark.asyncio
    async def test_pii_redaction_middleware(self):
        """Test PII redaction middleware."""
        pii_middleware = PIIRedactionMiddleware()

        # Test message with PII
        message = UserMessage(
            content="My email is john.doe@example.com and SSN is 123-45-6789",
            source="user",
        )

        context = MiddlewareContext(
            operation="model_call",
            agent_name="test_agent",
            agent_context=AgentContext(),
            data=[message],
        )

        processed_context = None
        async for item in pii_middleware.process_request(context):
            processed_context = item
        processed_message = processed_context.data[0]

        # PII should be redacted
        assert "[EMAIL-REDACTED]" in processed_message.content
        assert "[SSN-REDACTED]" in processed_message.content
        assert "john.doe@example.com" not in processed_message.content
        assert "123-45-6789" not in processed_message.content

    @pytest.mark.asyncio
    async def test_metrics_middleware(self):
        """Test metrics collection middleware."""
        metrics_middleware = MetricsMiddleware()

        context = MiddlewareContext(
            operation="model_call",
            agent_name="test_agent",
            agent_context=AgentContext(),
            data="test_data",
        )

        # Process request
        async for item in metrics_middleware.process_request(context):
            pass

        # Simulate some processing time
        await asyncio.sleep(0.01)

        # Process response
        async for item in metrics_middleware.process_response(context, "test_result"):
            pass

        metrics = metrics_middleware.get_metrics()

        assert metrics["total_operations"] == 1
        assert metrics["operations_by_type"]["model_call"] == 1
        assert metrics["total_duration"] > 0
        assert metrics["average_duration"] > 0

    @pytest.mark.asyncio
    async def test_guardrail_middleware(self):
        """Test guardrail middleware blocks dangerous operations."""
        guardrail = GuardrailMiddleware(
            blocked_tools=["dangerous_tool"], blocked_patterns=["rm -rf"]
        )

        # Test blocked tool
        tool_call = ToolCallRequest(
            tool_name="dangerous_tool", parameters={}, call_id="test_call"
        )

        context = MiddlewareContext(
            operation="tool_call",
            agent_name="test_agent",
            agent_context=AgentContext(),
            data=tool_call,
        )

        with pytest.raises(ValueError, match="blocked by guardrails"):
            async for item in guardrail.process_request(context):
                pass

        # Test blocked pattern in messages
        message = UserMessage(content="Please run: rm -rf /", source="user")
        context = MiddlewareContext(
            operation="model_call",
            agent_name="test_agent",
            agent_context=AgentContext(),
            data=[message],
        )

        with pytest.raises(ValueError, match="blocked pattern"):
            async for item in guardrail.process_request(context):
                pass


class TestAgentContext:
    """Test the AgentContext functionality."""

    def test_context_creation(self):
        """Test creating an empty context."""
        context = AgentContext()

        assert context.message_count == 0
        assert context.is_empty
        assert isinstance(context.messages, list)
        assert isinstance(context.metadata, dict)
        assert isinstance(context.shared_state, dict)

    def test_context_with_data(self):
        """Test creating context with initial data."""
        context = AgentContext(metadata={"user_id": "123"}, session_id="session456")

        assert context.metadata["user_id"] == "123"
        assert context.session_id == "session456"

    def test_add_message(self):
        """Test adding messages to context."""
        context = AgentContext()
        message = UserMessage(content="Test", source="user")

        context.add_message(message)

        assert context.message_count == 1
        assert not context.is_empty
        assert context.messages[0] == message

    def test_get_last_messages(self):
        """Test getting last messages by type."""
        context = AgentContext()

        user_msg = UserMessage(content="User message", source="user")
        assistant_msg = AssistantMessage(
            content="Assistant message", source="assistant"
        )

        context.add_message(user_msg)
        context.add_message(assistant_msg)

        assert context.get_last_user_message() == user_msg
        assert context.get_last_assistant_message() == assistant_msg

    def test_context_reset(self):
        """Test resetting context."""
        context = AgentContext(
            metadata={"key": "value"}, shared_state={"state": "data"}
        )

        context.add_message(UserMessage(content="Test", source="user"))
        context.reset()

        assert context.message_count == 0
        assert len(context.metadata) == 0
        assert len(context.shared_state) == 0

    def test_from_messages(self):
        """Test creating context from messages list."""
        messages = [
            UserMessage(content="Test 1", source="user"),
            AssistantMessage(content="Test 2", source="assistant"),
        ]

        context = AgentContext.from_messages(messages)

        assert context.message_count == 2
        assert context.messages == messages


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