"""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


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)
    bucket_name = os.getenv('MINIO_BUCKET_NAME', default='')
    assert len(bucket_name) > 0
    subfolder = 'models'
    # save image to buffer
    buffer = BytesIO()
    torch.save(model.state_dict(), buffer)
    # get md5 of image
    checksum = md5(buffer.getbuffer()).hexdigest()
    # prepare for saving
    num_bytes = buffer.tell()
    buffer.seek(0)
    # send data to bucket
    object_name = f'{subfolder}/{checksum}'
    try:
        client.put_object(
            bucket_name=bucket_name,
            object_name=object_name,
            length=num_bytes,
            data=buffer,
        )
    except Exception as exc:
        logging.error('failed saving data to MinIO')
        raise exc
    logging.debug('saved data to %s', object_name)
    return checksum