from pymongo.collection import Collection
import logging

from .classes import (
    VisualCommunication,
    NoDocumentFoundException,
    ModelOutputs
)


def total_documents(
    collection: Collection
) -> int:
    """Get total number of documents in database."""
    return collection.count_documents(filter={})


def total_annotated(
    collection: Collection
) -> int:
    """Get total number of annotated documents in database."""
    query = {
        "annotation": { "$ne": None }
    }
    return collection.count_documents(filter=query)


def get_visual_communication(
        collection: Collection,
        with_annotation: bool = False
    ) -> VisualCommunication:
    """Get a random visual communication from the database."""
    query = {}
    if with_annotation:
        query["annotation"] = {"$ne": None}
    else:
        query["annotation"] = None
    data = collection.aggregate([
        { "$match": query },   # find using filters
        { "$sample": { "size": 1 } }    # get one random
    ])
    data = list(data)   # read data from cursor object
    if len(data) == 0:
        logging.error("failed getting visual communication")
        raise NoDocumentFoundException()
    data = data[0]
    logging.info("finished")
    return VisualCommunication.model_validate(data)


def upsert_predictions(
    collection: Collection,
    vis_com_name: str,
    predictions: ModelOutputs,
) -> None:
    """Upsert prediction data in the database."""
    query = {
        "name": vis_com_name
    }
    update = {
        "$set": { "prediction": predictions.model_dump() }
    }
    res = collection.update_one(
        filter=query,
        update=update,
        upsert=True,
    )
    logging.debug("upserted document: %s", res)
    logging.info("finished")


def upsert_annotations(
    collection: Collection,
    vis_com_name: str,
    annotations: ModelOutputs,
) -> None:
    """Upserts annotation data in the database."""
    query = {
        "name": vis_com_name
    }
    update = {
        "$set": { "annotation": annotations.model_dump() }
    }
    res = collection.update_one(
        filter=query,
        update=update,
        upsert=True,
    )
    logging.info("upserted document: %s", res)
    logging.info("finished")