"""
Tests for the RAG query pipeline in litellm/rag/main.py.

The RAG pipeline forwards its kwargs (including the parent litellm_logging_obj)
into @client-decorated sub-calls (vector store search, completion). Each logging
object allows exactly one async_success event, so if sub-calls are not marked as
internal, the vector store search consumes the slot first and the LLM
completion's usage/cost is never logged (spend tracking and budget enforcement
are bypassed). These tests pin the invariant that the single billing event for
aquery carries the completion response with real usage and cost.
"""

import asyncio
from unittest.mock import patch

import pytest

import litellm
from litellm._internal_context import is_internal_call
from litellm.integrations.custom_logger import CustomLogger
from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
from litellm.types.utils import CallTypes, ModelResponse


async def _drain_logging_worker() -> None:
    """Run every queued logging task to completion on the current event loop.

    The success event is delivered through the fire-and-forget GLOBAL_LOGGING_WORKER
    singleton, whose queue survives across tests. start() rebinds any tasks left over
    from a previous test's event loop onto the current one, and flush() waits until
    the queue is fully processed, so tests neither miss their own event nor observe
    a neighbour's
    """
    await asyncio.sleep(0)
    GLOBAL_LOGGING_WORKER.start()
    await asyncio.wait_for(GLOBAL_LOGGING_WORKER.flush(), timeout=10.0)


class RecordingLogger(CustomLogger):
    def __init__(self):
        super().__init__()
        self.success_events = []

    async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
        self.success_events.append({"kwargs": kwargs, "response_obj": response_obj})


@pytest.mark.asyncio
@pytest.mark.parametrize("use_router", [False, True])
async def test_aquery_single_billing_event_carries_completion_usage_and_cost(use_router):
    """
    litellm.aquery must produce exactly one success event, and that event must
    carry the LLM completion (a ModelResponse with non-zero usage and cost),
    not the vector store search response. The proxy always passes a router, so
    both the router and non-router completion branches are pinned.
    """
    await _drain_logging_worker()
    recording_logger = RecordingLogger()
    original_callbacks = litellm.callbacks
    litellm.callbacks = [recording_logger]

    router_kwargs = {}
    if use_router:
        router_kwargs["router"] = litellm.Router(
            model_list=[
                {
                    "model_name": "gpt-4o-mini",
                    "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "test-key"},
                }
            ]
        )

    try:
        response = await litellm.aquery(
            model="gpt-4o-mini",
            messages=[{"role": "user", "content": "What is the secret project codename?"}],
            retrieval_config={"vector_store_id": "vs_test_123", "custom_llm_provider": "openai"},
            mock_response="The secret project codename is AZURE-FALCON-42.",
            **router_kwargs,
        )

        assert isinstance(response, ModelResponse)
        assert is_internal_call.get() is False

        await _drain_logging_worker()
    finally:
        litellm.callbacks = original_callbacks

    assert len(recording_logger.success_events) == 1
    event = recording_logger.success_events[0]

    response_obj = event["response_obj"]
    assert isinstance(response_obj, ModelResponse)
    assert response_obj.usage.total_tokens > 0

    standard_logging_object = event["kwargs"]["standard_logging_object"]
    assert standard_logging_object["call_type"] == "aquery"
    assert standard_logging_object["total_tokens"] > 0
    assert standard_logging_object["prompt_tokens"] > 0
    assert standard_logging_object["completion_tokens"] > 0
    assert standard_logging_object["response_cost"] > 0


@pytest.mark.asyncio
async def test_aquery_response_hidden_params_carry_completion_cost():
    """
    The aquery response must expose the completion's response_cost via hidden
    params, so the proxy can return the x-litellm-response-cost header.
    """
    response = await litellm.aquery(
        model="gpt-4o-mini",
        messages=[{"role": "user", "content": "hello"}],
        retrieval_config={"vector_store_id": "vs_test_123", "custom_llm_provider": "openai"},
        mock_response="hi there",
    )

    await _drain_logging_worker()

    assert isinstance(response, ModelResponse)
    response_cost = response._hidden_params.get("response_cost")
    assert response_cost is not None
    assert response_cost > 0


