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

import pytest
import requests

from tests.utils import VLLM_PATH, RemoteOpenAIServer
from vllm.entrypoints.pooling.score.protocol import RerankResponse, ScoreResponse
from vllm.multimodal.utils import encode_image_url, fetch_image

MODEL_NAME = "Qwen/Qwen3-VL-Reranker-2B"
HF_OVERRIDES = {
    "architectures": ["Qwen3VLForSequenceClassification"],
    "classifier_from_token": ["no", "yes"],
    "is_original_qwen3_reranker": True,
}

query = "A cat standing in the snow."
document = "This product was excellent and exceeded my expectations."
image_url = "https://vllm-public-assets.s3.us-west-2.amazonaws.com/multimodal_asset/cat_snow.jpg"
documents = [
    {
        "type": "text",
        "text": document,
    },
    {
        "type": "image_url",
        "image_url": {"url": image_url},
    },
    {
        "type": "image_url",
        "image_url": {"url": encode_image_url(fetch_image(image_url))},
    },
]

TEXT_VS_TEXT = 0.10040374100208282
TEXT_VS_IMAGE = 0.7423753142356873
TEXT_VS_TEXT_PLUS_IMAGE = 0.5298863053321838
TOL = 0.05


@pytest.fixture(scope="module")
def server():
    args = [
        "--enforce-eager",
        "--max-model-len",
        "8192",
        "--chat-template",
        str(VLLM_PATH / "examples/pooling/score/template/qwen3_vl_reranker.jinja"),
    ]

    with RemoteOpenAIServer(
        MODEL_NAME, args, override_hf_configs=HF_OVERRIDES
    ) as remote_server:
        yield remote_server


def test_score_api_queries_str_documents_str(server: RemoteOpenAIServer):
    score_response = requests.post(
        server.url_for("score"),
        json={
            "model": MODEL_NAME,
            "queries": query,
            "documents": document,
        },
    )
    score_response.raise_for_status()
    score = ScoreResponse.model_validate(score_response.json())

    assert score.id is not None
    assert score.data is not None
    assert len(score.data) == 1
    assert score.usage.prompt_tokens == 81
    assert score.data[0].score == pytest.approx(TEXT_VS_TEXT, rel=TOL)


def test_score_api_queries_str_documents_text_content(server: RemoteOpenAIServer):
    score_response = requests.post(
        server.url_for("score"),
        json={
            "model": MODEL_NAME,
            "queries": query,
            "documents": {"content": [documents[0]]},
        },
    )
    score_response.raise_for_status()
    score = ScoreResponse.model_validate(score_response.json())

    assert score.id is not None
    assert score.data is not None
    assert len(score.data) == 1
    assert score.usage.prompt_tokens == 81
    assert score.data[0].score == pytest.approx(TEXT_VS_TEXT, rel=TOL)


def test_score_api_queries_str_documents_image_url_content(server: RemoteOpenAIServer):
    score_response = requests.post(
        server.url_for("score"),
        json={
            "model": MODEL_NAME,
            "queries": query,
            "documents": {"content": [documents[1]]},
        },
    )
    score_response.raise_for_status()
    score = ScoreResponse.model_validate(score_response.json())

    assert score.id is not None
    assert score.data is not None
    assert len(score.data) == 1
    assert score.usage.prompt_tokens == 98
    assert score.data[0].score == pytest.approx(TEXT_VS_IMAGE, rel=TOL)


def test_score_api_queries_str_documents_image_base64_content(
    server: RemoteOpenAIServer,
):
    score_response = requests.post(
        server.url_for("score"),
        json={
            "model": MODEL_NAME,
            "queries": query,
            "documents": {"content": [documents[2]]},
        },
    )
    score_response.raise_for_status()
    score = ScoreResponse.model_validate(score_response.json())

    assert score.id is not None
    assert score.data is not None
    assert len(score.data) == 1
    assert score.usage.prompt_tokens == 98
    assert score.data[0].score == pytest.approx(TEXT_VS_IMAGE, rel=TOL)


