# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright The Geneva Authors

# module for packaging workspace and UDFs
# also data spec for persisting the artifacts

import abc
import base64
import hashlib
import json
import logging
import uuid

import attrs
import cloudpickle
from typing_extensions import Self

from geneva import DockerWorkspacePackager
from geneva.transformer import UDF

_LOG = logging.getLogger(__name__)


class UDFBackend(abc.ABC):
    @classmethod
    @abc.abstractmethod
    def packager(cls) -> "UDFPackager":
        """Return the packager for this backend."""

    def to_bytes(self) -> bytes:
        return json.dumps(attrs.asdict(self)).encode()

    @classmethod
    def from_bytes(cls, data: bytes) -> Self:
        return cls(**json.loads(data.decode()))


@attrs.define
class UDFSpec:
    """Specification for a user-defined function.

    This is an holder of an arbitrary user-defined function,
    which can use an backend for marshalling.

    The most common is likely Docker + some kind of workspace
    persistence. However, we want to support more than just
    Docker, so we create this "out most" abstraction to allow
    for more flexibility.
    """

    # the name of the udf
    name: str = attrs.field(validator=attrs.validators.min_len(1))

    # the packaging backend for the udf
    backend: str = attrs.field()

    @backend.validator
    def _validate_backend(self, attribute, value) -> None:
        backend_names = [cls.__name__ for cls in UDFBackend.__subclasses__()]
        unique_backend_names = set(backend_names)

        # ruff: noqa: ERA001
        # TODO - deduplication not working right sometimes
        # if len(backend_names) != len(unique_backend_names):
        #     duplicates = []
        #     while unique_backend_names:
        #         name = backend_names.pop()
        #         if name not in unique_backend_names:
        #             duplicates.append(name)
        #         else:
        #             unique_backend_names.discard(name)
        #     raise ValueError(
        #         "Duplicate backend names found, "
        #         "please don't define the same backend name twice."
        #         f"\nDuplicates: {duplicates}"
        #     )

        if value not in unique_backend_names:
            raise ValueError(f"Unknown backend: {value}")

    udf_payload: bytes = attrs.field()

    # the payload for the runner -- This is a HACK for allowing phalanx knowing
    # how to dispatch the UDF job. Make sure changes here are compatible to
    # parsing in phalanx.
    runner_payload: bytes | None = attrs.field(default=None)

    @classmethod
    def udf_from_spec(cls, data) -> UDF:
        # TODO: load the spec and find the backend,
        # then call the packager to do the next level unmarshalling
        pass


@attrs.define
class DockerUDFSpecV1(UDFBackend):
    """Specification for a user-defined function that runs in a Docker container.
    -- Version 1

    In this packaging spec, the python interpreter is assumed to be correctly
    setup in the container, and the user-defined function is expected to load
    using cloudpickle. With the option of downloading additional workspace
    files from a remote location (S3, GCS, etc).
    """

    # the image to run the udf in
    image: str = attrs.field()

    # the tag of the image
    tag: str | None = attrs.field()

    # spec from https://docs.docker.com/reference/cli/docker/image/tag/#description
    @tag.validator
    def _validate_tag(self, attribute, value: str | None) -> None:
        if value is None:
            return
        if len(value) > 128:
            raise ValueError("Tag must be less than 128 characters.")
        if not all(c.isalpha() or c.isnumeric or c in {"_", ".", "-"} for c in value):
            raise ValueError("Tag must be valid alphanumeric.")

    # optionally have a zip of the workspace and store it separately
    # this should be the path to the zip file on (S3, GCS, etc)
    workspace_zip: str | None = attrs.field()

    # the checksum of the workspace zip
    workspace_checksum: str | None = attrs.field()

    @workspace_checksum.validator
    def _validate_workspace_checksum(self, attribute, value) -> None:
        if self.workspace_zip and not value:
            raise ValueError(
                "Workspace checksum must not be empty when a workspace is provided."
            )

    # the udf pickle to run
    udf_pickle: bytes = attrs.field()

    @udf_pickle.validator
    def _validate_udf_pickle(self, attribute, value) -> None:
        if not value:
            raise ValueError("UDF pickle must not be empty.")
        udf = cloudpickle.loads(value)
        if not isinstance(udf, UDF):
            raise ValueError("UDF pickle must contain a UDF object.")

    # the checksum of the udf pickle
    udf_checksum: str = attrs.field(init=False)

    def __attrs_post_init__(self) -> None:
        self.udf_checksum = hashlib.sha256(self.udf_pickle).hexdigest()

    @classmethod
    def packager(cls) -> "UDFPackager":
        return DockerUDFPackager()

    def to_bytes(self) -> bytes:
        self_as_dict = attrs.asdict(self)
        self_as_dict["udf_pickle"] = base64.b64encode(
            self_as_dict["udf_pickle"]
        ).decode("utf-8")
        return json.dumps(self_as_dict).encode()

    @classmethod
    def from_bytes(cls, data: bytes) -> Self:
        self_as_dict = json.loads(data.decode())
        self_as_dict["udf_pickle"] = base64.b64decode(
            self_as_dict["udf_pickle"].encode("utf-8")
        )

        checksum = self_as_dict.pop("udf_checksum")  # not part of the init
        val = cls(**self_as_dict)
        val.udf_checksum = checksum

        return val


class UDFPackager(abc.ABC):
    """Packager for user-defined functions."""

    @abc.abstractmethod
    def marshal(self, udf: UDF) -> UDFSpec:
        """Marshal a user-defined function."""

    @abc.abstractmethod
    def unmarshal(self, spec: UDFSpec) -> UDF:
        """Unmarshal a user-defined function."""


class DockerUDFPackager(UDFPackager):
    def marshal(self, udf: UDF) -> UDFSpec:
        image_tag = uuid.uuid4().hex

        packager = DockerWorkspacePackager()
        _LOG.info("Building UDF image: %s", image_tag)
        image_name = packager.build(image_tag, platform="linux/amd64", cuda=udf.cuda)
        tag = image_name.split(":")[-1]
        image_name = ":".join(image_name.split(":")[:-1])
        _LOG.info("Pushing UDF image: %s", image_tag)
        packager.push(image_tag)
        _LOG.info("Pushed UDF image: %s", image_tag)

        udf_pickle = cloudpickle.dumps(udf)

        return UDFSpec(
            name=udf.name,
            backend=DockerUDFSpecV1.__name__,
            udf_payload=DockerUDFSpecV1(
                image=image_name,
                tag=tag,
                workspace_zip=None,
                workspace_checksum=None,
                udf_pickle=udf_pickle,
            ).to_bytes(),
            runner_payload=json.dumps(
                {
                    "image": image_name,
                }
            ).encode(),
        )

    def unmarshal(self, spec: UDFSpec) -> UDF:
        if spec.backend != DockerUDFSpecV1.__name__:
            raise ValueError("Invalid backend for UDF spec.")

        docker_spec = DockerUDFSpecV1.from_bytes(spec.udf_payload)
        udf = cloudpickle.loads(docker_spec.udf_pickle)
        if not isinstance(udf, UDF):
            raise ValueError("UDF pickle must contain a UDF object.")
        return udf
