# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Unit tests for EmbedIOProcessor."""

import pytest
import torch
from pydantic import TypeAdapter, ValidationError

from vllm import PoolingParams
from vllm.entrypoints.pooling.embed.io_processor import EmbedIOProcessor
from vllm.entrypoints.pooling.embed.protocol import (
    CohereEmbedContent,
    CohereEmbedInput,
    CohereEmbedRequest,
    EmbeddingBatchChatInputRequest,
    EmbeddingBatchChatRequest,
    EmbeddingChatInputRequest,
    EmbeddingChatRequest,
    EmbeddingCompletionRequest,
    EmbeddingRequest,
)
from vllm.entrypoints.pooling.typing import PoolingServeContext
from vllm.outputs import PoolingOutput, PoolingRequestOutput


class TestEmbeddingRequestParsing:
    """Unit tests for OpenAI embedding request parsing."""

    def test_input_messages_parses_as_chat_request(self):
        request = TypeAdapter(EmbeddingRequest).validate_python(
            {
                "model": "test",
                "input": [{"role": "user", "content": "hello"}],
                "chat_template_kwargs": {"instruction": "Represent the query: "},
            }
        )

        assert isinstance(request, EmbeddingChatInputRequest)
        assert request.input == [{"role": "user", "content": "hello"}]
        assert request.messages == [{"role": "user", "content": "hello"}]
        assert request.chat_template_kwargs == {"instruction": "Represent the query: "}

    def test_batched_input_messages_parses_as_batch_chat_input_request(self):
        request = TypeAdapter(EmbeddingRequest).validate_python(
            {
                "model": "test",
                "input": [
                    [{"role": "user", "content": "hello"}],
                    [{"role": "user", "content": "goodbye"}],
                ],
                "chat_template_kwargs": {"instruction": "Represent the query: "},
            }
        )

        assert isinstance(request, EmbeddingBatchChatInputRequest)
        assert request.input == [
            [{"role": "user", "content": "hello"}],
            [{"role": "user", "content": "goodbye"}],
        ]
        assert request.messages == [
            [{"role": "user", "content": "hello"}],
            [{"role": "user", "content": "goodbye"}],
        ]
        assert request.chat_template_kwargs == {"instruction": "Represent the query: "}

    def test_token_ids_still_parse_as_completion_request(self):
        request = TypeAdapter(EmbeddingRequest).validate_python(
            {
                "model": "test",
                "input": [[1, 2, 3], [4, 5]],
            }
        )

        assert isinstance(request, EmbeddingCompletionRequest)
        assert request.input == [[1, 2, 3], [4, 5]]

    def test_messages_still_parses_as_chat_request(self):
        request = TypeAdapter(EmbeddingRequest).validate_python(
            {
                "model": "test",
                "messages": [{"role": "user", "content": "hello"}],
                "chat_template_kwargs": {"instruction": "Represent the query: "},
            }
        )

        assert isinstance(request, EmbeddingChatRequest)
        assert request.messages == [{"role": "user", "content": "hello"}]
        assert request.chat_template_kwargs == {"instruction": "Represent the query: "}

    def test_batched_messages_parses_as_batch_chat_request(self):
        request = TypeAdapter(EmbeddingRequest).validate_python(
            {
                "model": "test",
                "messages": [
                    [{"role": "user", "content": "hello"}],
                    [{"role": "user", "content": "goodbye"}],
                ],
                "chat_template_kwargs": {"instruction": "Represent the query: "},
            }
        )

        assert isinstance(request, EmbeddingBatchChatRequest)
        assert request.messages == [
            [{"role": "user", "content": "hello"}],
            [{"role": "user", "content": "goodbye"}],
        ]
        assert request.chat_template_kwargs == {"instruction": "Represent the query: "}


