"""Script to move all images from MongoDB to MinIO."""

from __future__ import annotations

import logging
from io import BytesIO
from pathlib import Path

from bson import ObjectId
from dotenv import load_dotenv
from pymongo.collection import Collection

from shared.data_store import connect_minio, put
from shared.mongodb import connect_mongodb
from shared.mongodb.classes import VisualCommunication
from shared.utils import check_env, setup_logging
from web_ui.src.main import NECESSARY_ENV_VAR_LIST


def list_mongo_document_ids(
    collection: Collection,
) -> list[ObjectId]:
    """Get list of all VisualCommunication documents in MongoDB."""
    # prepare query
    query = collection.find(
        filter={},
        projection={'_id': True},
    )
    # execute query
    res_list = list(query)
    logging.debug('got %s document(s)', len(res_list))
    # convert result
    id_list = [elem['_id'] for elem in res_list]
    return id_list


def get_visual_communication(
    collection: Collection,
    doc_id: ObjectId,
) -> VisualCommunication:
    """Get image from MongoDB."""
    # prepare query
    query = collection.find(
        filter={'_id': doc_id},
        projection={'_id': False},
    )
    # execute query
    res_list = list(query)
    logging.debug('got %s document(s)', len(res_list))
    # convert result
    vis_com = VisualCommunication.model_validate(res_list[0])
    return vis_com


def update_visual_communication(
    collection: Collection,
    vis_com: VisualCommunication,
) -> None:
    """Update VisualCommunication document in MongoDB."""
    try:
        collection.update_one(
            filter={
                'name': vis_com.name,
            },
            update={
                '$set': vis_com.model_dump(),
            },
        )
    except Exception as exc:
        logging.error('failed saving %s', vis_com.name)
        logging.debug(exc)
        raise exc


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(NECESSARY_ENV_VAR_LIST)
    # setup logging
    setup_logging()
    # connect to minIO
    minio_client = connect_minio()
    # connect to MongoDB
    collection, db, client = connect_mongodb()
    # list documents in mongoDB
    id_list = list_mongo_document_ids(collection)
    for doc_id in id_list:
        try:
            # get image from MongoDB
            vis_com = get_visual_communication(collection, doc_id)
        except Exception as exc:
            logging.debug(exc)
            logging.error('failed getting image from document: %s', doc_id)
            continue
        try:
            # save image to buffer
            buffer = BytesIO()
            vis_com.image.save(buffer, 'png')  # type: ignore
            # put buffer in minio
            object_name = put(
                client=minio_client,
                buffer=buffer,
            )
        except Exception as exc:
            logging.debug(exc)
            logging.error('failed saving image to minio')
            continue
        # update visual communication in mongodb
        try:
            vis_com.image = None  # type: ignore
            vis_com.object_name = object_name
            update_visual_communication(collection, vis_com)
        except Exception as exc:
            logging.debug(exc)
            logging.error(
                'failed updating visual communication %s',
                vis_com.name,
            )
            continue
        logging.debug('updated visual communication %s', vis_com.name)