from base64 import b64encode
from hashlib import md5

from botocore.awsrequest import AWSRequest

from s3_manager.s3storage.config import OPERATIONS_NEEDING_MD5_HEADER


def add_md5_header(request: AWSRequest, **kwargs: dict) -> None:
    """
    Добавляет Content-MD5 заголовок для операций, изменяющих данные в S3.
    """
    if (
        kwargs.get("operation_name") not in OPERATIONS_NEEDING_MD5_HEADER
        or not request.body
    ):
        return

    # Получаем содержимое тела
    body = request.body
    if hasattr(body, "read"):
        pos = body.tell()
        body.seek(0)
        content = body.read()
        body.seek(pos)
    elif isinstance(body, bytes):
        content = body
    elif isinstance(body, str):
        content = body.encode("utf-8")
    else:
        try:
            content = str(body).encode("utf-8")
        except Exception:
            return

    # Вычисляем и добавляем MD5
    request.headers["Content-MD5"] = b64encode(md5(content).digest()).decode("utf-8")
