from __future__ import annotations

import logging

from pymongo.collection import Collection

from shared.dto import ModelData


def upsert_prediction(
    collection: Collection,
    vis_com_name: str,
    predictions: ModelData,
) -> 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')