from importlib.metadata import version

import pytest
from pymongo import AsyncMongoClient, IndexModel

from beanie import Document, Indexed, init_beanie
from beanie import __version__ as beanie_version
from beanie.exceptions import CollectionWasNotInitialized
from beanie.odm.utils.projection import get_projection
from tests.odm.models import (
    Color,
    DocumentTestModel,
    DocumentTestModelIndexFlagsAnnotated,
    DocumentTestModelStringImport,
    DocumentTestModelWithComplexIndex,
    DocumentTestModelWithCustomCollectionName,
    DocumentTestModelWithDroppedIndex,
    DocumentTestModelWithIndexFlags,
    DocumentTestModelWithIndexFlagsAliases,
    DocumentTestModelWithSimpleIndex,
    DocumentToBeLinked,
    DocumentWithCustomInit,
    DocumentWithIndexMerging2,
    DocumentWithLink,
    DocumentWithListLink,
    DocumentWithOptionalTypingOptionalBackLink,
    DocumentWithOptionalTypingOptionalLink,
    DocumentWithUnionTypeExpressionOptionalBackLink,
    DocumentWithUnionTypeExpressionOptionalLink,
)


def test_init_collection_was_not_initialized():
    class NewDocument(Document):
        test_str: str

    with pytest.raises(CollectionWasNotInitialized):
        NewDocument(test_str="test")


async def test_init_connection_string(settings):
    class NewDocumentCS(Document):
        test_str: str

    try:
        await init_beanie(
            connection_string=settings.mongodb_dsn,
            document_models=[NewDocumentCS],
        )
        assert (
            NewDocumentCS.get_pymongo_collection().database.name
            == settings.mongodb_dsn.split("/")[-1]
        )

    finally:
        await NewDocumentCS.get_pymongo_collection().database.client.close()


async def test_init_wrong_params(settings, db):
    class NewDocumentCS(Document):
        test_str: str

    with pytest.raises(ValueError):
        await init_beanie(
            database=db,
            connection_string=settings.mongodb_dsn,
            document_models=[NewDocumentCS],
        )

    with pytest.raises(ValueError):
        await init_beanie(document_models=[NewDocumentCS])

    with pytest.raises(ValueError):
        await init_beanie(connection_string=settings.mongodb_dsn)


async def test_metadata_connection_string(settings):
    class NewDocument(Document):
        test_str: str

    try:
        await init_beanie(
            connection_string=settings.mongodb_dsn,
            document_models=[NewDocument],
        )

        metadata = NewDocument.get_pymongo_collection().database.client.options.pool_options.metadata
        assert "beanie" in metadata["driver"]["name"]
        assert beanie_version in metadata["driver"]["version"]

    finally:
        await NewDocument.get_pymongo_collection().database.client.close()


@pytest.mark.skipif(
    version("pymongo") < "4.14",
    reason="append_metadata was added in PyMongo 4.14",
)
async def test_metadata_database(settings):
    class NewDocument(Document):
        test_str: str

    async with AsyncMongoClient(settings.mongodb_dsn) as client:
        db = client[settings.mongodb_db_name]
        await init_beanie(database=db, document_models=[NewDocument])

        metadata = NewDocument.get_pymongo_collection().database.client.options.pool_options.metadata
        assert "beanie" in metadata["driver"]["name"]
        assert beanie_version in metadata["driver"]["version"]


def test_collection_with_custom_name():
    collection = (
        DocumentTestModelWithCustomCollectionName.get_pymongo_collection()
    )
    assert collection.name == "custom"


async def test_simple_index_creation():
    collection = DocumentTestModelWithSimpleIndex.get_pymongo_collection()
    index_info = await collection.index_information()
    assert index_info["test_int_1"] == {"key": [("test_int", 1)], "v": 2}
    assert index_info["test_str_text"]["key"] == [
        ("_fts", "text"),
        ("_ftsx", 1),
    ]


async def test_flagged_index_creation():
    collection = DocumentTestModelWithIndexFlags.get_pymongo_collection()
    index_info = await collection.index_information()
    assert index_info["test_int_1"] == {
        "key": [("test_int", 1)],
        "sparse": True,
        "v": 2,
    }
    assert index_info["test_str_-1"] == {
        "key": [("test_str", -1)],
        "unique": True,
        "v": 2,
    }


async def test_flagged_index_creation_with_alias():
    collection = (
        DocumentTestModelWithIndexFlagsAliases.get_pymongo_collection()
    )
    index_info = await collection.index_information()
    assert index_info["testInt_1"] == {
        "key": [("testInt", 1)],
        "sparse": True,
        "v": 2,
    }
    assert index_info["testStr_-1"] == {
        "key": [("testStr", -1)],
        "unique": True,
        "v": 2,
    }