class TestCohereEmbedRequestParsing:
    """Unit tests for Cohere embed request parsing."""

    @pytest.mark.parametrize(
        "request_body",
        [
            {"model": "test"},
            {"model": "test", "texts": ["hello"], "images": ["image-uri"]},
            {
                "model": "test",
                "texts": ["hello"],
                "inputs": [
                    {"content": [{"type": "text", "text": "hello"}]},
                ],
            },
            {
                "model": "test",
                "images": ["image-uri"],
                "inputs": [
                    {"content": [{"type": "text", "text": "hello"}]},
                ],
            },
            {"model": "test", "texts": []},
            {"model": "test", "images": []},
            {"model": "test", "inputs": []},
        ],
    )
    def test_rejects_invalid_input_field_combinations(self, request_body):
        with pytest.raises(
            ValidationError,
            match="Exactly one of texts, images, or inputs must be provided",
        ):
            CohereEmbedRequest(**request_body)

    @pytest.mark.parametrize(
        "request_body",
        [
            {"model": "test", "texts": ["hello"]},
            {"model": "test", "images": ["image-uri"]},
            {
                "model": "test",
                "inputs": [
                    {"content": [{"type": "text", "text": "hello"}]},
                ],
            },
            {
                "model": "test",
                "inputs": [
                    {
                        "content": [
                            {"type": "image_url", "image_url": {"url": "image-uri"}}
                        ]
                    },
                ],
            },
        ],
    )
    def test_accepts_exactly_one_non_empty_input_field(self, request_body):
        request = CohereEmbedRequest(**request_body)

        assert request.model == "test"

    @pytest.mark.parametrize(
        ("content", "error"),
        [
            (
                {"type": "text"},
                "CohereEmbedContent with type='text' requires text",
            ),
            (
                {"type": "image_url"},
                "CohereEmbedContent with type='image_url' requires image_url.url",
            ),
            (
                {"type": "image_url", "image_url": {}},
                "CohereEmbedContent with type='image_url' requires image_url.url",
            ),
            (
                {"type": "image_url", "image_url": {"url": ""}},
                "CohereEmbedContent with type='image_url' requires image_url.url",
            ),
        ],
    )
    def test_rejects_invalid_mixed_content_payloads(self, content, error):
        with pytest.raises(ValidationError, match=error):
            CohereEmbedRequest(
                model="test",
                inputs=[
                    {
                        "content": [content],
                    },
                ],
            )


class TestResolveTruncation:
    """Unit tests for EmbedIOProcessor._resolve_cohere_truncation."""

    @staticmethod
    def _make_request(**kwargs) -> CohereEmbedRequest:
        defaults = {
            "model": "test",
            "input_type": "search_document",
            "texts": ["hello"],
        }
        return CohereEmbedRequest(**(defaults | kwargs))

    def test_truncate_end_default(self):
        req = self._make_request()
        tokens, side = EmbedIOProcessor._resolve_cohere_truncation(req)
        assert tokens == -1
        assert side is None

    def test_truncate_end_explicit(self):
        req = self._make_request(truncate="END")
        tokens, side = EmbedIOProcessor._resolve_cohere_truncation(req)
        assert tokens == -1
        assert side is None

    def test_truncate_end_with_max_tokens(self):
        req = self._make_request(truncate="END", max_tokens=128)
        tokens, side = EmbedIOProcessor._resolve_cohere_truncation(req)
        assert tokens == 128
        assert side is None

    def test_truncate_none(self):
        req = self._make_request(truncate="NONE")
        tokens, side = EmbedIOProcessor._resolve_cohere_truncation(req)
        assert tokens is None
        assert side is None

    def test_truncate_none_with_max_tokens(self):
        """truncate=NONE should NOT set truncate_prompt_tokens; the
        max_tokens limit is enforced separately via _check_max_tokens."""
        req = self._make_request(truncate="NONE", max_tokens=10)
        tokens, side = EmbedIOProcessor._resolve_cohere_truncation(req)
        assert tokens is None
        assert side is None

    def test_truncate_start(self):
        req = self._make_request(truncate="START")
        tokens, side = EmbedIOProcessor._resolve_cohere_truncation(req)
        assert tokens == -1
        assert side == "left"

    def test_truncate_start_with_max_tokens(self):
        req = self._make_request(truncate="START", max_tokens=64)
        tokens, side = EmbedIOProcessor._resolve_cohere_truncation(req)
        assert tokens == 64
        assert side == "left"


class TestApplyStPrompt:
    """Unit tests for EmbedIOProcessor._apply_task_instruction."""

    @staticmethod
    def _make_handler(task_instructions: dict[str, str] | None):
        handler = object.__new__(EmbedIOProcessor)
        handler.task_instructions = task_instructions
        return handler

    def test_no_prompts_configured(self):
        handler = self._make_handler(None)
        texts = ["hello", "world"]
        assert handler._apply_task_instruction(texts, "query") is texts

    def test_matching_input_type(self):
        handler = self._make_handler({"query": "search_query: "})
        result = handler._apply_task_instruction(["hello"], "query")
        assert result == ["search_query: hello"]

    def test_non_matching_input_type(self):
        handler = self._make_handler({"query": "search_query: "})
        texts = ["hello"]
        assert handler._apply_task_instruction(texts, "document") is texts

    def test_multiple_texts(self):
        handler = self._make_handler(
            {"query": "Represent this sentence for searching: "}
        )
        result = handler._apply_task_instruction(["a", "b", "c"], "query")
        assert result == [
            "Represent this sentence for searching: a",
            "Represent this sentence for searching: b",
            "Represent this sentence for searching: c",
        ]

    def test_empty_prefix_returns_unchanged(self):
        handler = self._make_handler({"passage": ""})
        texts = ["hello"]
        assert handler._apply_task_instruction(texts, "passage") is texts


