"""Definition of put_model function."""

import logging
import os
from hashlib import md5
from io import BytesIO

import torch
from minio import Minio
from torch.nn import Module

from shared.utils import check_env

from .put import put


def put_model(
    client: Minio,
    model: Module,
) -> str:
    """Put model in model subfolder in bucket in Minio and return MD5 checksum
    used as object name."""
    assert isinstance(client, Minio)
    assert isinstance(model, Module)
    # get bucket name from env
    check_env({'MINIO_BUCKET_NAME'})
    bucket_name = str(os.getenv('MINIO_BUCKET_NAME'))
    # save data to buffer
    buffer = BytesIO()
    torch.save(model.state_dict(), buffer)
    # get md5 of image
    checksum = md5(buffer.getbuffer()).hexdigest()
    # set object name
    object_name = f'models/{checksum}'
    # send data to bucket
    put(
        client=client,
        buffer=buffer,
        bucket_name=bucket_name,
        object_name=object_name,
    )
    logging.debug('finished')
    return checksum