async def test_annotated_index_creation():
    collection = DocumentTestModelIndexFlagsAnnotated.get_pymongo_collection()
    index_info = await collection.index_information()
    assert index_info["str_index_text"]["key"] == [
        ("_fts", "text"),
        ("_ftsx", 1),
    ]
    assert index_info["str_index_annotated_1"] == {
        "key": [("str_index_annotated", 1)],
        "v": 2,
    }

    assert index_info["uuid_index_annotated_1"] == {
        "key": [("uuid_index_annotated", 1)],
        "unique": True,
        "v": 2,
    }
    if "uuid_index" in index_info:
        assert index_info["uuid_index"] == {
            "key": [("uuid_index", 1)],
            "unique": True,
            "v": 2,
        }


async def test_complex_index_creation():
    collection = DocumentTestModelWithComplexIndex.get_pymongo_collection()
    index_info = await collection.index_information()
    assert index_info == {
        "_id_": {"key": [("_id", 1)], "v": 2},
        "test_int_1": {"key": [("test_int", 1)], "v": 2},
        "test_int_1_test_str_-1": {
            "key": [("test_int", 1), ("test_str", -1)],
            "v": 2,
        },
        "test_string_index_DESCENDING": {"key": [("test_str", -1)], "v": 2},
    }


async def test_index_dropping_is_allowed(db):
    await init_beanie(
        database=db, document_models=[DocumentTestModelWithComplexIndex]
    )
    collection = DocumentTestModelWithComplexIndex.get_pymongo_collection()

    await init_beanie(
        database=db,
        document_models=[DocumentTestModelWithDroppedIndex],
        allow_index_dropping=True,
    )

    collection = DocumentTestModelWithComplexIndex.get_pymongo_collection()
    index_info = await collection.index_information()
    assert index_info == {
        "_id_": {"key": [("_id", 1)], "v": 2},
        "test_int_1": {"key": [("test_int", 1)], "v": 2},
    }


async def test_index_dropping_is_not_allowed(db):
    await init_beanie(
        database=db, document_models=[DocumentTestModelWithComplexIndex]
    )
    await init_beanie(
        database=db,
        document_models=[DocumentTestModelWithDroppedIndex],
        allow_index_dropping=False,
    )

    collection = DocumentTestModelWithComplexIndex.get_pymongo_collection()
    index_info = await collection.index_information()
    assert index_info == {
        "_id_": {"key": [("_id", 1)], "v": 2},
        "test_int_1": {"key": [("test_int", 1)], "v": 2},
        "test_int_1_test_str_-1": {
            "key": [("test_int", 1), ("test_str", -1)],
            "v": 2,
        },
        "test_string_index_DESCENDING": {"key": [("test_str", -1)], "v": 2},
    }


async def test_index_dropping_is_not_allowed_as_default(db):
    await init_beanie(
        database=db, document_models=[DocumentTestModelWithComplexIndex]
    )
    await init_beanie(
        database=db,
        document_models=[DocumentTestModelWithDroppedIndex],
    )

    collection = DocumentTestModelWithComplexIndex.get_pymongo_collection()
    index_info = await collection.index_information()
    assert index_info == {
        "_id_": {"key": [("_id", 1)], "v": 2},
        "test_int_1": {"key": [("test_int", 1)], "v": 2},
        "test_int_1_test_str_-1": {
            "key": [("test_int", 1), ("test_str", -1)],
            "v": 2,
        },
        "test_string_index_DESCENDING": {"key": [("test_str", -1)], "v": 2},
    }


async def test_document_string_import(db):
    await init_beanie(
        database=db,
        document_models=[
            "tests.odm.models.DocumentTestModelStringImport",
        ],
    )
    document = DocumentTestModelStringImport(test_int=1)
    assert document.id is None
    await document.insert()
    assert document.id is not None

    with pytest.raises(ValueError):
        await init_beanie(
            database=db,
            document_models=[
                "tests",
            ],
        )

    with pytest.raises(AttributeError):
        await init_beanie(
            database=db,
            document_models=[
                "tests.wrong",
            ],
        )


def test_projection():
    projection = get_projection(DocumentTestModel)
    assert projection == {
        "_id": 1,
        "test_int": 1,
        "test_list": 1,
        "test_str": 1,
        "test_doc": 1,
        "revision_id": 1,
    }


