# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import weakref

import pytest
import torch

from tests.models.utils import check_embeddings_close
from vllm import LLM, PoolingParams, PoolingRequestOutput
from vllm.config import PoolerConfig
from vllm.exceptions import VLLMValidationError
from vllm.pooling_params import LateChunkingParams
from vllm.tasks import PoolingTask

MODEL_NAME = "intfloat/multilingual-e5-small"

prompt = "The chef prepared a delicious meal."
prompt_token_ids = [0, 581, 21861, 133888, 10, 8, 150, 60744, 109911, 5, 2]
embedding_size = 384


@pytest.fixture(scope="module")
def llm(vllm_runner):
    with vllm_runner(
        MODEL_NAME,
        max_model_len=None,
        pooler_config=PoolerConfig(task="token_embed"),
        max_num_batched_tokens=32768,
        tensor_parallel_size=1,
        gpu_memory_utilization=0.75,
        enforce_eager=True,
        seed=0,
        enable_chunked_prefill=None,
    ) as runner:
        assert embedding_size == runner.llm.model_config.embedding_size
        # pytest caches yielded fixtures until after teardown, so use a proxy to
        # avoid retaining the LLM while VllmRunner.__exit__ releases ROCm memory.
        yield weakref.proxy(runner.llm)


@pytest.mark.skip_global_cleanup
def test_str_prompts(llm: LLM):
    outputs = llm.encode(prompt, pooling_task="token_embed", use_tqdm=False)
    assert len(outputs) == 1
    assert isinstance(outputs[0], PoolingRequestOutput)
    assert outputs[0].outputs.data.shape == (11, 384)


@pytest.mark.skip_global_cleanup
def test_token_ids_prompts(llm: LLM):
    outputs = llm.encode([prompt_token_ids], pooling_task="token_embed", use_tqdm=False)
    assert len(outputs) == 1
    assert isinstance(outputs[0], PoolingRequestOutput)
    assert outputs[0].outputs.data.shape == (11, 384)


@pytest.mark.parametrize("task", ["embed", "classify", "token_classify", "plugin"])
def test_unsupported_tasks(llm: LLM, task: PoolingTask, caplog_vllm):
    if task == "plugin":
        err_msg = "No IOProcessor plugin installed."
    elif task == "embed":
        err_msg = "Try switching the model's pooling_task via.+"
    else:
        err_msg = "Classification API is not supported by this model.+"

    with pytest.raises(ValueError, match=err_msg):
        llm.encode(prompt, pooling_task=task, use_tqdm=False)


def test_bge_m3_late_chunking_keeps_special_tokens(vllm_runner):
    with vllm_runner(
        "BAAI/bge-m3",
        hf_overrides={"architectures": ["BgeM3EmbeddingModel"]},
        pooler_config=PoolerConfig(task="token_embed"),
        dtype="float32",
        max_model_len=128,
        # Keep encoder batch size fixed for the numerical reference.
        max_num_seqs=1,
        gpu_memory_utilization=0.5,
        enforce_eager=True,
    ) as runner:
        llm = runner.llm
        text = "A document about Berlin and its museums."
        sizes = (1, 3, 128, None)
        tokens, *outputs = llm.encode(
            [text] * len(sizes),
            pooling_task="token_embed",
            pooling_params=[
                PoolingParams(
                    use_activation=size not in (1, None),
                    late_chunking_params=LateChunkingParams(size) if size else None,
                )
                for size in sizes
            ],
            use_tqdm=False,
        )
        assert tokens.outputs.data.shape[0] == len(tokens.prompt_token_ids)
        assert tokens.late_chunking.chunks[0].char_range is None
        for size, output in zip(sizes[1:], outputs):
            if size is None:
                # Ordinary BGE-M3 token embeddings still drop BOS only.
                expected = tokens.outputs.data[1:]
                assert output.late_chunking is None
            else:
                # BGE-M3's projector is linear, so its unnormalized token
                # outputs also provide a reference for projected chunk means.
                expected = torch.nn.functional.normalize(
                    torch.stack([p.mean(0) for p in tokens.outputs.data.split(size)]),
                    dim=-1,
                )
                assert len(output.late_chunking.chunks) == len(expected)
            torch.testing.assert_close(
                output.outputs.data, expected, atol=2e-5, rtol=2e-4
            )


