# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Tests for the Generative Scoring API.

Tests cover:
1. Protocol models (request/response construction)
2. Probability computation (softmax normalization)
3. Input validation
4. Score formula: P(token[0]) / (P(token[0]) + P(token[1]))
5. Prompt building and item ordering
"""

import math
from dataclasses import dataclass, field
from typing import Any
from unittest.mock import MagicMock

import pytest

from vllm.config.multimodal import MultiModalConfig
from vllm.entrypoints.generate.generative_scoring.serving import (
    GenerativeScoringItemResult,
    GenerativeScoringRequest,
    GenerativeScoringResponse,
    ServingGenerativeScoring,
)
from vllm.entrypoints.openai.engine.protocol import ErrorResponse
from vllm.entrypoints.openai.models.protocol import BaseModelPath
from vllm.entrypoints.openai.models.serving import OpenAIServingModels
from vllm.logprobs import Logprob
from vllm.outputs import CompletionOutput, RequestOutput
from vllm.tokenizers import get_tokenizer
from vllm.v1.engine.async_llm import AsyncLLM

MODEL_NAME = "Qwen/Qwen3-0.6B"
BASE_MODEL_PATHS = [BaseModelPath(name=MODEL_NAME, model_path=MODEL_NAME)]


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


@dataclass
class MockModelConfig:
    task = "generate"
    runner_type = "generate"
    tokenizer = MODEL_NAME
    trust_remote_code = False
    tokenizer_mode = "auto"
    max_model_len = 100
    tokenizer_revision = None
    multimodal_config = MultiModalConfig()
    hf_config = MockHFConfig()
    logits_processor_pattern = None
    logits_processors: list[str] | None = None
    diff_sampling_param: dict | None = None
    allowed_local_media_path: str = ""
    allowed_media_domains: list[str] | None = None
    encoder_config = None
    generation_config: str = "auto"
    media_io_kwargs: dict[str, dict[str, Any]] = field(default_factory=dict)
    skip_tokenizer_init = False
    vocab_size = 151936

    def get_diff_sampling_param(self):
        return self.diff_sampling_param or {}

    def get_vocab_size(self):
        return self.vocab_size


def _create_mock_engine():
    """Create a mock AsyncLLM engine."""
    mock_engine = MagicMock(spec=AsyncLLM)
    mock_engine.get_tokenizer.return_value = get_tokenizer(MODEL_NAME)
    mock_engine.errored = False
    mock_engine.model_config = MockModelConfig()
    mock_engine.input_processor = MagicMock()

    # renderer is accessed by OpenAIServing.__init__ and serving.py
    mock_renderer = MagicMock()
    mock_renderer.tokenizer = get_tokenizer(MODEL_NAME)
    mock_engine.renderer = mock_renderer

    return mock_engine


def _create_serving(mock_engine) -> ServingGenerativeScoring:
    """Create an ServingGenerativeScoring instance with mocks."""
    models = OpenAIServingModels(
        engine_client=mock_engine,
        base_model_paths=BASE_MODEL_PATHS,
    )
    return ServingGenerativeScoring(mock_engine, models, request_logger=None)


def _create_mock_request_output(logprobs_dict: dict[int, float]) -> RequestOutput:
    """Create a mock RequestOutput with specified logprobs."""
    logprobs_with_objs = {
        tid: Logprob(logprob=lp, rank=i + 1)
        for i, (tid, lp) in enumerate(logprobs_dict.items())
    }
    completion_output = CompletionOutput(
        index=0,
        text="",
        token_ids=[100],
        cumulative_logprob=-1.0,
        logprobs=[logprobs_with_objs],
        finish_reason="length",
    )
    return RequestOutput(
        request_id="test-request",
        prompt="test prompt",
        prompt_token_ids=[1, 2, 3],
        prompt_logprobs=None,
        outputs=[completion_output],
        finished=True,
    )


class TestProtocolModels:
    """Tests for GenerativeScoringRequest and GenerativeScoringResponse."""

    def test_request_and_response_all_fields(self):
        """Test request construction with all field types and response structure."""
        # Test request with string inputs
        req_str = GenerativeScoringRequest(
            query="Is this the capital?",
            items=["Paris", "London"],
            label_token_ids=[9454, 2753],
        )
        assert req_str.query == "Is this the capital?"
        assert req_str.items == ["Paris", "London"]
        assert req_str.label_token_ids == [9454, 2753]
        assert req_str.apply_softmax is True  # default
        assert req_str.item_first is False  # default
        assert req_str.add_special_tokens is True  # default

        # Test request with pre-tokenized inputs and custom options
        req_tok = GenerativeScoringRequest(
            query=[100, 200, 300],
            items=[[400, 500], [600, 700]],
            label_token_ids=[1234, 5678],
            apply_softmax=False,
            item_first=True,
            add_special_tokens=False,
        )
        assert req_tok.query == [100, 200, 300]
        assert req_tok.items == [[400, 500], [600, 700]]
        assert req_tok.apply_softmax is False
        assert req_tok.item_first is True
        assert req_tok.add_special_tokens is False

        # Test response structure
        response = GenerativeScoringResponse(
            model="test-model",
            data=[
                GenerativeScoringItemResult(index=0, score=0.7),
                GenerativeScoringItemResult(index=1, score=0.4),
            ],
            usage={"prompt_tokens": 10, "total_tokens": 12, "completion_tokens": 2},
        )
        assert response.object == "list"
        assert response.model == "test-model"
        assert len(response.data) == 2
        assert response.data[0].score == 0.7
        assert response.data[0].object == "score"
        assert response.data[1].score == 0.4
        assert response.usage.prompt_tokens == 10


class TestProbabilityComputation:
    """Tests for _compute_probabilities with both softmax modes."""

    @pytest.mark.parametrize(
        "label_logprobs,apply_softmax,should_sum_to_one",
        [
            ({100: -1.0, 200: -2.0}, True, True),
            ({100: -100.0, 200: -100.5}, True, True),  # numerical stability
            ({100: -1.0, 200: -2.0}, False, False),
        ],
        ids=["softmax_basic", "softmax_extreme_values", "true_probs"],
    )
    def test_compute_probabilities(
        self, label_logprobs, apply_softmax, should_sum_to_one
    ):
        """Test probability computation for softmax and true probability modes."""
        serving = ServingGenerativeScoring.__new__(ServingGenerativeScoring)
        probs = serving._compute_probabilities(
            label_logprobs, apply_softmax=apply_softmax
        )

        # Verify sum behavior
        total = sum(probs.values())
        if should_sum_to_one:
            assert abs(total - 1.0) < 1e-6
        else:
            assert total < 1.0

        # Verify math
        if apply_softmax:
            max_lp = max(label_logprobs.values())
            exp_vals = {k: math.exp(v - max_lp) for k, v in label_logprobs.items()}
            sum_exp = sum(exp_vals.values())
            for tid, lp in label_logprobs.items():
                assert abs(probs[tid] - exp_vals[tid] / sum_exp) < 1e-9
        else:
            for tid, lp in label_logprobs.items():
                assert abs(probs[tid] - math.exp(lp)) < 1e-9

    def test_score_formula(self):
        """Test the score formula: P(token[0]) / (P(token[0]) + P(token[1]))."""
        serving = ServingGenerativeScoring.__new__(ServingGenerativeScoring)

        # With logprobs -0.5 and -2.0, softmax gives higher prob to first token
        logprobs = {9454: -0.5, 2753: -2.0}
        probs = serving._compute_probabilities(logprobs, apply_softmax=True)

        # Score = P(9454) / (P(9454) + P(2753)) = P(9454) since they sum to 1
        score = probs[9454]

        # Manual calculation
        exp_0 = math.exp(-0.5)
        exp_1 = math.exp(-2.0)
        expected_score = exp_0 / (exp_0 + exp_1)

        assert abs(score - expected_score) < 1e-9
        assert score > 0.5  # First token has higher logprob, so higher probability


class TestValidation:
    """Tests for input validation errors."""

    @pytest.mark.asyncio
    @pytest.mark.parametrize(
        "request_kwargs,expected_error",
        [
            (
                {"query": "q", "items": ["i"], "label_token_ids": [999999, 999998]},
                "out of vocabulary",
            ),
            (
                {"query": "q", "items": [], "label_token_ids": [100, 200]},
                "at least one item",
            ),
        ],
        ids=["invalid_token_id", "empty_items"],
    )
    async def test_validation_errors(self, request_kwargs, expected_error):
        """Test that invalid inputs return appropriate errors."""
        mock_engine = _create_mock_engine()
        serving = _create_serving(mock_engine)
        request = GenerativeScoringRequest(model=MODEL_NAME, **request_kwargs)
        result = await serving.create_generative_scoring(request, None)

        assert isinstance(result, ErrorResponse)
        assert expected_error in result.error.message.lower()


class TestPromptBuilding:
    """Tests for prompt construction and item ordering."""

    @pytest.mark.asyncio
    @pytest.mark.parametrize(
        "item_first,expected",
        [
            (False, [[100, 101, 200, 201], [100, 101, 300, 301]]),  # query + item
            (True, [[200, 201, 100, 101], [300, 301, 100, 101]]),  # item + query
        ],
        ids=["query_first", "item_first"],
    )
    async def test_item_ordering(self, item_first, expected):
        """Test that item_first flag controls prompt concatenation order."""
        mock_engine = _create_mock_engine()
        serving = _create_serving(mock_engine)

        request = GenerativeScoringRequest(
            query=[100, 101],
            items=[[200, 201], [300, 301]],
            label_token_ids=[500, 501],
            item_first=item_first,
        )
        engine_inputs, _ = await serving._build_prompts(
            request, MagicMock(), max_model_len=4096
        )

        for i, exp in enumerate(expected):
            assert engine_inputs[i]["prompt_token_ids"] == exp


class TestGeneration:
    """Tests for the full generation flow with mocked engine."""

    @pytest.mark.asyncio
    async def test_successful_generation(self):
        """Test successful score generation returns valid response."""
        mock_engine = _create_mock_engine()
        serving = _create_serving(mock_engine)

        mock_logprobs = {1234: -0.5, 5678: -2.0, 100: -3.0}
        mock_output = _create_mock_request_output(mock_logprobs)

        async def mock_generate(*args, **kwargs):
            yield mock_output

        mock_engine.generate = mock_generate

        request = GenerativeScoringRequest(
            model=MODEL_NAME,
            query="Is Paris the capital?",
            items=["Yes", "No"],
            label_token_ids=[1234, 5678],
        )
        result = await serving.create_generative_scoring(request, None)

        assert isinstance(result, GenerativeScoringResponse)
        assert len(result.data) == 2
        for item_result in result.data:
            assert 0.0 <= item_result.score <= 1.0


if __name__ == "__main__":
    pytest.main([__file__, "-v"])
