"""Mock LLM implementations for testing without API calls."""

from typing import List, Optional, Any
from langchain_core.messages import AIMessage, BaseMessage
from langchain_core.language_models import BaseChatModel
from langchain_core.outputs import ChatGeneration, ChatResult
import pytest


class MockLLM(BaseChatModel):
    """
    Mock LLM for testing without making actual API calls.
    
    This mock can be configured with predefined responses and tracks
    the number of times it's been called.
    
    Example:
        >>> mock = MockLLM(responses=["Hello!", "How can I help?"])
        >>> response = await mock.ainvoke([HumanMessage(content="Hi")])
        >>> assert response.content == "Hello!"
    """
    
    responses: List[str]
    call_count: int = 0
    
    def __init__(self, responses: Optional[List[str]] = None, **kwargs):
        """
        Initialize mock LLM with predefined responses.
        
        Args:
            responses: List of responses to return in sequence.
                      If None, returns "Mock response" for all calls.
        """
        super().__init__(**kwargs)
        self.responses = responses or ["Mock response"]
        self.call_count = 0
    
    @property
    def _llm_type(self) -> str:
        """Return identifier for this LLM type."""
        return "mock"
    
    def _generate(
        self,
        messages: List[BaseMessage],
        stop: Optional[List[str]] = None,
        **kwargs: Any,
    ) -> ChatResult:
        """Generate a mock response synchronously."""
        response = self.responses[self.call_count % len(self.responses)]
        self.call_count += 1
        
        message = AIMessage(content=response)
        generation = ChatGeneration(message=message)
        
        return ChatResult(generations=[generation])
    
    async def _agenerate(
        self,
        messages: List[BaseMessage],
        stop: Optional[List[str]] = None,
        **kwargs: Any,
    ) -> ChatResult:
        """Generate a mock response asynchronously."""
        return self._generate(messages, stop, **kwargs)
    
    def reset(self):
        """Reset call count for reuse in tests."""
        self.call_count = 0


class MockLLMWithError(BaseChatModel):
    """
    Mock LLM that raises errors for testing error handling.
    
    Example:
        >>> mock = MockLLMWithError(error_message="API timeout")
        >>> with pytest.raises(Exception):
        ...     await mock.ainvoke([HumanMessage(content="Hi")])
    """
    
    error_message: str
    
    def __init__(self, error_message: str = "Mock LLM error", **kwargs):
        """
        Initialize mock LLM that raises errors.
        
        Args:
            error_message: Error message to raise
        """
        super().__init__(**kwargs)
        self.error_message = error_message
    
    @property
    def _llm_type(self) -> str:
        """Return identifier for this LLM type."""
        return "mock_error"
    
    def _generate(
        self,
        messages: List[BaseMessage],
        stop: Optional[List[str]] = None,
        **kwargs: Any,
    ) -> ChatResult:
        """Raise an error."""
        raise Exception(self.error_message)
    
    async def _agenerate(
        self,
        messages: List[BaseMessage],
        stop: Optional[List[str]] = None,
        **kwargs: Any,
    ) -> ChatResult:
        """Raise an error asynchronously."""
        raise Exception(self.error_message)


@pytest.fixture
def mock_llm():
    """
    Pytest fixture providing a basic mock LLM.
    
    Returns:
        MockLLM instance with default response
    """
    return MockLLM()


@pytest.fixture
def mock_llm_with_responses():
    """
    Pytest fixture factory for creating mock LLM with custom responses.
    
    Returns:
        Function that creates MockLLM with specified responses
        
    Example:
        def test_conversation(mock_llm_with_responses):
            llm = mock_llm_with_responses(["Hi!", "Goodbye!"])
            # Use llm in test
    """
    def _create_mock(responses: List[str]) -> MockLLM:
        return MockLLM(responses=responses)
    return _create_mock


@pytest.fixture
def mock_llm_with_error():
    """
    Pytest fixture providing a mock LLM that raises errors.
    
    Returns:
        MockLLMWithError instance
    """
    return MockLLMWithError()