@pytest.mark.asyncio
async def test_aquery_billed_cost_includes_priced_vector_store_search():
    """
    When the vector store provider prices search calls (e.g. per-query cost),
    that cost must be folded into the aquery billing instead of being dropped
    with the suppressed sub-call event.
    """
    await _drain_logging_worker()
    recording_logger = RecordingLogger()
    original_callbacks = litellm.callbacks
    litellm.callbacks = [recording_logger]

    try:
        with patch("litellm.rag.main.vector_store_search_cost", return_value=(0.002, 0.0)):
            response = await litellm.aquery(
                model="gpt-4o-mini",
                messages=[{"role": "user", "content": "hello"}],
                retrieval_config={"vector_store_id": "vs_test_123", "custom_llm_provider": "openai"},
                mock_response="hi there",
            )

        await _drain_logging_worker()
    finally:
        litellm.callbacks = original_callbacks

    assert isinstance(response, ModelResponse)
    total_cost = response._hidden_params.get("response_cost")
    assert total_cost is not None
    assert total_cost > 0.002

    assert len(recording_logger.success_events) == 1
    standard_logging_object = recording_logger.success_events[0]["kwargs"]["standard_logging_object"]
    assert standard_logging_object["response_cost"] == total_cost


@pytest.mark.asyncio
async def test_aquery_with_rerank_bills_once_and_folds_rerank_cost():
    """
    When rerank is enabled, its sub-call must run under the internal-call
    context (no standalone billing event) and its cost must be folded into
    the single aquery billing event.
    """
    from litellm.types.rerank import RerankResponse

    await _drain_logging_worker()
    recording_logger = RecordingLogger()
    original_callbacks = litellm.callbacks
    litellm.callbacks = [recording_logger]
    rerank_seen = {}

    async def fake_arerank(**kwargs):
        rerank_seen["internal"] = is_internal_call.get()
        rerank_result = RerankResponse(id="rr_1", results=[{"index": 0, "relevance_score": 0.9}], meta={})
        rerank_result._hidden_params["response_cost"] = 0.001
        return rerank_result

    try:
        with patch("litellm.arerank", side_effect=fake_arerank):
            response = await litellm.aquery(
                model="gpt-4o-mini",
                messages=[{"role": "user", "content": "hello"}],
                retrieval_config={"vector_store_id": "vs_test_123", "custom_llm_provider": "openai"},
                rerank={"enabled": True, "model": "cohere/rerank-english-v3.0", "top_n": 1},
                mock_response="hi there",
            )

        await _drain_logging_worker()
    finally:
        litellm.callbacks = original_callbacks

    assert rerank_seen["internal"] is True
    assert is_internal_call.get() is False

    assert isinstance(response, ModelResponse)
    total_cost = response._hidden_params.get("response_cost")
    assert total_cost is not None
    assert total_cost > 0.001

    assert len(recording_logger.success_events) == 1
    standard_logging_object = recording_logger.success_events[0]["kwargs"]["standard_logging_object"]
    assert standard_logging_object["call_type"] == "aquery"
    assert standard_logging_object["response_cost"] == total_cost


@pytest.mark.asyncio
async def test_aquery_streaming_bills_sub_call_costs_into_final_event():
    """
    On the streaming path the response cost is computed from the assembled
    chunks after the pipeline returns, so there is no response object to fold
    sub-call costs into. The pipeline must instead carry the accumulated
    search and rerank cost through the logging object so the single streamed
    billing event includes it; otherwise a caller passing stream=true incurs
    priced vector search and rerank costs that never reach spend tracking.
    """
    from litellm.types.rerank import RerankResponse

    await _drain_logging_worker()
    recording_logger = RecordingLogger()
    original_callbacks = litellm.callbacks
    litellm.callbacks = [recording_logger]
    rerank_seen = {}

    async def fake_arerank(**kwargs):
        rerank_seen["internal"] = is_internal_call.get()
        rerank_result = RerankResponse(id="rr_1", results=[{"index": 0, "relevance_score": 0.9}], meta={})
        rerank_result._hidden_params["response_cost"] = 0.001
        return rerank_result

    try:
        with (
            patch("litellm.rag.main.vector_store_search_cost", return_value=(0.002, 0.0)),
            patch("litellm.arerank", side_effect=fake_arerank),
        ):
            response = await litellm.aquery(
                model="gpt-4o-mini",
                messages=[{"role": "user", "content": "hello"}],
                retrieval_config={"vector_store_id": "vs_test_123", "custom_llm_provider": "openai"},
                rerank={"enabled": True, "model": "cohere/rerank-english-v3.0", "top_n": 1},
                mock_response="hi there",
                stream=True,
            )
            async for _ in response:
                pass

        await _drain_logging_worker()
    finally:
        litellm.callbacks = original_callbacks

    assert rerank_seen["internal"] is True
    assert is_internal_call.get() is False

    assert len(recording_logger.success_events) == 1
    standard_logging_object = recording_logger.success_events[0]["kwargs"]["standard_logging_object"]
    assert standard_logging_object["call_type"] == "aquery"
    assert standard_logging_object["response_cost"] >= 0.003


