"""
BaseWorker methods for training.
"""

import functools
from collections.abc import Generator
from contextlib import contextmanager
from pathlib import Path
from typing import NewType
from uuid import UUID

from arkindex.multipart import MultipartUpload
from arkindex_worker import logger
from arkindex_worker.utils import close_delete_file, create_tar_zst_archive

DirPath = NewType("DirPath", Path)
"""Path to a directory"""

Hash = NewType("Hash", str)
"""MD5 Hash"""

FileSize = NewType("FileSize", int)
"""File size"""


@contextmanager
def create_archive(path: DirPath) -> Generator[tuple[Path, FileSize, Hash]]:
    """
    Create a tar archive from the files at the given location then compress it to a zst archive.

    Yield its location, its size and its hash.

    :param path: Create a compressed tar archive from the files
    :returns: The location of the created archive, its size and its hash
    """
    assert path.is_dir(), "create_archive needs a directory"

    zst_descriptor, zst_archive, archive_hash = create_tar_zst_archive(path)

    # Get content hash, archive size and hash
    yield zst_archive, zst_archive.stat().st_size, archive_hash

    # Remove the zst archive
    close_delete_file(zst_descriptor, zst_archive)


def build_clean_payload(**kwargs):
    """
    Remove null attributes from an API body payload
    """
    return {key: value for key, value in kwargs.items() if value is not None}


def skip_if_read_only(func):
    """
    Return shortly in case the is_read_only property is evaluated to True
    """

    @functools.wraps(func)
    def wrapper(self, *args, **kwargs):
        if getattr(self, "is_read_only", False):
            logger.warning(
                "Cannot perform this operation as the worker is in read-only mode"
            )
            return
        return func(self, *args, **kwargs)

    return wrapper


class TrainingMixin:
    """
    A mixin helper to create a new model version easily.
    You may use `publish_model_version` to publish a ready model version directly, or
    separately create the model version then publish it (e.g to store training metrics).
    Stores the currently handled model version as `self.model_version`.
    """

    model_version = None

    @property
    def is_finetuning(self) -> bool:
        """
        Whether or not this worker is fine-tuning an existing model version.
        """
        return bool(self.model_version_id)

    @skip_if_read_only
    def publish_model_version(
        self,
        model_path: DirPath,
        model_id: str,
        tag: str | None = None,
        description: str | None = None,
        configuration: dict | None = None,
        parent: str | UUID | None = None,
    ):
        """
        Publish a unique version of a model in Arkindex, identified by its hash.
        In case the `create_model_version` method has been called, reuses that model
        instead of creating a new one.

        :param model_path: Path to the directory containing the model version's files.
        :param model_id: ID of the model
        :param tag: Tag of the model version
        :param description: Description of the model version
        :param configuration: Configuration of the model version
        :param parent: ID of the parent model version
        """

        configuration = configuration or {}

        # Create the zst archive, get its hash and size
        with create_archive(path=model_path) as (
            path_to_archive,
            size,
            hash,
        ):
            # Update an existing model version with hash, size and any other defined attribute
            if self.model_version:
                assert self.model_version.get("model_id") == model_id, (
                    "Given `model_id` does not match the current model version"
                )
                self.update_model_version(
                    size=size,
                    archive_hash=hash,
                    tag=tag,
                    description=description,
                    configuration=configuration,
                    parent=parent,
                )

            # Create a new model version with hash and size
            else:
                self.create_model_version(
                    model_id=model_id,
                    size=size,
                    archive_hash=hash,
                    tag=tag,
                    description=description,
                    configuration=configuration,
                    parent=parent,
                )

            # Upload the archive in multiple parts (supports huge files)
            self.upload_to_s3(path_to_archive)

    @skip_if_read_only
    def create_model_version(
        self,
        model_id: str,
        size: FileSize,
        archive_hash: Hash,
        tag: str | None = None,
        description: str | None = None,
        configuration: dict | None = None,
        parent: str | UUID | None = None,
    ):
        """
        Create a new version of the specified model with its base attributes.
        Once successfully created, the model version is accessible via `self.model_version`.

        :param size: Size of uploaded archive
        :param hash: MD5 hash of the uploaded archive
        :param tag: Tag of the model version
        :param description: Description of the model version
        :param configuration: Configuration of the model version
        :param parent: ID of the parent model version
        """
        assert not self.model_version, "A model version has already been created."

        configuration = configuration or {}
        self.model_version = self.api_client.request(
            "CreateModelVersion",
            id=model_id,
            body=build_clean_payload(
                size=size,
                archive_hash=archive_hash,
                tag=tag,
                description=description,
                configuration=configuration,
                parent=parent,
            ),
        )

        logger.info(
            f"Model version ({self.model_version['id']}) was successfully created."
        )

    @skip_if_read_only
    def update_model_version(
        self,
        size: FileSize,
        archive_hash: Hash,
        tag: str | None = None,
        description: str | None = None,
        configuration: dict | None = None,
        parent: str | UUID | None = None,
    ):
        """
        Update the current model version with the given attributes.

        :param size: Size of uploaded archive
        :param hash: MD5 hash of the uploaded archive
        :param tag: Tag of the model version
        :param description: Description of the model version
        :param configuration: Configuration of the model version
        :param parent: ID of the parent model version
        """
        assert self.model_version, "No model version has been created yet."
        self.model_version = self.api_client.request(
            "UpdateModelVersion",
            id=self.model_version["id"],
            body=build_clean_payload(
                size=size,
                archive_hash=archive_hash,
                tag=tag,
                description=description,
                configuration=configuration,
                parent=parent,
            ),
        )
        logger.info(
            f"Model version ({self.model_version['id']}) was successfully updated."
        )

    @skip_if_read_only
    def upload_to_s3(self, archive_path: Path) -> None:
        """
        Upload the archive of the model's files to an Amazon s3 compatible storage in multiple parts
        """
        assert self.model_version, (
            "You must create the model version before uploading an archive."
        )
        assert self.model_version["state"] != "Available", (
            "The model version is already marked as available."
        )

        multipart = MultipartUpload(
            client=self.api_client,
            file_path=archive_path,
            object_type="model_version",
            object_id=str(self.model_version["id"]),
        )
        try:
            multipart.upload()
            multipart.complete()
        except Exception:
            multipart.abort()
            raise
        else:
            logger.info(
                f"Model version ({self.model_version['id']}) archive was successfully uploaded and is now available."
            )
