"""Tests for Ollama client module."""

from unittest.mock import patch, MagicMock

from shellgenius.ollama_client import OllamaClient, get_ollama_client


class TestOllamaClient:
    def test_init(self):
        """Test client initialization."""
        mock_config = MagicMock()
        mock_config.get.side_effect = lambda key, default=None: {
            "ollama_host": "localhost:11434",
            "ollama_model": "codellama",
        }.get(key, default)
        
        with patch('shellgenius.ollama_client.get_config', return_value=mock_config):
            client = OllamaClient()
            
            assert client.host == "localhost:11434"
            assert client.model == "codellama"
    
    def test_is_available(self):
        """Test availability check."""
        mock_config = MagicMock()
        mock_config.get.side_effect = lambda key, default=None: {
            "ollama_host": "localhost:11434",
            "ollama_model": "codellama",
        }.get(key, default)
        
        with patch('shellgenius.ollama_client.get_config', return_value=mock_config):
            client = OllamaClient()
            
            with patch.object(client, 'list_models', return_value=["codellama"]):
                assert client.is_available() is True
    
    def test_list_models(self):
        """Test listing models."""
        mock_config = MagicMock()
        mock_config.get.side_effect = lambda key, default=None: {
            "ollama_host": "localhost:11434",
            "ollama_model": "codellama",
        }.get(key, default)
        
        with patch('shellgenius.ollama_client.get_config', return_value=mock_config):
            client = OllamaClient()
            
            mock_response = {"models": [{"name": "codellama"}, {"name": "llama2"}]}
            
            with patch.object(client.client, 'list', return_value=mock_response):
                models = client.list_models()
                
                assert len(models) == 2
                assert "codellama" in models
    
    def test_generate(self):
        """Test text generation."""
        mock_config = MagicMock()
        mock_config.get.side_effect = lambda key, default=None: {
            "ollama_host": "localhost:11434",
            "ollama_model": "codellama",
        }.get(key, default)
        
        with patch('shellgenius.ollama_client.get_config', return_value=mock_config):
            client = OllamaClient()
            
            mock_response = {"response": "Generated text"}
            
            with patch.object(client.client, 'generate', return_value=mock_response):
                result = client.generate("test prompt")
                
                assert result["success"] is True
                assert "Generated text" in str(result["response"])


class TestGetOllamaClient:
    def test_convenience_function(self):
        """Test the convenience function for getting client."""
        with patch('shellgenius.ollama_client.get_config') as mock_config:
            mock_config.return_value = {}
            
            client = get_ollama_client(host="custom:9999", model="custom-model")
            
            assert client.host == "custom:9999"
            assert client.model == "custom-model"