class TestLoadTaskInstructions:
    """Unit tests for EmbedIOProcessor._load_task_instructions."""

    def test_no_attribute(self):
        class FakeConfig:
            pass

        assert EmbedIOProcessor._load_task_instructions(FakeConfig()) is None

    def test_with_task_instructions(self):
        class FakeConfig:
            task_instructions = {
                "retrieval.query": "Represent the query: ",
                "retrieval.passage": "",
            }

        result = EmbedIOProcessor._load_task_instructions(FakeConfig())
        assert result == {
            "retrieval.query": "Represent the query: ",
            "retrieval.passage": "",
        }

    def test_empty_dict(self):
        class FakeConfig:
            task_instructions = {}

        assert EmbedIOProcessor._load_task_instructions(FakeConfig()) is None

    def test_non_dict(self):
        class FakeConfig:
            task_instructions = "not a dict"

        assert EmbedIOProcessor._load_task_instructions(FakeConfig()) is None


class TestCheckMaxTokens:
    """Unit tests for EmbedIOProcessor._check_cohere_max_tokens."""

    @staticmethod
    def _fake_output(n_tokens: int):
        class _Out:
            def __init__(self, n: int):
                self.prompt_token_ids = list(range(n))

        return _Out(n_tokens)

    def test_none_check_is_noop(self):
        outs = [self._fake_output(100)]
        EmbedIOProcessor._check_cohere_max_tokens(outs, None)

    def test_within_limit(self):
        outs = [self._fake_output(5), self._fake_output(3)]
        EmbedIOProcessor._check_cohere_max_tokens(outs, 5)

    def test_exceeds_limit(self):
        outs = [self._fake_output(3), self._fake_output(10)]
        with pytest.raises(ValueError, match="exceeds max_tokens=5"):
            EmbedIOProcessor._check_cohere_max_tokens(outs, 5)

    def test_exact_limit(self):
        outs = [self._fake_output(5)]
        EmbedIOProcessor._check_cohere_max_tokens(outs, 5)


class TestValidateInputType:
    """Unit tests for EmbedIOProcessor._validate_input_type."""

    @staticmethod
    def _make_handler(task_instructions: dict[str, str] | None):
        handler = object.__new__(EmbedIOProcessor)
        handler.task_instructions = task_instructions
        return handler

    def test_none_input_type_always_accepted(self):
        handler = self._make_handler(None)
        handler._validate_input_type(None)
        handler_with = self._make_handler({"query": "q: "})
        handler_with._validate_input_type(None)

    def test_no_prompts_rejects(self):
        handler = self._make_handler(None)
        with pytest.raises(ValueError, match="does not define any input_type"):
            handler._validate_input_type("anything")

    def test_known_type_accepted(self):
        handler = self._make_handler({"query": "q: ", "document": "d: "})
        handler._validate_input_type("query")
        handler._validate_input_type("document")

    def test_unknown_type_rejected(self):
        handler = self._make_handler({"query": "q: ", "document": "d: "})
        with pytest.raises(ValueError, match="Unsupported input_type 'other'"):
            handler._validate_input_type("other")

    def test_error_lists_supported(self):
        handler = self._make_handler({"a": "", "b": ""})
        with pytest.raises(ValueError, match="Supported values: a, b"):
            handler._validate_input_type("z")


