from __future__ import annotations

from pathlib import Path

from dotenv import load_dotenv

from shared.data_store import connect_minio
from shared.database import connect as connect_mongo
from shared.database import upsert_prediction
from shared.database.classes import VisualCommunication
from shared.utils import check_env, setup_logging

if __name__ == '__main__':
    # load in env file
    env_path = Path(__file__).parent.parent / 'server.env'
    assert env_path.exists()
    load_dotenv(env_path)
    # ensure env vars set
    check_env()
    # setup logging
    setup_logging()
    # connect to minIO
    minio_client = connect_minio()
    # connect to MongoDB
    collection, db, client = connect_mongo()
    # get list of image paths
    test_dir = Path(__file__).parent
    img_dir = test_dir / 'imgs'
    img_path_list = [path for path in img_dir.glob('*.jpeg') if path.is_file()]
    # instantiate data object
    vis_com_list = [
        VisualCommunication.from_file(path, minio_client=minio_client)
        for path in img_path_list
    ]
    # generate random predictions
    for vis_com in vis_com_list:
        vis_com.generate_random_prediction()
    # upload visual communication
    for vis_com in vis_com_list:
        if vis_com.prediction is None:
            continue
        upsert_prediction(
            collection=collection,
            vis_com_name=vis_com.name,
            predictions=vis_com.prediction,
        )