import pytest
from topk_sdk import error

from . import ProjectContext
from .utils import dataset


def test_get_from_non_existent_collection(ctx: ProjectContext):
    with pytest.raises(error.CollectionNotFoundError):
        ctx.client.collection("missing").get(["doc1"])


def test_get_non_existent_document(ctx: ProjectContext):
    collection = dataset.books.setup(ctx)

    docs = ctx.client.collection(collection.name).get(["missing"])
    assert docs == {}


def test_get_document(ctx: ProjectContext):
    collection = dataset.books.setup(ctx)

    docs = ctx.client.collection(collection.name).get(["lotr"])

    assert docs == {
        "lotr": {
            "_id": "lotr",
            "title": "The Lord of the Rings: The Fellowship of the Ring",
            "published_year": 1954,
            "summary": "A group of unlikely heroes sets out to destroy a powerful, evil ring.",
            "summary_embedding": [9.0] * 16,
            "scalar_i8_embedding": [-100] * 16,
            "sparse_f32_embedding": {9: 1.0, 10: 2.0, 11: 3.0},
            "sparse_u8_embedding": {9: 1, 10: 2, 11: 3},
            "f8_embedding": [9.0] * 16,
            "f16_embedding": [9.0] * 16,
            "tags": ["lord of the rings", "fellowship", "magic", "wizard", "elves"],
            "codes": [
                "ISBN 978-0-547-92821-0",
                "ISBN 0-547-92821-2",
                "OCLC 434394005",
                "LCCN 2004558654",
                "Barcode 0618346252",
            ],
            "user_ratings": ["epic", "legendary", "good"],
        }
    }


def test_get_multiple_documents(ctx: ProjectContext):
    collection = dataset.books.setup(ctx)

    docs = ctx.client.collection(collection.name).get(["lotr", "moby"])

    assert docs == {
        "lotr": {
            "_id": "lotr",
            "title": "The Lord of the Rings: The Fellowship of the Ring",
            "published_year": 1954,
            "summary": "A group of unlikely heroes sets out to destroy a powerful, evil ring.",
            "summary_embedding": [9.0] * 16,
            "scalar_i8_embedding": [-100] * 16,
            "sparse_f32_embedding": {9: 1.0, 10: 2.0, 11: 3.0},
            "sparse_u8_embedding": {9: 1, 10: 2, 11: 3},
            "f8_embedding": [9.0] * 16,
            "f16_embedding": [9.0] * 16,
            "tags": ["lord of the rings", "fellowship", "magic", "wizard", "elves"],
            "codes": [
                "ISBN 978-0-547-92821-0",
                "ISBN 0-547-92821-2",
                "OCLC 434394005",
                "LCCN 2004558654",
                "Barcode 0618346252",
            ],
            "user_ratings": ["epic", "legendary", "good"],
        },
        "moby": {
            "_id": "moby",
            "title": "Moby-Dick",
            "published_year": 1851,
            "summary": "A sailor's obsessive quest to hunt a great white whale leads to tragic consequences.",
            "summary_embedding": [6.0] * 16,
            "sparse_f32_embedding": {6: 1.0, 7: 2.0, 8: 3.0},
            "sparse_u8_embedding": {6: 1, 7: 2, 8: 3},
            "f8_embedding": [6.0] * 16,
            "f16_embedding": [6.0] * 16,
            "nullable_importance": 5.0,
            "tags": ["whale", "obsession", "tragedy", "sailing", "ocean"],
            "codes": [],
            "reprint_years": [],
            "user_ratings": [-5, 2, -1, 1],
        },
    }


def test_get_document_fields(ctx: ProjectContext):
    collection = dataset.books.setup(ctx)

    docs = ctx.client.collection(collection.name).get(
        ["lotr"],
        fields=["title", "published_year"],
    )

    assert docs == {
        "lotr": {
            "_id": "lotr",
            "title": "The Lord of the Rings: The Fellowship of the Ring",
            "published_year": 1954,
        }
    }