class TestChunkedEmbeddingProcessing:
    """Unit tests for chunked embedding aggregation."""

    class _FakeModelConfig:
        max_model_len = 3

    @classmethod
    def _make_handler(cls):
        handler = object.__new__(EmbedIOProcessor)
        handler.model_config = cls._FakeModelConfig()
        return handler

    @staticmethod
    def _make_context() -> PoolingServeContext[EmbeddingCompletionRequest]:
        request = TypeAdapter(EmbeddingRequest).validate_python(
            {
                "model": "test",
                "input": [[0, 1, 2, 3, 4], [10, 11]],
            }
        )
        assert isinstance(request, EmbeddingCompletionRequest)
        return PoolingServeContext(
            request=request,
            pooling_params=PoolingParams(),
            model_name="test",
            request_id="embd-client-prompt-999-chunk-888",
            engine_inputs=[
                {"prompt_token_ids": [0, 1, 2, 3, 4]},
                {"prompt_token_ids": [10, 11]},
            ],
        )

    @staticmethod
    def _make_output(
        request_id: str,
        prompt_token_ids: list[int],
        embedding: list[float],
    ) -> PoolingRequestOutput:
        return PoolingRequestOutput(
            request_id=request_id,
            outputs=PoolingOutput(data=torch.tensor(embedding)),
            prompt_token_ids=prompt_token_ids,
            num_cached_tokens=0,
            finished=True,
        )

    def test_aggregation_uses_metadata_not_request_id_parsing(self):
        handler = self._make_handler()
        ctx = self._make_context()

        handler._pre_process_chunked(ctx)

        assert ctx.prompt_request_ids == [
            "embd-client-prompt-999-chunk-888-prompt-0-chunk-0",
            "embd-client-prompt-999-chunk-888-prompt-0-chunk-1",
            "embd-client-prompt-999-chunk-888-prompt-1-chunk-0",
        ]
        assert ctx.chunked_embedding_metadata is not None
        assert [
            (item.prompt_index, item.chunk_index)
            for item in ctx.chunked_embedding_metadata
        ] == [(0, 0), (0, 1), (1, 0)]

        ctx.final_res_batch = [
            self._make_output(ctx.prompt_request_ids[0], [0, 1, 2], [1.0, 1.0]),
            self._make_output(ctx.prompt_request_ids[1], [3, 4], [4.0, 7.0]),
            self._make_output(ctx.prompt_request_ids[2], [10, 11], [9.0, 9.0]),
        ]

        handler._post_process_chunked(ctx)

        assert len(ctx.final_res_batch) == 2
        assert ctx.final_res_batch[0].request_id == (
            "embd-client-prompt-999-chunk-888-prompt-0"
        )
        assert ctx.final_res_batch[0].prompt_token_ids == [0, 1, 2, 3, 4]
        assert torch.allclose(
            ctx.final_res_batch[0].outputs.data,
            torch.tensor([2.2, 3.4]),
        )
        assert ctx.final_res_batch[1].request_id == (
            "embd-client-prompt-999-chunk-888-prompt-1"
        )
        assert ctx.final_res_batch[1].prompt_token_ids == [10, 11]
        assert torch.allclose(
            ctx.final_res_batch[1].outputs.data,
            torch.tensor([9.0, 9.0]),
        )


class TestPreProcessCohereOnline:
    """Unit tests for EmbedIOProcessor._pre_process_cohere_online."""

    @staticmethod
    def _make_context(**request_kwargs) -> PoolingServeContext[CohereEmbedRequest]:
        return PoolingServeContext(
            request=CohereEmbedRequest(model="test", **request_kwargs),
            pooling_params=PoolingParams(),
            model_name="test",
            request_id="embd-test",
        )

    @staticmethod
    def _make_handler():
        handler = object.__new__(EmbedIOProcessor)
        handler._validate_input_type = lambda _input_type: None
        return handler

    def test_text_only_without_task_prefix_uses_completion_path(self):
        handler = self._make_handler()
        ctx = self._make_context(texts=["hello"])
        calls: list[tuple[str, object]] = []

        def preprocess_cmpl_online(request, prompt_input, prompt_embeds):
            calls.append(("completion", prompt_input))
            return ["completion"]

        handler._get_task_instruction_prefix = lambda _input_type: None
        handler._has_chat_template = lambda: False
        handler._preprocess_cmpl_online = preprocess_cmpl_online
        handler._batch_render_chat = lambda *_args, **_kwargs: pytest.fail(
            "text-only request should not require chat rendering"
        )

        handler._pre_process_cohere_online(ctx)

        assert ctx.engine_inputs == ["completion"]
        assert calls == [("completion", ["hello"])]

    def test_text_only_falls_back_to_prefixed_completion_without_template(self):
        handler = self._make_handler()
        ctx = self._make_context(texts=["hello"], input_type="query")
        calls: list[tuple[str, object]] = []

        def preprocess_cmpl(request, prompt_input, prompt_embeds):
            calls.append(("completion", prompt_input))
            return ["fallback"]

        handler._get_task_instruction_prefix = lambda _input_type: "query: "
        handler._has_chat_template = lambda: False
        handler._batch_render_chat = lambda *_args, **_kwargs: pytest.fail(
            "chat rendering should be skipped without a template"
        )
        handler._preprocess_cmpl_online = preprocess_cmpl

        handler._pre_process_cohere_online(ctx)

        assert ctx.engine_inputs == ["fallback"]
        assert calls == [("completion", ["query: hello"])]

    def test_text_only_with_template_uses_chat_path(self):
        handler = self._make_handler()
        ctx = self._make_context(texts=["hello"], input_type="query")
        calls: list[tuple[str, object]] = []

        def batch_render_chat(
            request,
            all_messages,
            truncate_prompt_tokens,
            truncation_side,
        ):
            calls.append(
                (
                    "chat",
                    {
                        "request": request,
                        "all_messages": all_messages,
                        "truncate_prompt_tokens": truncate_prompt_tokens,
                        "truncation_side": truncation_side,
                    },
                )
            )
            return ["chat"]

        handler._get_task_instruction_prefix = lambda _input_type: "query: "
        handler._has_chat_template = lambda: True
        handler._batch_render_chat = batch_render_chat
        handler._preprocess_cmpl_online = lambda *_args, **_kwargs: pytest.fail(
            "completion path should be skipped when a template exists"
        )

        handler._pre_process_cohere_online(ctx)

        assert ctx.engine_inputs == ["chat"]
        assert calls == [
            (
                "chat",
                {
                    "request": ctx.request,
                    "all_messages": [
                        handler._mixed_input_to_messages(
                            CohereEmbedInput(
                                content=[CohereEmbedContent(type="text", text="hello")]
                            ),
                            task_prefix="query: ",
                        )
                    ],
                    "truncate_prompt_tokens": -1,
                    "truncation_side": None,
                },
            )
        ]


