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

import io
from dataclasses import dataclass
from typing import Any
from unittest.mock import AsyncMock

import pybase64
import pytest
import torch

from vllm.inputs.data import is_embeds_prompt
from vllm.renderers import TokenizeParams
from vllm.renderers.hf import HfRenderer
from vllm.tokenizers.registry import tokenizer_args_from_config

MODEL_NAME = "openai-community/gpt2"


@dataclass
class MockHFConfig:
    model_type: str = "any"


@dataclass
class MockModelConfig:
    runner_type = "generate"
    model: str = MODEL_NAME
    tokenizer: str = MODEL_NAME
    trust_remote_code: bool = False
    max_model_len: int = 100
    tokenizer_revision = None
    tokenizer_mode = "auto"
    hf_config = MockHFConfig()
    encoder_config: dict[str, Any] | None = None
    enable_prompt_embeds: bool = True
    skip_tokenizer_init: bool = False


@pytest.fixture
def mock_model_config():
    return MockModelConfig()


@pytest.fixture
def mock_async_tokenizer():
    return AsyncMock()


@pytest.fixture
def renderer(mock_model_config):
    _, tokenizer_name, _, kwargs = tokenizer_args_from_config(mock_model_config)

    return HfRenderer(
        mock_model_config,
        tokenizer_kwargs={**kwargs, "tokenizer_name": tokenizer_name},
    )


class TestValidatePrompt:
    STRING_INPUTS = [
        "",
        "foo",
        "foo bar",
        "foo baz bar",
        "foo bar qux baz",
    ]

    TOKEN_INPUTS = [
        [-1],
        [1],
        [1, 2],
        [1, 3, 4],
        [1, 2, 4, 3],
    ]

    INPUTS_SLICES = [
        slice(None, None, -1),
        slice(None, None, 2),
        slice(None, None, -2),
    ]

    # Test that a nested mixed-type list of lists raises a TypeError.
    def test_empty_input(self, renderer):
        with pytest.raises(ValueError, match="at least one prompt"):
            renderer.render_completions([])

    def test_invalid_type(self, renderer):
        with pytest.raises(TypeError, match="string or an array of tokens"):
            renderer.render_completions([[1, 2], ["foo", "bar"]])

    @pytest.mark.parametrize("string_input", STRING_INPUTS)
    def test_string_consistent(self, renderer, string_input: str):
        assert renderer.render_completions(string_input) == renderer.render_completions(
            [string_input]
        )

    @pytest.mark.parametrize("token_input", TOKEN_INPUTS)
    def test_token_consistent(self, renderer, token_input: list[int]):
        assert renderer.render_completions(token_input) == renderer.render_completions(
            [token_input]
        )

    @pytest.mark.parametrize("inputs_slice", INPUTS_SLICES)
    def test_string_slice(self, renderer, inputs_slice: slice):
        assert renderer.render_completions(self.STRING_INPUTS)[
            inputs_slice
        ] == renderer.render_completions(self.STRING_INPUTS[inputs_slice])


