import json

from pydantic import BaseModel
from pymongo import MongoClient
from pymongo.collection import Collection
from pymongo.database import Database
from pymongo.results import DeleteResult, InsertOneResult, UpdateResult

from ...sdk.common.config.environment import demand_env_var, demand_env_var_as_int
from ...sdk.common.utils import address_type, get_document_id_from_address
from ...sdk.contracts.dtos.entities.entity import Entity
from ...sdk.contracts.dtos.tiles.address import Address
from ...sdk.contracts.dtos.tiles.chunk import Chunk
from ...sdk.contracts.dtos.tiles.tile import Tile
from ...sdk.contracts.dtos.tiles.world import World
from ...sdk.contracts.errors.server.dao.doesnotexist import DaoDoesNotExistError
from ...sdk.contracts.errors.server.dao.unknown import DaoUnknownError
from ...sdk.contracts.types.dao_document import DaoDocumentType


def get_mongodb() -> Database:
    hostname: str = demand_env_var(name="DARKNESS_MONGODB_HOST")
    port: int = demand_env_var_as_int(name="DARKNESS_MONGODB_PORT")
    database: str = demand_env_var(name="DARKNESS_MONGODB_DATABASE")
    return MongoClient(hostname, port)[database]


class DaoClient(BaseModel):
    database: Database
    collections: dict[DaoDocumentType, Collection] = {}

    def __init__(self, **kwargs) -> None:
        super().__init__(**kwargs)

        # our collections will match our document types
        for member in DaoDocumentType:
            # Create a collection for each document type
            self.collections[DaoDocumentType[member.name]] = self.database[member.value]

            # We add a single column index on the business logic layer `id` for our lookups:
            # db.<collection>.createIndex( { <field>: <sortOrder> } )
            self.collections[DaoDocumentType[member.name]].create_index({"id": 1})

    class Config:
        arbitrary_types_allowed = True

    # Single document delete
    async def delete(self, address: Address) -> DeleteResult:
        doc_type = address_type(address)
        document_id: str = get_document_id_from_address(address=address, doc_type=doc_type)
        return self.collections[doc_type].delete_one({"id": document_id})

    # single document get with typing!
    async def get(self, address: Address) -> World | Chunk | Tile | Entity:
        doc_type: DaoDocumentType = address_type(address=address)
        document_id: str = get_document_id_from_address(address=address, doc_type=doc_type)
        result: dict | None = self.collections[doc_type].find_one({"id": document_id})

        if result is None:
            raise DaoDoesNotExistError("no document found")

        if doc_type == DaoDocumentType.ENTITY:
            return Entity.model_validate(result)

        if doc_type == DaoDocumentType.TILE:
            return Tile.model_validate(result)

        if doc_type == DaoDocumentType.CHUNK:
            return Chunk.model_validate(result)

        if doc_type == DaoDocumentType.WORLD:
            return World.model_validate(result)

        raise DaoUnknownError("invalid document type")

    async def get_all(self, doc_type: DaoDocumentType) -> list[World] | list[Chunk] | list[Tile] | list[Entity]:
        results = self.collections[doc_type].find()
        documents: list[World] | list[Chunk] | list[Tile] | list[Entity] = []
        for result in results:
            if doc_type == DaoDocumentType.ENTITY:
                documents.append(Entity.model_validate(result))
            elif doc_type == DaoDocumentType.TILE:
                documents.append(Tile.model_validate(result))
            elif doc_type == DaoDocumentType.CHUNK:
                documents.append(Chunk.model_validate(result))
            elif doc_type == DaoDocumentType.WORLD:
                documents.append(World.model_validate(result))
            else:
                raise DaoUnknownError("invalid document type")
        return documents

    async def get_multi(
        self, addresses: list[Address], doc_type: DaoDocumentType
    ) -> list[World] | list[Chunk] | list[Tile] | list[Entity]:
        doc_ids: list = []
        for address in addresses:
            document_id: str = get_document_id_from_address(address=address, doc_type=doc_type)
            doc_ids.append(document_id)
        results = self.collections[doc_type].find({"id": {"$in": doc_ids}})
        documents: list[World] | list[Chunk] | list[Tile] | list[Entity] = []
        for result in results:
            if doc_type == DaoDocumentType.ENTITY:
                documents.append(Entity.model_validate(result))
            elif doc_type == DaoDocumentType.TILE:
                documents.append(Tile.model_validate(result))
            elif doc_type == DaoDocumentType.CHUNK:
                documents.append(Chunk.model_validate(result))
            elif doc_type == DaoDocumentType.WORLD:
                documents.append(World.model_validate(result))
            else:
                raise DaoUnknownError("invalid document type")
        return documents

    async def post(self, address: Address, document: World | Chunk | Tile | Entity) -> InsertOneResult:
        doc_type: DaoDocumentType = address_type(address=address)
        return self.collections[doc_type].insert_one(json.loads(document.model_dump_json()))

    async def patch(self, address: Address, document: dict) -> UpdateResult:
        # document CAN NOT CONTAIN set variables.  It MUST be serialized before this call, so help you
        doc_type: DaoDocumentType = address_type(address=address)
        document_id: str = get_document_id_from_address(address=address, doc_type=doc_type)
        return self.collections[doc_type].update_one(filter={"id": document_id}, update={"$set": document})
