"""Tests for query rewrite module."""

from unittest.mock import AsyncMock, patch

from nlweb_platform_handlers.query_rewrite import (
    QueryRewriteResponse,
    rewrite_query,
)


class TestQueryRewriteResponse:
    """Tests for QueryRewriteResponse model."""

    def test_valid_queries(self):
        resp = QueryRewriteResponse(rewritten_queries=["tea pot", "green tea"])
        assert resp.rewritten_queries == ["tea pot", "green tea"]

    def test_filters_empty_strings(self):
        resp = QueryRewriteResponse(
            rewritten_queries=["tea pot", "", "  ", "green tea"]
        )
        assert resp.rewritten_queries == ["tea pot", "green tea"]

    def test_limits_to_five(self):
        queries = [f"query {i}" for i in range(10)]
        resp = QueryRewriteResponse(rewritten_queries=queries)
        assert len(resp.rewritten_queries) == 5

    def test_strips_whitespace(self):
        resp = QueryRewriteResponse(rewritten_queries=["  tea pot  ", " green tea "])
        assert resp.rewritten_queries == ["tea pot", "green tea"]


class TestRewriteQuery:
    """Tests for rewrite_query function."""

    @patch("nlweb_platform_handlers.query_rewrite.ask_llm", new_callable=AsyncMock)
    async def test_returns_rewritten_queries(self, mock_ask_llm):
        mock_ask_llm.return_value = QueryRewriteResponse(
            rewritten_queries=["vegetable plates", "serving plates"]
        )
        result = await rewrite_query("plates for serving vegetables")
        assert result == ["vegetable plates", "serving plates"]
        mock_ask_llm.assert_called_once()

    @patch("nlweb_platform_handlers.query_rewrite.ask_llm", new_callable=AsyncMock)
    async def test_falls_back_to_original_on_error(self, mock_ask_llm):
        mock_ask_llm.side_effect = Exception("LLM unavailable")
        result = await rewrite_query("tea pot for green tea")
        assert result == ["tea pot for green tea"]

    @patch("nlweb_platform_handlers.query_rewrite.ask_llm", new_callable=AsyncMock)
    async def test_falls_back_on_validation_error(self, mock_ask_llm):
        # LLM returns invalid response that fails Pydantic validation
        from pydantic import ValidationError

        mock_ask_llm.side_effect = ValidationError.from_exception_data(
            title="QueryRewriteResponse",
            line_errors=[],
        )
        result = await rewrite_query("original query")
        assert result == ["original query"]
