"""
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_worker import logger
from arkindex_worker.utils import close_delete_file, create_tar_zst_archive
from teklia_toolbox.uploads import MultipartUpload

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