from contextlib import asynccontextmanager
from unittest.mock import AsyncMock, patch

import httpx
import pytest

from haiku.rag.client.downloads import download_models
from haiku.rag.config import Config


@pytest.fixture
def mock_to_thread():
    """Patch asyncio.to_thread to skip docling/tokenizer downloads."""
    with patch("haiku.rag.client.downloads.asyncio.to_thread", new_callable=AsyncMock):
        yield


@asynccontextmanager
async def _mock_httpx_client(stream_fn):
    """Create a mock httpx.AsyncClient context manager with a given stream function."""
    mock_client = AsyncMock()
    mock_client.stream = stream_fn
    yield mock_client


async def test_download_models_ollama_connect_error(mock_to_thread):
    """When Ollama is not running, download_models raises ConnectionError."""

    @asynccontextmanager
    async def failing_stream(method, url, **kwargs):
        raise httpx.ConnectError("All connection attempts failed")
        yield  # unreachable, but needed for generator syntax

    with patch(
        "haiku.rag.client.downloads.httpx.AsyncClient",
        return_value=_mock_httpx_client(failing_stream),
    ):
        with pytest.raises(
            ConnectionError, match="Cannot connect to Ollama"
        ) as exc_info:
            async for _ in download_models(Config):
                pass

        assert "ollama serve" in str(exc_info.value)


async def test_download_models_ollama_pulls_models(mock_to_thread):
    """download_models yields correct progress events for Ollama model pulls."""
    stream_lines = [
        '{"status": "pulling manifest"}',
        "",
        '{"status": "downloading", "digest": "sha256:abc", "total": 1000, "completed": 500}',
        '{"status": "downloading", "digest": "sha256:abc", "total": 1000, "completed": 1000}',
        "not valid json",
        '{"status": "verifying sha256 digest"}',
        '{"status": "writing manifest"}',
        '{"status": "success"}',
    ]

    @asynccontextmanager
    async def mock_stream(method, url, **kwargs):
        mock_resp = AsyncMock()

        async def aiter_lines():
            for line in stream_lines:
                yield line

        mock_resp.aiter_lines = aiter_lines
        yield mock_resp

    with patch(
        "haiku.rag.client.downloads.httpx.AsyncClient",
        return_value=_mock_httpx_client(mock_stream),
    ):
        events = []
        async for progress in download_models(Config):
            events.append(progress)

    # Default config has embeddings=qwen3-embedding:4b, qa/research=gpt-oss
    ollama_models = {"gpt-oss", "qwen3-embedding:4b"}
    ollama_events = [e for e in events if e.model in ollama_models]
    pulling_events = [e for e in ollama_events if e.status == "pulling"]
    done_events = [e for e in ollama_events if e.status == "done"]
    download_events = [e for e in ollama_events if e.status == "downloading"]

    assert len(pulling_events) == 2
    assert len(done_events) == 2
    assert len(download_events) > 0

    for de in download_events:
        assert de.digest == "sha256:abc"
        assert de.total == 1000
        assert de.completed > 0


async def test_download_models_no_ollama_models(mock_to_thread):
    """When no Ollama models are configured, no Ollama pull events are yielded."""
    from haiku.rag.config import AppConfig

    config = AppConfig()
    config.embeddings.model.provider = "openai"
    config.qa.model.provider = "openai"
    config.research.model.provider = "openai"

    events = []
    async for progress in download_models(config):
        events.append(progress)

    models = {e.model for e in events}
    assert "qwen3-embedding:4b" not in models
    assert "gpt-oss" not in models