"""
NoSQL storage backend CRUD operations for ActivityPub data.
"""

from collections.abc import Awaitable, Callable
from logging import Logger
from typing import Any, cast, override

from pydantic import Field, field_serializer
from pymongo.asynchronous.collection import AsyncCollection
from pymongo.asynchronous.cursor import AsyncCursor

from phederation.models import CRUD, T, object_from_type
from phederation.utils import ObjectId
from phederation.utils.base import AccessType, UrlType, assemble_id_url
from phederation.utils.exceptions import StorageError, catch_exceptions


def nosql_build_search_query(
    collection_id: ObjectId | None, ids: list[ObjectId] | None = None, access: AccessType = AccessType.PUBLIC
) -> dict[str, list[dict[str, ObjectId] | dict[str, list[dict[str, str]]]]]:
    query: dict[str, list[dict[str, ObjectId | str] | dict[str, list[dict[str, str]]]]] | None = {"$and": []}
    if collection_id is not None:
        query["$and"].append({"collectionId": collection_id})
    if ids is not None:
        query["$and"].append({"$or": [{"storedId": id} for id in ids]})

    if access == AccessType.PUBLIC:
        query["$and"].append({"visibility": AccessType.PUBLIC.value})
    elif access == AccessType.FOLLOWER:
        query["$and"].append({"$or": [{"visibility": AccessType.PUBLIC.value}, {"visibility": AccessType.FOLLOWER.value}]})
    elif access == AccessType.PRIVATE:
        pass  # use all items
    if len(query["$and"]) == 0:
        query = {}
    return query


class NoSqlCRUD(CRUD[T]):
    get_collection: Callable[[str], Awaitable[AsyncCollection[Any]]]
    table_name: str
    base_url: ObjectId
    logger: Logger = Field(default=..., exclude=True)
    MODEL: type

    @field_serializer("logger")
    def serialize_logger(self, logger: Logger):
        logger = logger
        return "None"

    @override
    async def get_table_name(self) -> str:
        return self.table_name

    @catch_exceptions(StorageError, "NoSQL: Failed to create_many")
    @override
    async def create_many(self, data: list[T]) -> list[ObjectId]:
        data_dict = [d.serialize(include_context=False) for d in data]
        object_ids: list[ObjectId | str | None] = []
        for d in data_dict:
            if "id" in d:
                object_id = d.get("id")
                d["_id"] = object_id  # necessary for mongodb
            else:
                object_id = None

            object_ids.append(object_id)
        collection = await self.get_collection(self.table_name)
        new_items = await collection.insert_many(data_dict)

        # update the id string according to the new primary key
        new_object_ids: list[ObjectId | str] = []
        for inserted_id, object_id in zip(cast(list[ObjectId], new_items.inserted_ids), object_ids):
            if object_id is None:
                object_id = assemble_id_url(type=UrlType(self.table_name), primary=(inserted_id), base_url=self.base_url)
                query_filter = {"_id": (inserted_id)}
                update_operation = {"$set": {"id": object_id}}
                _ = await collection.update_one(query_filter, update_operation)
            new_object_ids.append(object_id)
        return [object_id for object_id in object_ids if object_id]

    @catch_exceptions(StorageError, "NoSQL: Failed to create")
    @override
    async def create(self, data: T, raise_if_exists: bool = True) -> ObjectId:
        data_dict = data.serialize(include_context=False)
        if "id" in data_dict:
            object_id = data_dict.get("id")
            data_dict["_id"] = object_id  # necessary for mongodb
        else:
            object_id = None

        collection = await self.get_collection(self.table_name)

        # if there is already an item in the collection with the given id, change the id and insert
        try:
            new_item = await collection.insert_one(data_dict)
        except Exception as e:
            # this can happen if an instance receives a message for multiple actors.
            # TODO: this is not a very efficient way to resolve multiple messages
            self.logger.warning(f"Data '{data}' with id '{object_id}' already exists in storage")
            if raise_if_exists or not object_id:
                raise e
            else:
                return cast(str, object_id)

        # update the id string according to the new primary key
        if object_id is None:
            new_inserted_id = cast(ObjectId, new_item.inserted_id)
            object_id = assemble_id_url(type=UrlType(self.table_name), primary=new_inserted_id, base_url=self.base_url)
            query_filter = {"_id": new_inserted_id}
            update_operation = {"$set": {"id": object_id}}
            _ = await collection.update_one(query_filter, update_operation)
        return object_id

    @catch_exceptions(StorageError, "NoSQL: Failed to read")
    @override
    async def read(self, id: ObjectId) -> T | None:
        collection = await self.get_collection(self.table_name)
        result = cast(dict[str, Any], await collection.find_one({"id": id}))

        if result:
            result.pop("_id", None)
            obj = cast(T, object_from_type(self.MODEL.deserialize(result)))  # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType]
            return obj
        else:
            return None

    @catch_exceptions(StorageError, "NoSQL: Failed to select")
    @override
    async def select(
        self, collection_id: ObjectId | None = None, stored_ids: list[ObjectId] | None = None, access: AccessType = AccessType.PUBLIC
    ) -> list[T]:
        collection = await self.get_collection(self.table_name)
        results: list[T] = []
        query = nosql_build_search_query(collection_id=collection_id, ids=stored_ids, access=access)
        async with collection.find(query) as cursor:
            async for result in cast(AsyncCursor[dict[str, Any]], cursor):
                if result:
                    result.pop("_id", None)
                    obj = cast(
                        T, object_from_type(self.MODEL.deserialize(result))  # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType]
                    )
                    results.append(obj)
        return results

    @catch_exceptions(StorageError, "Failed to count")
    @override
    async def count(self, collection_id: ObjectId | None = None, access: AccessType = AccessType.PUBLIC) -> int:
        collection = await self.get_collection(self.table_name)
        query = nosql_build_search_query(collection_id=collection_id, access=access)
        total_items = await collection.count_documents(query)
        return total_items

    @catch_exceptions(StorageError, "NoSQL: Failed to update")
    @override
    async def upsert(self, id: ObjectId, data: T) -> ObjectId:
        """Upsert an object with given id.
        If the id cannot be found, inserts the new object with the .id field set to the parameter id.

        Args:
            id (ObjectId): Id of the object to update.
            data (T): Object information.

        Returns:
            ObjectId: The input id.
        """
        collection = await self.get_collection(self.table_name)
        query_filter = {"id": id}
        data.id = id
        data_dict = data.serialize(include_context=False)
        update_operation = {"$set": data_dict}
        _ = await collection.update_one(query_filter, update_operation, upsert=True)
        return id

    @catch_exceptions(StorageError, "NoSQL: Failed to delete")
    @override
    async def delete(self, id: ObjectId | list[ObjectId]) -> bool:
        collection = await self.get_collection(self.table_name)
        if isinstance(id, str):
            _ = await collection.delete_one({"id": id})
        else:
            if len(id) > 0:
                filter = {"$or": [{"id": i} for i in id]}
                _ = await collection.delete_many(filter)
        return True
