"""
Mock LLM for Testing TestGen AI.

This module provides mock LLM responses for testing without making
actual API calls.
"""

from typing import Optional, Dict, Any, List
from dataclasses import dataclass
import re


@dataclass
class MockResponse:
    """Mock LLM response."""
    
    content: str
    model: str = "mock-model"
    prompt_tokens: int = 0
    completion_tokens: int = 0
    
    @property
    def total_tokens(self) -> int:
        """Total tokens used."""
        return self.prompt_tokens + self.completion_tokens


class MockLLM:
    """
    Mock LLM for testing without API calls.
    
    Implements Task 45 requirements:
    - Mock LiteLLM responses
    - Test prompt construction
    - Test response parsing
    
    Example:
        >>> mock = MockLLM()
        >>> response = mock.generate("def add(a, b): return a + b")
        >>> assert "def test_" in response.content
    """
    
    def __init__(
        self,
        model: str = "mock-gpt-4",
        default_response: Optional[str] = None
    ):
        """
        Initialize mock LLM.
        
        Args:
            model: Model name to simulate
            default_response: Default response template
        """
        self.model = model
        self.default_response = default_response or self._get_default_test_template()
        self.call_count = 0
        self.call_history: List[Dict[str, Any]] = []
    
    def generate(
        self,
        prompt: str,
        max_tokens: int = 1000,
        temperature: float = 0.7,
        **kwargs
    ) -> MockResponse:
        """
        Generate mock response (Task 45 requirement).
        
        Args:
            prompt: Input prompt
            max_tokens: Maximum tokens
            temperature: Temperature parameter
            **kwargs: Additional parameters
            
        Returns:
            MockResponse with generated content
            
        Example:
            >>> mock = MockLLM()
            >>> response = mock.generate("test prompt")
            >>> print(response.content)
        """
        self.call_count += 1
        
        # Record call
        self.call_history.append({
            "prompt": prompt,
            "max_tokens": max_tokens,
            "temperature": temperature,
            "kwargs": kwargs
        })
        
        # Generate response based on prompt
        content = self._generate_response(prompt)
        
        # Estimate tokens
        prompt_tokens = len(prompt.split()) * 1.3  # Rough estimate
        completion_tokens = len(content.split()) * 1.3
        
        return MockResponse(
            content=content,
            model=self.model,
            prompt_tokens=int(prompt_tokens),
            completion_tokens=int(completion_tokens)
        )
    
    def generate_test(
        self,
        source_code: str,
        language: str = "python",
        file_path: Optional[str] = None,
        **kwargs
    ) -> str:
        """
        Generate test code for given source code (WorkflowManager compatibility).
        
        Args:
            source_code: Source code to generate tests for
            language: Programming language
            file_path: Path to source file
            **kwargs: Additional arguments
            
        Returns:
            Generated test code as string
        """
        # Create prompt for test generation
        prompt = f"""Generate unit tests for the following {language} code:

```{language}
{source_code}
```

File: {file_path or 'unknown'}

Generate comprehensive unit tests with good coverage."""
        
        # Call generate and extract content
        response = self.generate(prompt, **kwargs)
        return response.content
    
    def _generate_response(self, prompt: str) -> str:
        """
        Generate appropriate response based on prompt.
        
        Args:
            prompt: Input prompt
            
        Returns:
            Generated test code
        """
        # Extract function name if present
        function_match = re.search(r'def (\w+)\(', prompt)
        
        if function_match:
            func_name = function_match.group(1)
            return self._generate_test_for_function(func_name, prompt)
        
        # Check for class
        class_match = re.search(r'class (\w+)', prompt)
        
        if class_match:
            class_name = class_match.group(1)
            return self._generate_test_for_class(class_name)
        
        # Default response
        return self.default_response
    
    def _generate_test_for_function(self, func_name: str, prompt: str) -> str:
        """Generate test for a function."""
        # Check if it's a simple operation
        if 'add' in func_name.lower() or '+' in prompt:
            return f'''import pytest

def test_{func_name}_basic():
    """Test basic {func_name} functionality."""
    result = {func_name}(2, 3)
    assert result == 5

def test_{func_name}_negative():
    """Test {func_name} with negative numbers."""
    result = {func_name}(-1, -1)
    assert result == -2

def test_{func_name}_zero():
    """Test {func_name} with zero."""
    result = {func_name}(0, 5)
    assert result == 5
'''
        
        elif 'subtract' in func_name.lower() or '-' in prompt:
            return f'''import pytest

def test_{func_name}_basic():
    """Test basic {func_name} functionality."""
    result = {func_name}(5, 3)
    assert result == 2

def test_{func_name}_negative():
    """Test {func_name} with negative result."""
    result = {func_name}(3, 5)
    assert result == -2
'''
        
        # Generic test
        return f'''import pytest

def test_{func_name}_basic():
    """Test {func_name} function."""
    result = {func_name}()
    assert result is not None

def test_{func_name}_with_args():
    """Test {func_name} with arguments."""
    # Add your test logic here
    pass
'''
    
    def _generate_test_for_class(self, class_name: str) -> str:
        """Generate test for a class."""
        return f'''import pytest

class Test{class_name}:
    """Tests for {class_name} class."""
    
    @pytest.fixture
    def instance(self):
        """Create {class_name} instance."""
        return {class_name}()
    
    def test_initialization(self, instance):
        """Test {class_name} initialization."""
        assert instance is not None
    
    def test_basic_functionality(self, instance):
        """Test basic functionality."""
        # Add your test logic here
        pass
'''
    
    def _get_default_test_template(self) -> str:
        """Get default test template."""
        return '''import pytest

def test_example():
    """Example test."""
    assert True

def test_another_example():
    """Another example test."""
    result = 1 + 1
    assert result == 2
'''
    
    def reset(self):
        """Reset mock state."""
        self.call_count = 0
        self.call_history = []
    
    def get_statistics(self) -> Dict[str, Any]:
        """Get mock usage statistics."""
        total_prompt_tokens = sum(
            len(call["prompt"].split()) * 1.3 
            for call in self.call_history
        )
        
        return {
            "total_calls": self.call_count,
            "total_prompt_tokens": int(total_prompt_tokens),
            "average_prompt_length": (
                sum(len(call["prompt"]) for call in self.call_history) / self.call_count
                if self.call_count > 0 else 0
            ),
            "model": self.model
        }


