"""Tests for ask_docs tool schema and executor.

Covers:
- Tool schema (name, required params, enum values)
- Executor with mocked HTTP: success, auth error, API error, timeout
"""
from __future__ import annotations

import json
from unittest.mock import AsyncMock, MagicMock, patch

import pytest

from wafer_core.rollouts.dtypes import ToolCall
from wafer_core.tools.ask_docs_tool import (
    ASK_DOCS_TOOL,
    VALID_CORPORA,
    exec_ask_docs,
)


class TestAskDocsToolSchema:
    """Tests for ASK_DOCS_TOOL schema definition."""

    def test_tool_name(self) -> None:
        assert ASK_DOCS_TOOL.function.name == "ask_docs"

    def test_required_params(self) -> None:
        assert "query" in ASK_DOCS_TOOL.function.required
        assert "corpus" in ASK_DOCS_TOOL.function.required

    def test_corpus_enum_values(self) -> None:
        corpus_prop = ASK_DOCS_TOOL.function.parameters.properties["corpus"]
        enum_values = corpus_prop["enum"]
        assert isinstance(enum_values, list)
        assert len(enum_values) == len(VALID_CORPORA)
        for name in VALID_CORPORA:
            assert name in enum_values

    def test_query_is_string(self) -> None:
        query_prop = ASK_DOCS_TOOL.function.parameters.properties["query"]
        assert query_prop["type"] == "string"


class TestExecAskDocsValidation:
    """Tests for exec_ask_docs input validation."""

    @pytest.mark.anyio
    async def test_empty_query_returns_error(self) -> None:
        tc = ToolCall(id="test-1", name="ask_docs", args={"query": "", "corpus": "cuda"})
        result = await exec_ask_docs(tc)
        assert result.is_error
        assert "query" in result.error

    @pytest.mark.anyio
    async def test_invalid_corpus_returns_error(self) -> None:
        tc = ToolCall(id="test-2", name="ask_docs", args={"query": "test", "corpus": "invalid"})
        result = await exec_ask_docs(tc)
        assert result.is_error
        assert "Unknown corpus" in result.error

    @pytest.mark.anyio
    async def test_missing_api_url_returns_error(self) -> None:
        """When WAFER_API_URL is not set, tool returns error."""
        with patch.dict("os.environ", {}, clear=True):
            tc = ToolCall(id="test-3", name="ask_docs", args={"query": "test", "corpus": "cuda"})
            result = await exec_ask_docs(tc)
            assert result.is_error
            assert "WAFER_API_URL" in result.error

    @pytest.mark.anyio
    async def test_missing_auth_token_returns_error(self) -> None:
        with patch.dict("os.environ", {"WAFER_API_URL": "http://test"}, clear=True):
            tc = ToolCall(id="test-4", name="ask_docs", args={"query": "test", "corpus": "cuda"})
            result = await exec_ask_docs(tc)
            assert result.is_error
            assert "WAFER_AUTH_TOKEN" in result.error


class TestExecAskDocsHTTP:
    """Tests for exec_ask_docs with mocked HTTP responses."""

    @pytest.mark.anyio
    async def test_success_accumulates_sse(self) -> None:
        """Successful SSE stream should be accumulated into content."""
        sse_lines = [
            'data: {"type": "text", "content": "Hello "}',
            'data: {"type": "text", "content": "world"}',
            'data: {"type": "done"}',
        ]

        mock_response = AsyncMock()
        mock_response.status_code = 200
        mock_response.aiter_lines = MagicMock(return_value=_async_iter(sse_lines))
        mock_response.__aenter__ = AsyncMock(return_value=mock_response)
        mock_response.__aexit__ = AsyncMock(return_value=False)

        mock_client = AsyncMock()
        mock_client.stream = MagicMock(return_value=mock_response)
        mock_client.__aenter__ = AsyncMock(return_value=mock_client)
        mock_client.__aexit__ = AsyncMock(return_value=False)

        with (
            patch.dict("os.environ", {"WAFER_API_URL": "http://test", "WAFER_AUTH_TOKEN": "tok"}),
            patch("wafer_core.tools.ask_docs_tool.httpx.AsyncClient", return_value=mock_client),
        ):
            tc = ToolCall(id="test-5", name="ask_docs", args={"query": "test", "corpus": "cuda"})
            result = await exec_ask_docs(tc)

        assert not result.is_error
        assert result.content == "Hello world"

    @pytest.mark.anyio
    async def test_auth_error_401(self) -> None:
        """401 response should return auth error."""
        mock_response = AsyncMock()
        mock_response.status_code = 401
        mock_response.__aenter__ = AsyncMock(return_value=mock_response)
        mock_response.__aexit__ = AsyncMock(return_value=False)

        mock_client = AsyncMock()
        mock_client.stream = MagicMock(return_value=mock_response)
        mock_client.__aenter__ = AsyncMock(return_value=mock_client)
        mock_client.__aexit__ = AsyncMock(return_value=False)

        with (
            patch.dict("os.environ", {"WAFER_API_URL": "http://test", "WAFER_AUTH_TOKEN": "bad"}),
            patch("wafer_core.tools.ask_docs_tool.httpx.AsyncClient", return_value=mock_client),
        ):
            tc = ToolCall(id="test-6", name="ask_docs", args={"query": "test", "corpus": "cuda"})
            result = await exec_ask_docs(tc)

        assert result.is_error
        assert "Authentication" in result.error

    @pytest.mark.anyio
    async def test_api_error_500(self) -> None:
        """500 response should return API error."""
        mock_response = AsyncMock()
        mock_response.status_code = 500
        mock_response.aread = AsyncMock(return_value=b"Internal Server Error")
        mock_response.__aenter__ = AsyncMock(return_value=mock_response)
        mock_response.__aexit__ = AsyncMock(return_value=False)

        mock_client = AsyncMock()
        mock_client.stream = MagicMock(return_value=mock_response)
        mock_client.__aenter__ = AsyncMock(return_value=mock_client)
        mock_client.__aexit__ = AsyncMock(return_value=False)

        with (
            patch.dict("os.environ", {"WAFER_API_URL": "http://test", "WAFER_AUTH_TOKEN": "tok"}),
            patch("wafer_core.tools.ask_docs_tool.httpx.AsyncClient", return_value=mock_client),
        ):
            tc = ToolCall(id="test-7", name="ask_docs", args={"query": "test", "corpus": "cuda"})
            result = await exec_ask_docs(tc)

        assert result.is_error
        assert "500" in result.error

    @pytest.mark.anyio
    async def test_timeout(self) -> None:
        """Timeout should return timeout error."""
        import httpx

        mock_client = AsyncMock()
        mock_client.stream = MagicMock(side_effect=httpx.TimeoutException("timed out"))
        mock_client.__aenter__ = AsyncMock(return_value=mock_client)
        mock_client.__aexit__ = AsyncMock(return_value=False)

        with (
            patch.dict("os.environ", {"WAFER_API_URL": "http://test", "WAFER_AUTH_TOKEN": "tok"}),
            patch("wafer_core.tools.ask_docs_tool.httpx.AsyncClient", return_value=mock_client),
        ):
            tc = ToolCall(id="test-8", name="ask_docs", args={"query": "test", "corpus": "cuda"})
            result = await exec_ask_docs(tc)

        assert result.is_error
        assert "timed out" in result.error


async def _async_iter(items: list[str]):
    """Helper to create an async iterator from a list."""
    for item in items:
        yield item