@pytest.mark.asyncio
async def test_aquery_forwards_provider_retrieval_config_and_router_to_search():
    """
    Regression: provider-specific retrieval_config keys (aws_region_name,
    embedding_model, vector_bucket_name, ...) and the router must be forwarded
    to the vector store search call. Pre-fix they were silently dropped, so
    /v1/rag/query failed with provider config errors (e.g. S3 Vectors
    "aws_region_name is required") even when the caller supplied them.
    """
    from unittest.mock import AsyncMock

    from litellm.types.vector_stores import VectorStoreSearchResponse

    router = litellm.Router(
        model_list=[
            {
                "model_name": "gpt-4o-mini",
                "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "test-key"},
            }
        ]
    )

    fake_search = AsyncMock(
        return_value=VectorStoreSearchResponse(
            object="vector_store.search_results.page", search_query="q", data=[]
        )
    )
    with patch("litellm.vector_stores.asearch", new=fake_search):  # test-quality-ok: asearch is the boundary the forwarding contract under test targets
        response = await litellm.aquery(
            model="gpt-4o-mini",
            messages=[{"role": "user", "content": "hello"}],
            retrieval_config={
                "vector_store_id": "bkt:idx",
                "custom_llm_provider": "s3_vectors",
                "top_k": 5,
                "aws_region_name": "eu-west-1",
                "embedding_model": "my-embed",
                "vector_bucket_name": "bkt",
            },
            router=router,
            mock_response="hi",
        )

    assert isinstance(response, ModelResponse)
    fake_search.assert_awaited_once()
    search_kwargs = fake_search.await_args.kwargs
    assert search_kwargs["vector_store_id"] == "bkt:idx"
    assert search_kwargs["custom_llm_provider"] == "s3_vectors"
    assert search_kwargs["max_num_results"] == 5
    assert search_kwargs["router"] is router
    # provider-specific extras forwarded
    assert search_kwargs["aws_region_name"] == "eu-west-1"
    assert search_kwargs["embedding_model"] == "my-embed"
    assert search_kwargs["vector_bucket_name"] == "bkt"
    # consumed keys are not duplicated into the spread
    assert "top_k" not in search_kwargs


@pytest.mark.asyncio
async def test_aquery_minimal_retrieval_config_forwards_no_extras():
    """
    A minimal retrieval_config must not leak consumed keys (or invent extras)
    into the vector store search call.
    """
    from unittest.mock import AsyncMock

    from litellm.types.vector_stores import VectorStoreSearchResponse

    fake_search = AsyncMock(
        return_value=VectorStoreSearchResponse(
            object="vector_store.search_results.page", search_query="q", data=[]
        )
    )
    with patch("litellm.vector_stores.asearch", new=fake_search):  # test-quality-ok: asearch is the boundary the forwarding contract under test targets
        await litellm.aquery(
            model="gpt-4o-mini",
            messages=[{"role": "user", "content": "hello"}],
            retrieval_config={"vector_store_id": "vs_test_123", "custom_llm_provider": "openai"},
            mock_response="hi",
        )

    fake_search.assert_awaited_once()
    search_kwargs = fake_search.await_args.kwargs
    assert search_kwargs["vector_store_id"] == "vs_test_123"
    assert search_kwargs["custom_llm_provider"] == "openai"
    assert search_kwargs["router"] is None
    leaked = {"top_k", "filters", "retrieval_filter", "aws_region_name", "embedding_model", "vector_bucket_name"}
    assert not (leaked & set(search_kwargs.keys()))