class TestRenderPrompt:
    @pytest.mark.asyncio
    async def test_token_input(self, renderer):
        tokens = [101, 7592, 2088]
        prompts = await renderer.render_completions_async(tokens)
        results = await renderer.tokenize_prompts_async(
            prompts,
            TokenizeParams(max_total_tokens=100),
        )

        assert len(results) == 1
        assert results[0]["prompt_token_ids"] == tokens

    @pytest.mark.asyncio
    async def test_token_list_input(self, renderer):
        token_lists = [[101, 7592, 2088], [102, 1234, 5678, 9012], [103, 4567]]
        prompts = await renderer.render_completions_async(token_lists)
        results = await renderer.tokenize_prompts_async(
            prompts,
            TokenizeParams(max_total_tokens=100),
        )

        assert len(results) == 3
        assert results[0]["prompt_token_ids"] == [101, 7592, 2088]
        assert results[1]["prompt_token_ids"] == [102, 1234, 5678, 9012]
        assert results[2]["prompt_token_ids"] == [103, 4567]

    @pytest.mark.asyncio
    async def test_text_input(self, renderer, mock_async_tokenizer):
        mock_async_tokenizer.encode.return_value = [101, 7592, 2088]
        renderer._async_tokenizer = mock_async_tokenizer

        prompts = await renderer.render_completions_async("Hello world")
        results = await renderer.tokenize_prompts_async(
            prompts,
            TokenizeParams(max_total_tokens=100),
        )

        assert len(results) == 1
        assert results[0]["prompt_token_ids"] == [101, 7592, 2088]
        mock_async_tokenizer.encode.assert_called_once()

    @pytest.mark.asyncio
    async def test_text_list_input(self, renderer, mock_async_tokenizer):
        mock_async_tokenizer.encode.return_value = [101, 7592, 2088]
        renderer._async_tokenizer = mock_async_tokenizer

        text_list_input = ["Hello world", "How are you?", "Good morning"]
        prompts = await renderer.render_completions_async(text_list_input)
        results = await renderer.tokenize_prompts_async(
            prompts,
            TokenizeParams(max_total_tokens=100),
        )

        assert len(results) == 3
        for result in results:
            assert result["prompt_token_ids"] == [101, 7592, 2088]
        assert mock_async_tokenizer.encode.call_count == 3

    @pytest.mark.asyncio
    async def test_no_truncation(self, renderer, mock_async_tokenizer):
        mock_async_tokenizer.encode.return_value = [101, 7592, 2088]
        renderer._async_tokenizer = mock_async_tokenizer

        prompts = await renderer.render_completions_async("Hello world")
        results = await renderer.tokenize_prompts_async(
            prompts,
            TokenizeParams(max_total_tokens=100),
        )

        assert len(results) == 1
        call_args = mock_async_tokenizer.encode.call_args
        assert (
            "truncation" not in call_args.kwargs
            or call_args.kwargs["truncation"] is False
        )

    @pytest.mark.asyncio
    async def test_truncation_positive(self, renderer, mock_async_tokenizer):
        mock_async_tokenizer.encode.return_value = [101, 7592, 2088]  # Truncated
        renderer._async_tokenizer = mock_async_tokenizer

        prompts = await renderer.render_completions_async("Hello world")
        results = await renderer.tokenize_prompts_async(
            prompts,
            TokenizeParams(
                max_total_tokens=200,
                truncate_prompt_tokens=50,
            ),
        )

        assert len(results) == 1
        call_args = mock_async_tokenizer.encode.call_args
        assert call_args.kwargs["truncation"] is True
        assert call_args.kwargs["max_length"] == 50

    @pytest.mark.asyncio
    async def test_truncation_negative(self, renderer, mock_async_tokenizer):
        # Test that negative truncation uses model's max_model_len
        mock_async_tokenizer.encode.return_value = [
            101,
            7592,
            2088,
        ]  # Truncated to max_model_len
        renderer._async_tokenizer = mock_async_tokenizer

        prompts = await renderer.render_completions_async("Hello world")
        results = await renderer.tokenize_prompts_async(
            prompts,
            TokenizeParams(
                max_total_tokens=200,
                truncate_prompt_tokens=-1,
            ),
        )

        assert len(results) == 1
        call_args = mock_async_tokenizer.encode.call_args
        assert call_args.kwargs["truncation"] is True
        assert call_args.kwargs["max_length"] == 200

    @pytest.mark.asyncio
    async def test_token_truncation_last_elements(self, renderer):
        # Test that token truncation keeps the last N elements
        long_tokens = [100, 101, 102, 103, 104, 105, 106, 107, 108, 109]  # 10 tokens
        prompts = await renderer.render_completions_async(long_tokens)
        results = await renderer.tokenize_prompts_async(
            prompts,
            TokenizeParams(
                max_total_tokens=100,
                truncate_prompt_tokens=5,
            ),
        )

        assert len(results) == 1
        # Should keep the last 5 tokens: [105, 106, 107, 108, 109]
        assert results[0]["prompt_token_ids"] == [105, 106, 107, 108, 109]

    @pytest.mark.asyncio
    async def test_max_length_exceeded(self, renderer):
        long_tokens = list(range(150))  # Exceeds max_model_len=100

        prompts = await renderer.render_completions_async(long_tokens)

        with pytest.raises(ValueError, match="context length is only"):
            await renderer.tokenize_prompts_async(
                prompts,
                TokenizeParams(max_total_tokens=100),
            )

    @pytest.mark.asyncio
    async def test_no_tokenizer_for_text(self, renderer):
        renderer_no_tokenizer = HfRenderer.from_config(
            MockModelConfig(skip_tokenizer_init=True),
            tokenizer_kwargs={},
        )

        prompts = await renderer_no_tokenizer.render_completions_async("Hello world")

        with pytest.raises(ValueError, match="`skip_tokenizer_init=True`"):
            await renderer_no_tokenizer.tokenize_prompts_async(
                prompts,
                TokenizeParams(max_total_tokens=100),
            )

    @pytest.mark.asyncio
    async def test_token_input_with_needs_detokenization(
        self, renderer, mock_async_tokenizer
    ):
        # When needs_detokenization=True for token inputs, renderer should
        # use the async tokenizer to decode and include the original text
        # in the returned prompt object.
        mock_async_tokenizer.decode = AsyncMock(return_value="decoded text")
        renderer._async_tokenizer = mock_async_tokenizer

        tokens = [1, 2, 3, 4]
        prompts = await renderer.render_completions_async(tokens)
        results = await renderer.tokenize_prompts_async(
            prompts,
            TokenizeParams(
                max_total_tokens=renderer.config.max_model_len,
                needs_detokenization=True,
            ),
        )

        assert len(results) == 1
        assert results[0]["prompt_token_ids"] == tokens
        assert results[0]["prompt"] == "decoded text"
        mock_async_tokenizer.decode.assert_awaited_once()


