"""Main script to be run by service."""
from __future__ import annotations

from hashlib import md5
from io import BytesIO

import torch
from data_store import connect
from data_store import put
from models import VisualCommunicationModel
from torchinfo import summary
from utils import check_env

from shared.utils import setup_logging


DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

if __name__ == '__main__':
    # check that environment variables are set
    check_env()
    # setup logging
    setup_logging()
    # instantiate model
    model = VisualCommunicationModel().to(DEVICE)
    # show model weights
    summary(model)
    # for param in model.parameters():
    #     logging.info(param.data)
    # save model to buffer
    buffer = BytesIO()
    torch.save(model.state_dict(), buffer)
    print(f"buffer size: {len(buffer.getvalue())}")
    # calculate buffer hash
    hash_str = md5(buffer.getbuffer()).hexdigest()
    print(f"hash string: {hash_str}")
    # connect to minio
    client = connect()
    # put buffer in minio bucket
    put(
        client=client,
        buffer=buffer,
        object_name=hash_str,
    )
    print('saved to Minio')