@pytest.mark.asyncio
async def test_aquery_does_not_forward_connection_override_keys_to_search():
    """
    Only allowlisted retrieval_config keys may reach the vector store search
    call. Caller-controlled connection overrides (api_base, api_key, arbitrary
    extras) must be dropped, otherwise a caller could redirect store
    credentials to an attacker-chosen host.
    """
    from unittest.mock import AsyncMock

    from litellm.types.vector_stores import VectorStoreSearchResponse

    fake_search = AsyncMock(
        return_value=VectorStoreSearchResponse(
            object="vector_store.search_results.page", search_query="q", data=[]
        )
    )
    with patch("litellm.vector_stores.asearch", new=fake_search):  # test-quality-ok: asearch is the boundary the forwarding contract under test targets
        await litellm.aquery(
            model="gpt-4o-mini",
            messages=[{"role": "user", "content": "hello"}],
            retrieval_config={
                "vector_store_id": "bkt:idx",
                "custom_llm_provider": "s3_vectors",
                "aws_region_name": "eu-west-1",
                "api_base": "https://attacker.example.com",
                "api_key": "attacker-key",
                "arbitrary_extra": "nope",
            },
            mock_response="hi",
        )

    fake_search.assert_awaited_once()
    search_kwargs = fake_search.await_args.kwargs
    assert search_kwargs["aws_region_name"] == "eu-west-1"
    blocked = {"api_base", "api_key", "arbitrary_extra"}
    assert not (blocked & set(search_kwargs.keys()))


@pytest.mark.asyncio
async def test_aquery_forwards_vector_store_params_to_search_but_not_completion():
    """
    Regression for LIT-6773: the server-trusted vector_store_params (a managed
    store's litellm_params) must reach the search call wholesale, including the
    connection keys the caller allowlist blocks, while the caller's own
    retrieval_config overrides stay blocked, the caller's top-level api_key and
    api_base stay on the completion only, and the completion never inherits the
    store's connection params.
    """
    from unittest.mock import AsyncMock

    from litellm.types.vector_stores import VectorStoreSearchResponse

    fake_search = AsyncMock(
        return_value=VectorStoreSearchResponse(
            object="vector_store.search_results.page", search_query="q", data=[]
        )
    )
    fake_completion = AsyncMock(
        return_value=ModelResponse(
            id="chatcmpl-test",
            choices=[{"index": 0, "message": {"role": "assistant", "content": "hi"}, "finish_reason": "stop"}],
            model="gpt-4o-mini",
        )
    )
    with (
        patch("litellm.vector_stores.asearch", new=fake_search),  # test-quality-ok: the search boundary under test
        patch("litellm.acompletion", new=fake_completion),  # test-quality-ok: the completion boundary under test
    ):
        await litellm.aquery(
            model="gpt-4o-mini",
            messages=[{"role": "user", "content": "hello"}],
            api_key="sk-llm-key",
            api_base="https://llm.example.com",
            retrieval_config={
                "vector_store_id": "customer_kb",
                "custom_llm_provider": "milvus",
                "api_base": "https://attacker.example.com",
                "api_key": "attacker-key",
            },
            vector_store_params={
                "vector_store_id": "customer_kb",
                "custom_llm_provider": "milvus",
                "api_base": "http://127.0.0.1:19530",
                "api_key": "root:Milvus",
                "milvus_text_field": "book_intro_text",
                "outputFields": ["book_intro_text"],
            },
        )

    fake_search.assert_awaited_once()
    search_kwargs = fake_search.await_args.kwargs
    assert search_kwargs["vector_store_id"] == "customer_kb"
    assert search_kwargs["custom_llm_provider"] == "milvus"
    assert search_kwargs["api_base"] == "http://127.0.0.1:19530"
    assert search_kwargs["api_key"] == "root:Milvus"
    assert search_kwargs["milvus_text_field"] == "book_intro_text"
    assert search_kwargs["outputFields"] == ["book_intro_text"]
    fake_completion.assert_awaited_once()
    completion_kwargs = fake_completion.await_args.kwargs
    assert completion_kwargs["api_key"] == "sk-llm-key"
    assert completion_kwargs["api_base"] == "https://llm.example.com"
    assert not ({"milvus_text_field", "outputFields"} & set(completion_kwargs))


def test_rag_call_types_are_registered():
    """
    query/aquery/ingest/aingest are @client-decorated entry points, so their
    function names must resolve to CallTypes members (deployment hooks and
    call-type driven logic silently no-op for unregistered call types).
    """
    assert CallTypes("query") is CallTypes.query
    assert CallTypes("aquery") is CallTypes.aquery
    assert CallTypes("ingest") is CallTypes.ingest
    assert CallTypes("aingest") is CallTypes.aingest