class TestRenderEmbedPrompt:
    def _create_test_embed_bytes(self, tensor: torch.Tensor) -> bytes:
        """Helper to create base64-encoded tensor bytes"""
        buffer = io.BytesIO()
        torch.save(tensor, buffer)
        buffer.seek(0)
        return pybase64.b64encode(buffer.read())

    @pytest.mark.asyncio
    async def test_single_prompt_embed(self, renderer):
        # Create a test tensor
        test_tensor = torch.randn(10, 768, dtype=torch.float32)
        embed_bytes = self._create_test_embed_bytes(test_tensor)

        prompts = await renderer.render_completions_async(prompt_embeds=embed_bytes)
        results = await renderer.tokenize_prompts_async(
            prompts,
            TokenizeParams(max_total_tokens=renderer.config.max_model_len),
        )

        assert len(results) == 1
        assert is_embeds_prompt(results[0])
        assert torch.allclose(results[0]["prompt_embeds"], test_tensor)

    @pytest.mark.asyncio
    async def test_multiple_prompt_embeds(self, renderer):
        # Create multiple test tensors
        test_tensors = [
            torch.randn(8, 512, dtype=torch.float32),
            torch.randn(12, 512, dtype=torch.float32),
        ]
        embed_bytes_list = [self._create_test_embed_bytes(t) for t in test_tensors]

        prompts = await renderer.render_completions_async(
            prompt_embeds=embed_bytes_list
        )
        results = await renderer.tokenize_prompts_async(
            prompts,
            TokenizeParams(max_total_tokens=renderer.config.max_model_len),
        )

        assert len(results) == 2
        for i, result in enumerate(results):
            assert is_embeds_prompt(result)
            assert torch.allclose(result["prompt_embeds"], test_tensors[i])

    @pytest.mark.asyncio
    async def test_prompt_embed_truncation(self, renderer):
        # Create tensor with more tokens than truncation limit
        test_tensor = torch.randn(20, 768, dtype=torch.float32)
        embed_bytes = self._create_test_embed_bytes(test_tensor)

        prompts = await renderer.render_completions_async(prompt_embeds=embed_bytes)
        results = await renderer.tokenize_prompts_async(
            prompts,
            TokenizeParams(
                max_total_tokens=renderer.config.max_model_len,
                truncate_prompt_tokens=10,
            ),
        )

        assert len(results) == 1
        # Should keep last 10 tokens
        expected = test_tensor[-10:]
        assert torch.allclose(results[0]["prompt_embeds"], expected)

    @pytest.mark.asyncio
    async def test_prompt_embed_different_dtypes(self, renderer):
        # Test different supported dtypes
        dtypes = [torch.float32, torch.float16, torch.bfloat16]

        for dtype in dtypes:
            test_tensor = torch.randn(5, 256, dtype=dtype)
            embed_bytes = self._create_test_embed_bytes(test_tensor)

            prompts = await renderer.render_completions_async(prompt_embeds=embed_bytes)
            results = await renderer.tokenize_prompts_async(
                prompts,
                TokenizeParams(max_total_tokens=renderer.config.max_model_len),
            )

            assert len(results) == 1
            assert results[0]["prompt_embeds"].dtype == dtype

    @pytest.mark.asyncio
    async def test_prompt_embed_squeeze_batch_dim(self, renderer):
        # Test tensor with batch dimension gets squeezed
        test_tensor = torch.randn(1, 10, 768, dtype=torch.float32)
        embed_bytes = self._create_test_embed_bytes(test_tensor)

        prompts = await renderer.render_completions_async(prompt_embeds=embed_bytes)
        results = await renderer.tokenize_prompts_async(
            prompts,
            TokenizeParams(max_total_tokens=renderer.config.max_model_len),
        )

        assert len(results) == 1
        # Should be squeezed to 2D
        assert results[0]["prompt_embeds"].shape == (10, 768)

    @pytest.mark.asyncio
    async def test_both_prompts_and_embeds(self, renderer, mock_async_tokenizer):
        # Set up text tokenization
        mock_async_tokenizer.encode.return_value = [101, 102, 103]
        renderer._async_tokenizer = mock_async_tokenizer

        # Create embed
        test_tensor = torch.randn(5, 256, dtype=torch.float32)
        embed_bytes = self._create_test_embed_bytes(test_tensor)

        prompts = await renderer.render_completions_async(
            "Hello world",
            prompt_embeds=embed_bytes,
        )
        results = await renderer.tokenize_prompts_async(
            prompts,
            TokenizeParams(max_total_tokens=renderer.config.max_model_len),
        )

        assert len(results) == 2
        # First should be embed prompt
        assert is_embeds_prompt(results[0])
        # Second should be tokens prompt
        assert "prompt_token_ids" in results[1]
        assert results[1]["prompt_token_ids"] == [101, 102, 103]