"""Tests for Ollama client."""
import pytest
from unittest.mock import patch, AsyncMock, MagicMock
import json

from src.models.ollama_client import OllamaClient, get_ollama_client, close_ollama_client


class TestOllamaClientInit:
    """Test OllamaClient initialization."""

    @patch("src.models.ollama_client.settings")
    def test_uses_settings_defaults(self, mock_settings):
        """Client should use settings for defaults."""
        mock_settings.ollama_base_url = "http://ollama:11434"
        mock_settings.ollama_timeout = 60

        client = OllamaClient()

        assert client.base_url == "http://ollama:11434"
        assert client.timeout == 60

    @patch("src.models.ollama_client.settings")
    def test_creates_http_client(self, mock_settings):
        """Client should create httpx AsyncClient."""
        mock_settings.ollama_base_url = "http://ollama:11434"
        mock_settings.ollama_timeout = 30

        client = OllamaClient()

        assert client.client is not None


class TestOllamaClientClose:
    """Test client close functionality."""

    @pytest.mark.asyncio
    @patch("src.models.ollama_client.settings")
    async def test_close_closes_client(self, mock_settings):
        """close should close the HTTP client."""
        mock_settings.ollama_base_url = "http://ollama:11434"
        mock_settings.ollama_timeout = 30

        client = OllamaClient()

        with patch.object(client.client, "aclose", new_callable=AsyncMock) as mock_close:
            await client.close()
            mock_close.assert_called_once()


class TestOllamaClientResolveModel:
    """Test model resolution."""

    @patch("src.models.ollama_client.settings")
    def test_resolves_aliased_model(self, mock_settings):
        """resolve_model should map alias to actual model."""
        mock_settings.ollama_base_url = "http://ollama:11434"
        mock_settings.ollama_timeout = 30
        mock_settings.model_aliases = {"gpt-3.5-turbo": "gemma:7b"}

        client = OllamaClient()
        result = client.resolve_model("gpt-3.5-turbo")

        assert result == "gemma:7b"

    @patch("src.models.ollama_client.settings")
    def test_returns_original_if_no_alias(self, mock_settings):
        """resolve_model should return original if no alias found."""
        mock_settings.ollama_base_url = "http://ollama:11434"
        mock_settings.ollama_timeout = 30
        mock_settings.model_aliases = {}

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

        assert result == "llama2"


class TestOllamaClientHealthCheck:
    """Test health check functionality."""

    @pytest.mark.asyncio
    @patch("src.models.ollama_client.settings")
    async def test_health_check_returns_true_on_200(self, mock_settings):
        """Health check should return True when Ollama responds 200."""
        mock_settings.ollama_base_url = "http://ollama:11434"
        mock_settings.ollama_timeout = 30

        client = OllamaClient()

        mock_response = MagicMock()
        mock_response.status_code = 200

        with patch.object(client.client, "get", new_callable=AsyncMock) as mock_get:
            mock_get.return_value = mock_response
            result = await client.health_check()

            assert result is True
            mock_get.assert_called_once()

    @pytest.mark.asyncio
    @patch("src.models.ollama_client.settings")
    async def test_health_check_returns_false_on_error(self, mock_settings):
        """Health check should return False on connection error."""
        mock_settings.ollama_base_url = "http://ollama:11434"
        mock_settings.ollama_timeout = 30

        client = OllamaClient()

        with patch.object(client.client, "get", new_callable=AsyncMock) as mock_get:
            mock_get.side_effect = Exception("Connection refused")
            result = await client.health_check()

            assert result is False

    @pytest.mark.asyncio
    @patch("src.models.ollama_client.settings")
    async def test_health_check_returns_false_on_non_200(self, mock_settings):
        """Health check should return False on non-200 status."""
        mock_settings.ollama_base_url = "http://ollama:11434"
        mock_settings.ollama_timeout = 30

        client = OllamaClient()

        mock_response = MagicMock()
        mock_response.status_code = 500

        with patch.object(client.client, "get", new_callable=AsyncMock) as mock_get:
            mock_get.return_value = mock_response
            result = await client.health_check()

            assert result is False


class TestOllamaClientListModels:
    """Test list models functionality."""

    @pytest.mark.asyncio
    @patch("src.models.ollama_client.settings")
    async def test_list_models_returns_dict(self, mock_settings):
        """list_models should return dict with models."""
        mock_settings.ollama_base_url = "http://ollama:11434"
        mock_settings.ollama_timeout = 30

        client = OllamaClient()

        models_data = {
            "models": [
                {"name": "llama2", "size": 1000000},
                {"name": "gemma:7b", "size": 2000000}
            ]
        }

        mock_response = MagicMock()
        mock_response.status_code = 200
        mock_response.json.return_value = models_data
        mock_response.raise_for_status = MagicMock()

        with patch.object(client.client, "get", new_callable=AsyncMock) as mock_get:
            mock_get.return_value = mock_response
            result = await client.list_models()

            assert result == models_data
            assert len(result["models"]) == 2

    @pytest.mark.asyncio
    @patch("src.models.ollama_client.settings")
    async def test_list_models_raises_on_error(self, mock_settings):
        """list_models should raise on error."""
        mock_settings.ollama_base_url = "http://ollama:11434"
        mock_settings.ollama_timeout = 30

        client = OllamaClient()

        with patch.object(client.client, "get", new_callable=AsyncMock) as mock_get:
            mock_get.side_effect = Exception("Connection error")

            with pytest.raises(Exception):
                await client.list_models()


