from __future__ import annotations

import logging

from pymongo.collection import Collection

from shared.database import Dataset


def save_dataset(
    collection: Collection,
    dataset: Dataset,
) -> None:
    """Save dataset to database."""
    res = collection.insert_one(
        document=dataset.model_dump(),
    )
    logging.debug('inserted document: %s', res)


if __name__ == '__main__':
    from dotenv import load_dotenv

    load_dotenv('local.env')
    from shared.database import connect_mongodb, list_names

    # connect to database
    collection, db, client = connect_mongodb()
    print(client.server_info())

    name_list = list_names(collection=collection, only_with_annotation=True)
    ds = Dataset.new_from_name_list(name_list=name_list)

    print(ds)
    # save_dataset(
    #     collection=collection,
    #     dataset=ds
    # )