import importlib
import pytest
from unittest.mock import patch
from os import getenv

from garak.attempt import Message, Turn, Conversation
from garak.exception import BadGeneratorException
from garak.generators.litellm import LiteLLMGenerator


@pytest.mark.skipif(
    not all(
        [importlib.util.find_spec(m) for m in LiteLLMGenerator.extra_dependency_names]
    ),
    reason="missing optional dependency",
)
@pytest.mark.skipif(
    getenv("OPENAI_API_KEY", None) is None,
    reason="OpenAI API key is not set in OPENAI_API_KEY",
)
def test_litellm_openai():
    target_name = "gpt-3.5-turbo"
    generator = LiteLLMGenerator(name=target_name)
    assert generator.name == target_name
    assert isinstance(generator.max_tokens, int)

    output = generator.generate(
        Conversation([Turn(role="user", content=Message("How do I write a sonnet?"))])
    )
    assert len(output) == 1  # expect 1 generation by default

    for item in output:
        assert isinstance(item, Message)


@pytest.mark.skipif(
    not all(
        [importlib.util.find_spec(m) for m in LiteLLMGenerator.extra_dependency_names]
    ),
    reason="missing optional dependency",
)
@pytest.mark.skipif(
    getenv("OPENROUTER_API_KEY", None) is None,
    reason="OpenRouter API key is not set in OPENROUTER_API_KEY",
)
def test_litellm_openrouter():
    target_name = "openrouter/google/gemma-7b-it"
    generator = LiteLLMGenerator(name=target_name)
    assert generator.name == target_name
    assert isinstance(generator.max_tokens, int)

    output = generator.generate(Message("How do I write a sonnet?"))
    assert len(output) == 1  # expect 1 generation by default

    for item in output:
        assert isinstance(item, Message)


@pytest.mark.skipif(
    not all(
        [importlib.util.find_spec(m) for m in LiteLLMGenerator.extra_dependency_names]
    ),
    reason="missing optional dependency",
)
def test_litellm_model_detection():
    custom_config = {
        "generators": {
            "litellm": {
                "api_base": "https://garak.example.com/v1",
            }
        }
    }
    target_name = "non-existent-model"
    generator = LiteLLMGenerator(name=target_name, config_root=custom_config)
    conv = Conversation(
        [Turn("user", Message("This should raise a BadGeneratorException"))]
    )
    with pytest.raises(BadGeneratorException):
        generator.generate(conv)
    generator = LiteLLMGenerator(name="openai/invalid-model", config_root=custom_config)
    with pytest.raises(BadGeneratorException):
        generator.generate(conv)


def test_litellm_suppressed_params():
    """Test that suppressed_params configuration works correctly."""
    # Create a mock response object that matches what litellm.completion returns
    mock_response = type('obj', (object,), {
        'choices': [
            type('obj', (object,), {
                'message': type('obj', (object,), {
                    'content': 'Mock response'
                })
            })
        ]
    })
    
    target_name = "gpt-4"
    
    # Test case 1: Suppress top_p parameter
    with patch('litellm.completion', return_value=mock_response) as mock_completion:
        custom_config = {
            "generators": {
                "litellm": {
                    "suppressed_params": ["top_p"]
                }
            }
        }
        generator = LiteLLMGenerator(name=target_name, config_root=custom_config)
        conv = Conversation([Turn("user", Message("Test message"))])
        generator.generate(conv)
        
        # Check that top_p was not included but temperature was
        args, kwargs = mock_completion.call_args
        assert 'temperature' in kwargs
        assert 'top_p' not in kwargs
        assert 'max_tokens' in kwargs
    
    # Test case 2: Suppress multiple parameters
    with patch('litellm.completion', return_value=mock_response) as mock_completion:
        custom_config = {
            "generators": {
                "litellm": {
                    "suppressed_params": ["top_p", "frequency_penalty", "presence_penalty"]
                }
            }
        }
        generator = LiteLLMGenerator(name=target_name, config_root=custom_config)
        conv = Conversation([Turn("user", Message("Test message"))])
        generator.generate(conv)
        
        # Check that suppressed params were not included
        args, kwargs = mock_completion.call_args
        assert 'temperature' in kwargs
        assert 'top_p' not in kwargs
        assert 'frequency_penalty' not in kwargs
        assert 'presence_penalty' not in kwargs
        assert 'max_tokens' in kwargs
        
    # Test case 3: No suppressed params (default behavior)
    with patch('litellm.completion', return_value=mock_response) as mock_completion:
        generator = LiteLLMGenerator(name=target_name)
        conv = Conversation([Turn("user", Message("Test message"))])
        generator.generate(conv)
        
        # Check that all standard parameters were included
        args, kwargs = mock_completion.call_args
        assert 'temperature' in kwargs
        assert 'top_p' in kwargs
        assert 'frequency_penalty' in kwargs
        assert 'presence_penalty' in kwargs
        assert 'max_tokens' in kwargs