class MockLiteLLM:
    """
    Mock for LiteLLM completion (Task 45 requirement).
    
    Simulates litellm.completion() interface.
    
    Example:
        >>> import testgen.core.mock_llm as mock_llm
        >>> mock_llm.install_mock()
        >>> # Now litellm.completion() will use mock
    """
    
    def __init__(self):
        """Initialize mock LiteLLM."""
        self.mock = MockLLM()
    
    def completion(
        self,
        model: str,
        messages: List[Dict[str, str]],
        max_tokens: Optional[int] = None,
        temperature: Optional[float] = None,
        **kwargs
    ) -> Any:
        """
        Mock completion method.
        
        Args:
            model: Model name
            messages: Chat messages
            max_tokens: Max tokens
            temperature: Temperature
            **kwargs: Additional args
            
        Returns:
            Mock response object
        """
        # Extract prompt from messages
        prompt = "\n".join(msg.get("content", "") for msg in messages)
        
        # Generate response
        response = self.mock.generate(
            prompt,
            max_tokens=max_tokens or 1000,
            temperature=temperature or 0.7,
            **kwargs
        )
        
        # Create mock response object matching LiteLLM format
        class Choice:
            def __init__(self, content):
                self.message = type('Message', (), {'content': content})()
        
        class Usage:
            def __init__(self, prompt_tokens, completion_tokens):
                self.prompt_tokens = prompt_tokens
                self.completion_tokens = completion_tokens
                self.total_tokens = prompt_tokens + completion_tokens
        
        class Response:
            def __init__(self, content, prompt_tokens, completion_tokens, model):
                self.choices = [Choice(content)]
                self.usage = Usage(prompt_tokens, completion_tokens)
                self.model = model
        
        return Response(
            response.content,
            response.prompt_tokens,
            response.completion_tokens,
            model
        )


# Global mock instance
_mock_llm = MockLiteLLM()


def get_mock() -> MockLLM:
    """
    Get mock LLM instance.
    
    Returns:
        MockLLM instance
        
    Example:
        >>> mock = get_mock()
        >>> mock.reset()
    """
    return _mock_llm.mock


def install_mock():
    """
    Install mock LLM (replaces litellm).
    
    Example:
        >>> from testgen.core import mock_llm
        >>> mock_llm.install_mock()
        >>> # Now all LLM calls use mock
    """
    global _mock_llm
    # In real implementation, this would patch litellm
    return _mock_llm


def uninstall_mock():
    """Uninstall mock LLM."""
    global _mock_llm
    _mock_llm = MockLiteLLM()
