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

from unittest.mock import Mock, patch, MagicMock

from shellgenius.ollama_client import OllamaClient, get_ollama_client


class TestOllamaClient:
    """Test Ollama client functionality."""

    @patch('shellgenius.ollama_client.ollama.Client')
    def test_client_initialization(self, mock_client_class):
        """Test client initialization."""
        mock_client = Mock()
        mock_client_class.return_value = mock_client

        client = OllamaClient(host="localhost:11434", model="codellama")

        assert client.host == "localhost:11434"
        assert client.model == "codellama"

    @patch('shellgenius.ollama_client.ollama.Client')
    def test_list_models(self, mock_client_class):
        """Test listing models."""
        mock_client = Mock()
        mock_client.list.return_value = {
            "models": [
                {"name": "codellama:latest"},
                {"name": "llama2:latest"},
            ]
        }
        mock_client_class.return_value = mock_client

        client = OllamaClient()
        models = client.list_models()

        assert len(models) == 2
        assert "codellama:latest" in models

    @patch('shellgenius.ollama_client.ollama.Client')
    def test_generate_success(self, mock_client_class):
        """Test successful generation."""
        mock_client = Mock()
        mock_client.generate.return_value = {
            "response": "echo hello"
        }
        mock_client_class.return_value = mock_client

        client = OllamaClient()
        result = client.generate("test prompt")

        assert result["success"] is True
        assert result["response"]["response"] == "echo hello"

    @patch('shellgenius.ollama_client.ollama.Client')
    def test_generate_failure(self, mock_client_class):
        """Test failed generation."""
        mock_client = Mock()
        mock_client.generate.side_effect = Exception("Connection failed")
        mock_client_class.return_value = mock_client

        client = OllamaClient()
        result = client.generate("test prompt")

        assert result["success"] is False
        assert "error" in result

    @patch('shellgenius.ollama_client.ollama.Client')
    def test_chat_success(self, mock_client_class):
        """Test successful chat."""
        mock_client = Mock()
        mock_client.chat.return_value = {
            "message": {"content": "Hello!"}
        }
        mock_client_class.return_value = mock_client

        client = OllamaClient()
        messages = [{"role": "user", "content": "Hi"}]
        result = client.chat(messages)

        assert result["success"] is True

    @patch('shellgenius.ollama_client.ollama.Client')
    def test_pull_model(self, mock_client_class):
        """Test pulling a model."""
        mock_client = Mock()
        mock_client.pull.return_value = {"status": "success"}
        mock_client_class.return_value = mock_client

        client = OllamaClient()
        result = client.pull_model("llama2")

        assert result is True
        mock_client.pull.assert_called_once_with("llama2")

    @patch('shellgenius.ollama_client.ollama.Client')
    def test_is_available_true(self, mock_client_class):
        """Test availability check when available."""
        mock_client = Mock()
        mock_client.list.return_value = {"models": []}
        mock_client_class.return_value = mock_client

        client = OllamaClient()
        assert client.is_available() is True

    @patch('shellgenius.ollama_client.ollama.Client')
    def test_is_available_false(self, mock_client_class):
        """Test availability check when not available."""
        mock_instance = MagicMock()
        mock_instance.list.side_effect = Exception("Connection refused")
        mock_client_class.return_value = mock_instance

        with patch.object(OllamaClient, 'list_models', side_effect=Exception("Connection refused")):
            client = OllamaClient()
            result = client.is_available()
            assert result is False


class TestGetOllamaClient:
    """Test get_ollama_client convenience function."""

    @patch('shellgenius.ollama_client.ollama.Client')
    def test_get_client(self, mock_client_class):
        """Test getting client instance."""
        mock_client = Mock()
        mock_client_class.return_value = mock_client

        client = get_ollama_client(host="localhost:11434")

        assert isinstance(client, OllamaClient)
        assert client.host == "localhost:11434"