from base64 import b64encode
from datetime import datetime, timedelta
from hashlib import sha256

from botocore.client import BaseClient, logger
from pydantic import TypeAdapter

from s3_manager.exception.main import ExceptionHandler
from s3_manager.s3storage.schema import PartMetadataSchema
from s3_manager.s3storage.session import s3_client_wrapper


@s3_client_wrapper
def create(s3_client: BaseClient, s3bucket: str, obj_key: str) -> dict:
    """
    Создает объект с экспайром
    """
    logger.info("Создаю объект для привязки к чанкам.")
    return s3_client.create_multipart_upload(
        Bucket=s3bucket,
        Key=obj_key,
        Expires=datetime.now() + timedelta(minutes=30),
    )


@s3_client_wrapper
def upload_part(
    s3_client: BaseClient,
    s3bucket: str,
    obj_key: str,
    chunk: bytes,
    upload_id: str,
    chunk_num: int,
) -> PartMetadataSchema:
    """
    Загружает чанк объекта
    """
    logger.info("Отправляю буфер на s3.")
    res = s3_client.upload_part(
        Bucket=s3bucket,
        Body=chunk,  # чанки не могут быть меньше 5 мегабайт, нужно копить их в буфер
        UploadId=upload_id,
        PartNumber=chunk_num,
        Key=obj_key,
        ChecksumAlgorithm="SHA256",
        ChecksumSHA256=b64encode(sha256(chunk).digest()).decode("utf-8"),
    )
    logger.info("Отправил буфер на s3.")
    return PartMetadataSchema(PartNumber=chunk_num, ETag=res.get("ETag"))


@s3_client_wrapper
def complete(
    s3_client: BaseClient,
    s3bucket: str,
    obj_key: str,
    upload_id: str,
    parts: list[PartMetadataSchema],
):
    """
    Помечает, что все чанки загружены
    """
    logger.info("Комитим чанки.")
    parts: list[dict] = TypeAdapter(list[PartMetadataSchema]).dump_python(parts)

    return s3_client.complete_multipart_upload(
        Bucket=s3bucket,
        Key=obj_key,
        UploadId=upload_id,
        MultipartUpload={"Parts": sorted(parts, key=lambda k: k.get("PartNumber"))},
    )


@ExceptionHandler(ex_handler_map={Exception: None})
@s3_client_wrapper
def abort(s3_client: BaseClient, s3bucket: str, obj_key: str, upload_id: str):
    """
    Откатывает все загруженные чанки и созданный объект
    """
    logger.info("Откатываем чанки.")
    s3_client.abort_multipart_upload(
        Bucket=s3bucket,
        Key=obj_key,
        UploadId=upload_id,
    )