def test_nomic_late_chunking_without_special_tokens(vllm_runner):
    """Keep source ranges without special tokens and reject empty tokenization."""
    with vllm_runner(
        "nomic-ai/nomic-embed-text-v1",
        revision="720244025c1a7e15661a174c63cce63c8218e52b",
        trust_remote_code=True,
        runner="pooling",
        pooler_config=PoolerConfig(task="token_embed"),
        dtype="float32",
        max_model_len=128,
        max_num_seqs=2,
        gpu_memory_utilization=0.3,
        enforce_eager=True,
        enable_prefix_caching=False,
        enable_chunked_prefill=False,
    ) as runner:
        llm = runner.llm
        text = "search_document: 中文 😀 café e\u0301"
        params = PoolingParams(
            late_chunking_params=LateChunkingParams(3), use_activation=False
        )
        kwargs = {"add_special_tokens": False}
        raw = llm.encode(
            text,
            pooling_task="token_embed",
            use_tqdm=False,
            pooling_params=PoolingParams(use_activation=False),
            tokenization_kwargs=kwargs,
        )[0]
        actual = llm.encode(
            text,
            pooling_task="token_embed",
            use_tqdm=False,
            pooling_params=params,
            tokenization_kwargs=kwargs,
        )[0]
        expected = torch.stack([part.mean(0) for part in raw.outputs.data.split(3)])
        torch.testing.assert_close(actual.outputs.data, expected, atol=1e-5, rtol=1e-4)
        encoded = llm.get_tokenizer()(text, return_offsets_mapping=True, **kwargs)
        assert actual.prompt_token_ids == encoded["input_ids"]
        for chunk in actual.late_chunking.chunks:
            start, end = chunk.token_range
            offsets = [(a, b) for a, b in encoded["offset_mapping"][start:end] if a < b]
            assert chunk.char_range == (offsets[0][0], offsets[-1][1])
        for rejected in ("long " * 256, "   "):
            with pytest.raises(VLLMValidationError):
                llm.encode(
                    rejected,
                    pooling_task="token_embed",
                    use_tqdm=False,
                    pooling_params=params,
                    tokenization_kwargs=kwargs,
                )
        assert (
            llm.encode(text, pooling_task="token_embed", use_tqdm=False)[
                0
            ].late_chunking
            is None
        )


@pytest.mark.skip_global_cleanup
@pytest.mark.parametrize("normalize", [False, True])
def test_e5_late_chunking_mixed_requests(llm: LLM, normalize: bool):
    texts = [prompt, "Short text.", "中文 😀 café"]
    raw = llm.encode(
        texts,
        pooling_task="token_embed",
        use_tqdm=False,
        pooling_params=PoolingParams(use_activation=False),
    )
    sizes = [3, None, 100]
    params = [
        PoolingParams(
            late_chunking_params=LateChunkingParams(size) if size is not None else None,
            use_activation=normalize,
        )
        for size in sizes
    ]
    actual = llm.encode(
        texts, pooling_task="token_embed", pooling_params=params, use_tqdm=False
    )
    for original, result, size in zip(raw, actual, sizes):
        expected = original.outputs.data
        if size is not None:
            expected = torch.stack(
                [part.float().mean(0) for part in expected.split(size)]
            )
            assert result.late_chunking is not None
            assert len(result.late_chunking.chunks) == len(expected)
            assert result.late_chunking.input_tokens == len(result.prompt_token_ids)
        else:
            assert result.late_chunking is None
        if normalize:
            expected = torch.nn.functional.normalize(expected, dim=-1)
        check_embeddings_close(
            embeddings_0_lst=result.outputs.data.tolist(),
            embeddings_1_lst=expected.tolist(),
            name_0="late_chunking",
            name_1="token_reference",
        )
        torch.testing.assert_close(
            result.outputs.data.norm(dim=-1),
            expected.norm(dim=-1),
            atol=1e-3,
            rtol=1e-3,
        )


@pytest.mark.parametrize("use_v2_runner", [False, True])
def test_late_chunking_recomputes_cached_prefix_with_chunked_prefill(
    vllm_runner, monkeypatch, use_v2_runner
):
    # Dummy weights suffice for testing scheduling and reduction of the same states.
    monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1" if use_v2_runner else "0")
    with vllm_runner(
        "openai-community/gpt2",
        runner="pooling",
        convert="embed",
        load_format="dummy",
        pooler_config=PoolerConfig(
            task="token_embed", seq_pooling_type="LAST", tok_pooling_type="ALL"
        ),
        dtype="float32",
        max_model_len=128,
        max_num_batched_tokens=16,
        max_num_seqs=2,
        enable_prefix_caching=True,
        enable_chunked_prefill=True,
        enforce_eager=True,
        gpu_memory_utilization=0.3,
        seed=0,
    ) as runner:
        text = "The document has several important words. " * 6
        llm = runner.llm
        raw = llm.encode(
            text,
            pooling_task="token_embed",
            use_tqdm=False,
            pooling_params=PoolingParams(use_activation=False),
        )[0]
        assert len(raw.prompt_token_ids) > 16
        cached = llm.encode(
            text,
            pooling_task="token_embed",
            use_tqdm=False,
            pooling_params=PoolingParams(
                use_activation=False, skip_reading_prefix_cache=False
            ),
        )[0]
        assert cached.num_cached_tokens > 0
        params = PoolingParams(
            late_chunking_params=LateChunkingParams(7),
            use_activation=False,
            skip_reading_prefix_cache=False,
        )
        actual = llm.encode(
            text, pooling_task="token_embed", pooling_params=params, use_tqdm=False
        )[0]
        assert actual.num_cached_tokens == 0
        assert actual.prompt_token_ids == raw.prompt_token_ids
        expected = torch.stack([part.mean(0) for part in raw.outputs.data.split(7)])
        torch.testing.assert_close(actual.outputs.data, expected, atol=1e-5, rtol=1e-4)
        assert actual.late_chunking.input_tokens == len(raw.prompt_token_ids)
        assert params.skip_reading_prefix_cache is False
        assert params.late_chunking_params is not None