def test_score_api_queries_str_documents_image_url_plus_text_content(
    server: RemoteOpenAIServer,
):
    score_response = requests.post(
        server.url_for("score"),
        json={
            "model": MODEL_NAME,
            "queries": query,
            "documents": {"content": [documents[0], documents[1]]},
        },
    )
    score_response.raise_for_status()
    score = ScoreResponse.model_validate(score_response.json())

    assert score.id is not None
    assert score.data is not None
    assert len(score.data) == 1
    assert score.usage.prompt_tokens == 108
    assert score.data[0].score == pytest.approx(TEXT_VS_TEXT_PLUS_IMAGE, rel=TOL)


def test_score_api_queries_str_documents_list(server: RemoteOpenAIServer):
    score_response = requests.post(
        server.url_for("score"),
        json={
            "model": MODEL_NAME,
            "queries": query,
            "documents": [
                document,
                {"content": [documents[0]]},
                {"content": [documents[1]]},
                {"content": [documents[0], documents[1]]},
            ],
        },
    )
    score_response.raise_for_status()
    score = ScoreResponse.model_validate(score_response.json())

    assert score.id is not None
    assert score.data is not None
    assert len(score.data) == 4
    assert score.usage.prompt_tokens == 368
    assert score.data[0].score == pytest.approx(TEXT_VS_TEXT, rel=TOL)
    assert score.data[1].score == pytest.approx(TEXT_VS_TEXT, rel=TOL)
    assert score.data[2].score == pytest.approx(TEXT_VS_IMAGE, rel=TOL)
    assert score.data[3].score == pytest.approx(TEXT_VS_TEXT_PLUS_IMAGE, rel=TOL)


def test_rerank_api_queries_str_documents_list(server: RemoteOpenAIServer):
    rerank_response = requests.post(
        server.url_for("rerank"),
        json={
            "model": MODEL_NAME,
            "query": query,
            "documents": [
                document,
                {"content": [documents[0]]},
                {"content": [documents[1]]},
                {"content": [documents[0], documents[1]]},
            ],
        },
    )
    rerank_response.raise_for_status()
    rerank = RerankResponse.model_validate(rerank_response.json())

    assert rerank.id is not None
    assert rerank.model is not None
    assert rerank.usage is not None
    assert len(rerank.results) == 4

    rerank.results.sort(key=lambda x: x.index)
    assert rerank.results[0].relevance_score == pytest.approx(TEXT_VS_TEXT, rel=TOL)
    assert rerank.results[1].relevance_score == pytest.approx(TEXT_VS_TEXT, rel=TOL)
    assert rerank.results[2].relevance_score == pytest.approx(TEXT_VS_IMAGE, rel=TOL)
    assert rerank.results[3].relevance_score == pytest.approx(
        TEXT_VS_TEXT_PLUS_IMAGE, rel=TOL
    )


def test_score_api_queries_list_documents_list(server: RemoteOpenAIServer):
    score_response = requests.post(
        server.url_for("score"),
        json={
            "model": MODEL_NAME,
            "queries": [query] * 4,
            "documents": [
                document,
                {"content": [documents[0]]},
                {"content": [documents[1]]},
                {"content": [documents[0], documents[1]]},
            ],
        },
    )
    score_response.raise_for_status()
    score = ScoreResponse.model_validate(score_response.json())

    assert score.id is not None
    assert score.data is not None
    assert len(score.data) == 4
    assert score.usage.prompt_tokens == 368
    assert score.data[0].score == pytest.approx(TEXT_VS_TEXT, rel=TOL)
    assert score.data[1].score == pytest.approx(TEXT_VS_TEXT, rel=TOL)
    assert score.data[2].score == pytest.approx(TEXT_VS_IMAGE, rel=TOL)
    assert score.data[3].score == pytest.approx(TEXT_VS_TEXT_PLUS_IMAGE, rel=TOL)