async def test_index_recreation(settings):
    class Sample1(Document):
        name: Indexed(str, unique=True)

        class Settings:
            name = "sample"

    class Sample2(Document):
        name: str
        status: str = "active"

        class Settings:
            indexes = [
                IndexModel(
                    "name",
                    unique=True,
                    partialFilterExpression={"is_active": {"$eq": "active"}},
                ),
            ]
            name = "sample"

    async with AsyncMongoClient(settings.mongodb_dsn) as client:
        db = client[settings.mongodb_db_name]
        await db.drop_collection("sample")

        await init_beanie(
            database=db,
            document_models=[Sample1],
        )

        await init_beanie(
            database=db, document_models=[Sample2], allow_index_dropping=True
        )

        await db.drop_collection("sample")


async def test_merge_indexes():
    assert (
        await DocumentWithIndexMerging2.get_pymongo_collection().index_information()
        == {
            "_id_": {"key": [("_id", 1)], "v": 2},
            "s0_1": {"key": [("s0", 1)], "v": 2},
            "s1_1": {"key": [("s1", 1)], "v": 2},
            "s2_-1": {"key": [("s2", -1)], "v": 2},
            "s3_index": {"key": [("s3", -1)], "v": 2},
            "s4_index": {"key": [("s4", 1)], "v": 2},
        }
    )


def test_custom_init():
    assert DocumentWithCustomInit.s == "TEST2"


async def test_index_on_custom_types(settings):
    class Sample1(Document):
        name: Indexed(Color, unique=True)

        class Settings:
            name = "sample"

    async with AsyncMongoClient(settings.mongodb_dsn) as client:
        db = client[settings.mongodb_db_name]
        await db.drop_collection("sample")

        await init_beanie(
            database=db,
            document_models=[Sample1],
        )

        await db.drop_collection("sample")


async def test_init_document_with_union_type_expression_optional_link(db):
    await init_beanie(
        database=db,
        document_models=[
            DocumentToBeLinked,
            DocumentWithUnionTypeExpressionOptionalLink,
        ],
    )

    assert (
        DocumentWithUnionTypeExpressionOptionalLink.get_link_fields().keys()
        == {"link", "link_list"}
    )


async def test_init_document_with_typing_optional_link(db):
    await init_beanie(
        database=db,
        document_models=[
            DocumentToBeLinked,
            DocumentWithOptionalTypingOptionalLink,
        ],
    )

    assert DocumentWithOptionalTypingOptionalLink.get_link_fields().keys() == {
        "link",
        "link_list",
    }


async def test_init_document_with_union_type_expression_optional_back_link(db):
    await init_beanie(
        database=db,
        document_models=[
            DocumentWithUnionTypeExpressionOptionalBackLink,
            DocumentWithListLink,
            DocumentWithLink,
        ],
    )

    assert (
        DocumentWithUnionTypeExpressionOptionalBackLink.get_link_fields().keys()
        == {
            "back_link_list",
            "back_link",
        }
    )


async def test_init_document_with_typing_optional_back_link(db):
    await init_beanie(
        database=db,
        document_models=[
            DocumentWithOptionalTypingOptionalBackLink,
            DocumentWithListLink,
            DocumentWithLink,
        ],
    )

    assert (
        DocumentWithOptionalTypingOptionalBackLink.get_link_fields().keys()
        == {"back_link_list", "back_link"}
    )


async def test_init_document_can_inherit_and_extend_settings(settings):
    class Sample1(Document):
        class Settings:
            name = "sample1"
            bson_encoders = {Color: lambda x: x.value}

    class Sample2(Sample1):
        class Settings(Sample1.Settings):
            name = "sample2"

    async with AsyncMongoClient(settings.mongodb_dsn) as client:
        db = client[settings.mongodb_db_name]
        await init_beanie(
            database=db,
            document_models=[Sample2],
        )

        assert Sample2.get_settings().bson_encoders != {}
        assert Sample2.get_settings().name == "sample2"


async def test_init_beanie_with_skip_indexes(settings):
    class NewDocument(Document):
        test_str: str

        class Settings:
            indexes = ["test_str"]

    async with AsyncMongoClient(settings.mongodb_dsn) as client:
        db = client[settings.mongodb_db_name]
        await init_beanie(
            database=db,
            document_models=[NewDocument],
            skip_indexes=True,
        )

        # To force collection creation
        await NewDocument(test_str="Roman Right").save()

        collection = NewDocument.get_pymongo_collection()
        index_info = await collection.index_information()
        assert (
            len(index_info) == 1
        )  # Only the default _id index should be present