class TestPreProcessOpenAIEmbeddingChatOnline:
    """Unit tests for OpenAI embedding chat preprocessing."""

    class _FakeModelConfig:
        max_model_len = 128
        encoder_config: dict[str, object] = {}
        pooler_config = None
        multimodal_config = None
        is_encoder_decoder = False

    class _FakeRenderer:
        tokenizer = object()

        def __init__(self):
            self.calls = []

        def render_chat(
            self,
            all_messages,
            chat_params,
            tok_params,
            prompt_extras=None,
        ):
            self.calls.append(
                {
                    "all_messages": all_messages,
                    "chat_params": chat_params,
                    "tok_params": tok_params,
                    "prompt_extras": prompt_extras,
                }
            )
            return all_messages, [
                {"prompt_token_ids": [index]} for index, _ in enumerate(all_messages)
            ]

    @classmethod
    def _make_handler(cls, renderer):
        handler = object.__new__(EmbedIOProcessor)
        handler.renderer = renderer
        handler.model_config = cls._FakeModelConfig()
        handler.chat_template = "template"
        handler.chat_template_content_format = "auto"
        handler.trust_request_chat_template = False
        handler.enable_chunked_processing = False
        return handler

    @staticmethod
    def _make_context(
        request: (
            EmbeddingChatRequest
            | EmbeddingBatchChatRequest
            | EmbeddingChatInputRequest
            | EmbeddingBatchChatInputRequest
        ),
    ) -> PoolingServeContext[
        EmbeddingChatRequest
        | EmbeddingBatchChatRequest
        | EmbeddingChatInputRequest
        | EmbeddingBatchChatInputRequest
    ]:
        return PoolingServeContext(
            request=request,
            pooling_params=PoolingParams(),
            model_name="test",
            request_id="embd-test",
        )

    def test_chat_template_kwargs_forwarded_for_batched_input_messages(self):
        request = TypeAdapter(EmbeddingRequest).validate_python(
            {
                "model": "test",
                "input": [
                    [{"role": "user", "content": "hello"}],
                    [{"role": "user", "content": "goodbye"}],
                ],
                "add_generation_prompt": True,
                "chat_template_kwargs": {"instruction": "Represent the query: "},
                "mm_processor_kwargs": {"max_pixels": 1},
                "cache_salt": "salt",
            }
        )
        assert isinstance(request, EmbeddingBatchChatInputRequest)

        renderer = self._FakeRenderer()
        handler = self._make_handler(renderer)
        ctx = self._make_context(request)

        handler.pre_process_online(ctx)

        assert ctx.engine_inputs == [
            {"prompt_token_ids": [0]},
            {"prompt_token_ids": [1]},
        ]
        assert len(renderer.calls) == 1

        call = renderer.calls[0]
        assert call["all_messages"] == request.messages
        assert call["prompt_extras"] == {
            "mm_processor_kwargs": {"max_pixels": 1},
            "cache_salt": "salt",
        }

        chat_template_kwargs = call["chat_params"].chat_template_kwargs
        assert chat_template_kwargs["instruction"] == "Represent the query: "
        assert chat_template_kwargs["add_generation_prompt"] is True
        assert chat_template_kwargs["continue_final_message"] is False
        assert "tools" not in chat_template_kwargs
        assert chat_template_kwargs["tokenize"] is False