class TestOllamaClientGenerateNonStreaming:
    """Test non-streaming generation."""

    @pytest.mark.asyncio
    @patch("src.models.ollama_client.settings")
    async def test_generate_non_streaming_returns_response(self, mock_settings):
        """generate_non_streaming should return response dict."""
        mock_settings.ollama_base_url = "http://ollama:11434"
        mock_settings.ollama_timeout = 30
        mock_settings.model_aliases = {}

        client = OllamaClient()

        response_data = {
            "message": {"content": "Hello! How can I help?"},
            "prompt_eval_count": 10,
            "eval_count": 20
        }

        mock_response = MagicMock()
        mock_response.status_code = 200
        mock_response.json.return_value = response_data
        mock_response.raise_for_status = MagicMock()

        with patch.object(client.client, "post", new_callable=AsyncMock) as mock_post:
            mock_post.return_value = mock_response
            result = await client.generate_non_streaming("llama2", "Hello")

            assert result["response"] == "Hello! How can I help?"
            assert result["tokens"]["prompt"] == 10
            assert result["tokens"]["completion"] == 20
            assert result["tokens"]["total"] == 30

    @pytest.mark.asyncio
    @patch("src.models.ollama_client.settings")
    async def test_generate_non_streaming_includes_max_tokens(self, mock_settings):
        """generate_non_streaming should include max_tokens in payload."""
        mock_settings.ollama_base_url = "http://ollama:11434"
        mock_settings.ollama_timeout = 30
        mock_settings.model_aliases = {}

        client = OllamaClient()

        mock_response = MagicMock()
        mock_response.status_code = 200
        mock_response.json.return_value = {"message": {"content": "Hi"}}
        mock_response.raise_for_status = MagicMock()

        with patch.object(client.client, "post", new_callable=AsyncMock) as mock_post:
            mock_post.return_value = mock_response
            await client.generate_non_streaming("llama2", "Hello", max_tokens=100)

            call_args = mock_post.call_args
            assert call_args[1]["json"]["options"]["num_predict"] == 100


class TestOllamaClientGenerateStreaming:
    """Test streaming generation."""

    @pytest.mark.asyncio
    @patch("src.models.ollama_client.settings")
    async def test_generate_streaming_yields_content(self, mock_settings):
        """generate_streaming should yield content chunks."""
        mock_settings.ollama_base_url = "http://ollama:11434"
        mock_settings.ollama_timeout = 30
        mock_settings.model_aliases = {}

        client = OllamaClient()

        # Create mock streaming response
        async def mock_aiter_lines():
            yield json.dumps({"message": {"content": "Hello"}})
            yield json.dumps({"message": {"content": " world"}})
            yield json.dumps({"done": True})

        mock_response = MagicMock()
        mock_response.status_code = 200
        mock_response.raise_for_status = MagicMock()
        mock_response.aiter_lines = mock_aiter_lines

        mock_stream_context = AsyncMock()
        mock_stream_context.__aenter__.return_value = mock_response
        mock_stream_context.__aexit__.return_value = None

        with patch.object(client.client, "stream", return_value=mock_stream_context):
            chunks = []
            async for chunk in client.generate_streaming("llama2", "Hi"):
                chunks.append(chunk)

            assert "Hello" in chunks
            assert " world" in chunks


class TestOllamaClientSingleton:
    """Test singleton pattern."""

    @patch("src.models.ollama_client.settings")
    def test_get_ollama_client_returns_same_instance(self, mock_settings):
        """get_ollama_client should return singleton."""
        mock_settings.ollama_base_url = "http://ollama:11434"
        mock_settings.ollama_timeout = 30

        import src.models.ollama_client as module
        module._ollama_client = None

        client1 = get_ollama_client()
        client2 = get_ollama_client()

        assert client1 is client2

    @pytest.mark.asyncio
    @patch("src.models.ollama_client.settings")
    async def test_close_ollama_client_clears_singleton(self, mock_settings):
        """close_ollama_client should clear the singleton."""
        mock_settings.ollama_base_url = "http://ollama:11434"
        mock_settings.ollama_timeout = 30

        import src.models.ollama_client as module
        module._ollama_client = None

        client = get_ollama_client()

        with patch.object(client.client, "aclose", new_callable=AsyncMock):
            await close_ollama_client()

        assert module._ollama_client is None