import pytest

from .helpers.collection_setup import drop_collection
from .helpers.helpers import request_with_validation


@pytest.fixture(autouse=True)
def setup(on_disk_vectors, collection_name):
    multivector_collection_setup(collection_name=collection_name, on_disk_vectors=on_disk_vectors)
    yield
    drop_collection(collection_name=collection_name)

def multivector_collection_setup(
        collection_name='test_collection',
        on_disk_vectors=False):
    drop_collection(collection_name=collection_name)

    response = request_with_validation(
        api='/collections/{collection_name}',
        method="PUT",
        path_params={'collection_name': collection_name},
        body={
            "vectors": {
                "size": 2,
                "distance": "Dot",
                "on_disk": on_disk_vectors,
                "multivector_config": {
                    "comparator": "max_sim"
                }
            },
        }
    )
    assert response.ok

    # batch upsert
    response = request_with_validation(
        api='/collections/{collection_name}/points',
        method="PUT",
        path_params={'collection_name': collection_name},
        query_params={'wait': 'true'},
        body={
            "points": [
                {
                    "id": 1,
                    "vector": [
                        [-1, 1],
                        [-2, 2],
                        [-3, 3]
                    ]

                },
                {
                    "id": 2,
                    "vector": [
                        [-1.5, 1.5],
                        [-2.5, 2.5],
                        [-3.5, 3.5]
                    ]
                },
                {
                    "id": 3,
                    "vector": [
                        [-1.7, 1.7],
                        [-2.7, 2.7],
                        [-3.7, 3.7]
                    ]
                },
            ]
        }
    )
    assert response.ok

    response = request_with_validation(
        api='/collections/{collection_name}',
        method="GET",
        path_params={'collection_name': collection_name},
    )
    assert response.ok

def test_multi_default_is_avg_vector(collection_name):
    examples = {
        "positive": [[1, 2]],
        "negative": [[3, 4]],
    }

    default_response = request_with_validation(
        api="/collections/{collection_name}/points/query",
        method="POST",
        path_params={"collection_name": collection_name},
        body={
            "query": {"recommend": examples},
            "params": {"exact": True},
            "limit": 10,
        },
    )
    assert default_response.ok

    # we should only get 3 because we don't have more
    assert len(default_response.json()["result"]["points"]) == 3

    avg_response = request_with_validation(
        api="/collections/{collection_name}/points/query",
        method="POST",
        path_params={"collection_name": collection_name},
        body={
            "query": {"recommend": {**examples, "strategy": "average_vector"}},
            "params": {"exact": True},
            "limit": 10,
        },
    )
    assert avg_response.ok
    assert len(avg_response.json()["result"]["points"]) == 3

    assert default_response.json()["result"] == avg_response.json()["result"]


def test_multi_single_vs_batch(collection_name):
    # Bunch of valid examples
    params_list = [
        {
            "query": {"recommend": {"positive": [1], "negative": [3]}},
            "limit": 1,
        },
        {
            # no positive because it's optional with this strategy
            "query": {"recommend": {"negative": [1, 2], "strategy": "best_score"}},
            "params": {"exact": True},
            "limit": 1,
        },
        {
            "query": {"recommend": {"positive": [1], "negative": [], "strategy": "average_vector"}},
            "params": {"exact": True},
            "limit": 1,
        },
    ]

    batch_response = request_with_validation(
        api="/collections/{collection_name}/points/query/batch",
        method="POST",
        path_params={"collection_name": collection_name},
        body={"searches": params_list},
    )

    assert batch_response.ok
    assert len(batch_response.json()["result"]) == len(params_list)

    # Compare against sequential single searches
    for i, params in enumerate(params_list):
        single_response = request_with_validation(
            api="/collections/{collection_name}/points/query",
            method="POST",
            path_params={"collection_name": collection_name},
            body=params,
        )
        assert single_response.ok
        assert single_response.json()["result"] == batch_response.json()["result"][i]


def test_multi_without_positives(collection_name):
    def req_with_positives(positive, strategy=None):
        recommend = {"positive": positive}
        if strategy is not None:
            recommend["strategy"] = strategy

        return request_with_validation(
            api="/collections/{collection_name}/points/query",
            method="POST",
            path_params={"collection_name": collection_name},
            body={
                "query": {"recommend": recommend},
                "limit": 2,
            },
        )

    # Assert this is valid
    response = req_with_positives([[1, 2]])
    assert response.ok

    # But all these are not. An empty vector is a malformed example, rejected
    # once the query runs...
    response = req_with_positives([[]])
    assert response.status_code == 400

    response = req_with_positives([[]], "average_vector")
    assert response.status_code == 400

    # ...whereas no examples at all violates a `RecommendInput` validation rule
    # and is rejected before that, hence 422.
    response = req_with_positives([], "best_score")
    assert response.status_code == 422


def test_multi_best_score_works_with_only_negatives(collection_name):
    response = request_with_validation(
        api="/collections/{collection_name}/points/query",
        method="POST",
        path_params={"collection_name": collection_name},
        body={
            "query": {"recommend": {"negative": [[1, 2]], "strategy": "best_score"}},
            "limit": 5,
        },
    )
    assert response.ok
    assert len(response.json()["result"]["points"]) == 3

    # All scores should be negative
    for result in response.json()["result"]["points"]:
        assert result["score"] < 0