"""
Mock LLM provider for testing with integrated tracing.
"""
import asyncio
import logging
from typing import Dict, Any, Optional

from .base import BaseLLMProvider

logger = logging.getLogger(__name__)

class MockLLMProvider(BaseLLMProvider):
    """Mock LLM provider for testing purposes with automatic call tracing."""
    
    def __init__(
        self,
        model: str = "mock-model",
        responses: Optional[Dict[str, str]] = None,
        delay: float = 0.1,
        **kwargs
    ):
        """
        Initialize mock provider.
        
        Args:
            model: Mock model name
            responses: Dictionary mapping prompts to responses
            delay: Artificial delay to simulate API calls
            **kwargs: Additional parameters
        """
        # Remove api_key from kwargs to avoid conflict, then pass it explicitly
        kwargs.pop('api_key', None)  # Remove if exists
        super().__init__(model=model, api_key="mock-key", **kwargs)
        
        # Predefined responses
        self.responses = responses or {}
        self.delay = delay
        
        # Default responses
        self.default_responses = {
            "default": "This is a mock response from the LLM.",
            "analyze": "Based on the data provided, here are the key insights: [mock analysis]",
            "summarize": "Summary: [mock summary of the content]",
            "error": "This is an error response for testing."
        }
        
        # Track calls for testing
        self.call_history = []
    
    async def _generate_impl(self, prompt: str, **kwargs) -> str:
        """
        Provider-specific implementation of mock text generation.
        
        This method contains the mock generation logic and is automatically
        wrapped with tracing by the base class generate() method.
        
        Args:
            prompt: Input prompt
            **kwargs: Optional parameters
            
        Returns:
            Mock response
        """
        # Record the call
        self.call_history.append({
            'prompt': prompt,
            'params': kwargs,
            'timestamp': asyncio.get_event_loop().time()
        })
        
        # Simulate API delay
        if self.delay > 0:
            await asyncio.sleep(self.delay)
        
        # Check for specific response
        if prompt in self.responses:
            return self.responses[prompt]
        
        # Check for keyword-based responses
        prompt_lower = prompt.lower()
        for keyword, response in self.default_responses.items():
            if keyword in prompt_lower:
                return response
        
        # Default response
        return f"Mock response for: {prompt[:50]}..."
    
    async def _generate_with_tools_single(
        self,
        messages: list,
        tools: list,
        **kwargs
    ) -> dict:
        """
        Mock implementation of single LLM call with tools.

        Returns a mock response without tool calls (final answer).
        Override this in tests if you need to test tool calling behavior.
        """
        # Extract the last user message
        user_message = ""
        for msg in reversed(messages):
            if msg.get("role") == "user":
                user_message = msg.get("content", "")
                break

        # Record the call
        self.call_history.append({
            'messages': messages,
            'tools': tools,
            'params': kwargs,
            'timestamp': asyncio.get_event_loop().time()
        })

        # Simulate API delay
        if self.delay > 0:
            await asyncio.sleep(self.delay)

        # Return final answer (no tool calls for basic mock)
        return {
            "content": f"Mock response for: {user_message[:50]}...",
            "tool_calls": None
        }

    def _get_last_token_usage(self) -> Dict[str, int]:
        """
        Override base class method to return mock token usage.

        Provides realistic but fake token counts for testing.
        """
        if self.call_history:
            # Get the last call to estimate tokens
            last_call = self.call_history[-1]
            prompt = last_call.get('prompt', '')

            # Mock realistic token counts
            estimated_prompt_tokens = max(5, len(prompt) // 4)  # Rough estimate
            # Assume a moderate response length for mocking
            estimated_completion_tokens = max(10, estimated_prompt_tokens // 2)

            return {
                'total_tokens': estimated_prompt_tokens + estimated_completion_tokens,
                'prompt_tokens': estimated_prompt_tokens,
                'completion_tokens': estimated_completion_tokens
            }

        # Fallback to default
        return super()._get_last_token_usage()
    
    def set_response(self, prompt: str, response: str) -> None:
        """Set a specific response for a prompt."""
        self.responses[prompt] = response
    
    def clear_history(self) -> None:
        """Clear call history."""
        self.call_history.clear()
    
    def get_last_call(self) -> Optional[Dict[str, Any]]:
        """Get the last call made to the provider."""
        return self.call_history[-1] if self.call_history else None
    
    @property
    def info(self) -> Dict[str, Any]:
        """Get information about the mock provider."""
        base_info = super().info
        base_info.update({
            'provider_name': 'Mock LLM (Testing)',
            'call_count': len(self.call_history),
            'configured_responses': len(self.responses),
            'delay_seconds': self.delay
        })
        return base_info