# Copyright 2024 The Debusine Developers
# See the AUTHORS file at the top-level directory of this distribution
#
# This file is part of Debusine. It is subject to the license terms
# in the LICENSE file found in the top-level directory of this
# distribution. No part of Debusine, including this file, may be copied,
# modified, propagated, or distributed except according to the terms
# contained in the LICENSE file.

"""Package upload workflow."""
try:
    import pydantic.v1 as pydantic
except ImportError:
    import pydantic  # type: ignore

from debusine.artifacts.models import ArtifactCategory
from debusine.client.models import LookupChildType
from debusine.db.models import WorkRequest
from debusine.server.collections.lookup import (
    lookup_multiple,
    lookup_single,
    reconstruct_lookup,
)
from debusine.server.tasks.models import (
    PackageUploadData,
    PackageUploadInput,
    PackageUploadTarget,
)
from debusine.server.tasks.wait.models import ExternalDebsignData
from debusine.server.workflows import Workflow, WorkflowValidationError
from debusine.server.workflows.models import (
    PackageUploadWorkflowData,
    WorkRequestWorkflowData,
)
from debusine.signing.tasks.models import DebsignData
from debusine.tasks.models import (
    BaseDynamicTaskData,
    LookupMultiple,
    LookupSingle,
    MakeSourcePackageUploadData,
    MakeSourcePackageUploadInput,
    MergeUploadsData,
    MergeUploadsInput,
    TaskTypes,
    empty_lookup_multiple,
)


class PackageUploadWorkflow(
    Workflow[PackageUploadWorkflowData, BaseDynamicTaskData]
):
    """Run package uploads for a set of architectures."""

    TASK_NAME = "package_upload"

    def __init__(self, work_request: "WorkRequest"):
        """Instantiate a Workflow with its database instance."""
        super().__init__(work_request)

        self._artifacts_pending_upload: list[LookupSingle] = []

    def validate_input(self) -> None:
        """Raise WorkflowValidationError if needed."""
        if (
            self.data.source_artifact is None
            and self.data.binary_artifacts == empty_lookup_multiple()
        ):
            raise WorkflowValidationError(
                '"source_artifact" or "binary_artifacts" must be set'
            )

        if self.data.source_artifact is not None:
            source_artifact = lookup_single(
                self.data.source_artifact,
                self.workspace,
                user=self.work_request.created_by,
                workflow_root=self.work_request.get_workflow_root(),
                expect_type=LookupChildType.ARTIFACT,
            ).artifact

            if (
                source_artifact.category == ArtifactCategory.SOURCE_PACKAGE
                and (self.data.vendor is None or self.data.codename is None)
            ):
                # MakeSourcePackageUpload will be used but
                # self.data.vendor or self.data.codename are not set
                raise WorkflowValidationError(
                    f'"vendor" and "codename" are required when source '
                    f'artifact category is {ArtifactCategory.SOURCE_PACKAGE}'
                )

        if self.data.merge_uploads and (
            self.data.vendor is None or self.data.codename is None
        ):
            # MergeUploads will be used but self.data.vendor or
            # self.data.codename are not set
            raise WorkflowValidationError(
                '"vendor" and "codename" are required when '
                '"merge_uploads" is set'
            )

    def populate(self) -> None:
        """Create package upload work requests."""
        if (source_artifact := self.data.source_artifact) is not None:
            artifact = lookup_single(
                source_artifact,
                self.workspace,
                user=self.work_request.created_by,
                workflow_root=self.work_request.get_workflow_root(),
                expect_type=LookupChildType.ARTIFACT,
            ).artifact
            since_version = self.data.since_version
            target_distribution = self.data.target_distribution

            assert self.data.source_artifact

            if artifact.category == ArtifactCategory.SOURCE_PACKAGE:
                package_upload_source = (
                    self._populate_make_source_package_upload(
                        self.data.source_artifact,
                        since_version=since_version,
                        target_distribution=target_distribution,
                    )
                )
                self._artifacts_pending_upload.append(package_upload_source)

            else:
                self._artifacts_pending_upload.append(self.data.source_artifact)

        # Upload binary_artifacts
        results = lookup_multiple(
            self.data.binary_artifacts,
            self.workspace,
            user=self.work_request.created_by,
            workflow_root=self.work_request.get_workflow_root(),
            expect_type=LookupChildType.ARTIFACT_OR_PROMISE,
        )

        for result in results:
            self._artifacts_pending_upload.append(reconstruct_lookup(result))

        self._populate_package_uploads()

    def _populate_make_source_package_upload(
        self,
        lookup: LookupSingle,
        *,
        since_version: str | None,
        target_distribution: str | None,
    ) -> LookupSingle:
        """
        Create work request for MakeSourcePackageUpload.

        :returns: Lookup of the signed artifact in the collection
        """
        wr = self.work_request_ensure_child(
            task_name="makesourcepackageupload",
            task_data=MakeSourcePackageUploadData(
                input=MakeSourcePackageUploadInput(source_artifact=lookup),
                target_distribution=target_distribution,
                since_version=since_version,
                environment=(
                    f"{self.data.vendor}/match:"
                    f"codename={self.data.codename}"
                ),
            ).dict(exclude_unset=True),
            workflow_data=WorkRequestWorkflowData(
                display_name="Make .changes file",
                step="source-package-upload",
            ),
        )

        artifact_name = "package-upload-source"
        self.provides_artifact(wr, ArtifactCategory.UPLOAD, artifact_name)
        return f"internal@collections/name:{artifact_name}"

    def _populate_signer(
        self,
        signer_identifier: str,
        unsigned_artifact: LookupSingle,
        key: LookupSingle | None,
    ) -> LookupSingle:
        """
        Create a signer: Debsign or ExternalDebsign depending on self.data.

        :param signer_identifier: identifier for the end user of what is going
          to be signed
        :param unsigned_artifact: unsigned artifact for the signer
        :param key: key to be used by the signer. If None, will use
          ExternalDebsign
        :returns: lookup of the signed artifact
        """
        if key is not None:
            signer = self.work_request_ensure_child(
                task_name="debsign",
                task_data=DebsignData(unsigned=unsigned_artifact, key=key).dict(
                    exclude_unset=True
                ),
                workflow_data=WorkRequestWorkflowData(
                    display_name=f"Sign upload for {signer_identifier}",
                    step=f"debsign-{signer_identifier}",
                ),
            )
            self.requires_artifact(signer, unsigned_artifact)

        else:
            signer = self.work_request_ensure_child(
                task_name="externaldebsign",
                task_data=ExternalDebsignData(unsigned=unsigned_artifact).dict(
                    exclude_unset=True
                ),
                workflow_data=WorkRequestWorkflowData(
                    display_name=(
                        "Wait for signature on "
                        f"upload for {signer_identifier}"
                    ),
                    step=f"external-debsign-{signer_identifier}",
                ),
            )

        name_signed_artifact = (
            f"package-upload-signed-{unsigned_artifact}".replace(":", "_")
            .replace("@", "_")
            .replace("/", "_")
        )

        self.provides_artifact(
            signer,
            ArtifactCategory.UPLOAD,
            name_signed_artifact,
        )

        return f"internal@collections/name:{name_signed_artifact}"

    def _populate_package_uploads(self) -> None:
        """
        Create work request(s) to upload the artifacts.

        Depending on self.data.merge_uploads will:
        -Create one work request for each self._artifacts_pending_upload
        -Create a single work request of type MergeUpload
        """
        if self.data.merge_uploads:
            wr = self.work_request_ensure_child(
                task_name="mergeuploads",
                task_data=MergeUploadsData(
                    environment=(
                        f"{self.data.vendor}/match:"
                        f"codename={self.data.codename}"
                    ),
                    input=MergeUploadsInput(
                        uploads=LookupMultiple.parse_obj(
                            self._artifacts_pending_upload
                        )
                    ),
                ).dict(exclude_unset=True),
                workflow_data=WorkRequestWorkflowData(
                    display_name="Merge uploads",
                    step="merge-uploads",
                ),
            )

            for upload in self._artifacts_pending_upload:
                self.requires_artifact(wr, upload)

            merged_artifact_name = "package-upload-merged"
            self.provides_artifact(
                wr, ArtifactCategory.UPLOAD, merged_artifact_name
            )

            self._artifacts_pending_upload = [
                f"internal@collections/name:{merged_artifact_name}"
            ]

        for upload in self._artifacts_pending_upload:
            identifier = str(upload)

            if self.data.key is not None or self.data.require_signature:

                upload = self._populate_signer(
                    identifier, upload, key=self.data.key
                )

            uploader = self.work_request_ensure_child(
                task_name="packageupload",
                task_data=PackageUploadData(
                    input=PackageUploadInput(upload=upload),
                    target=pydantic.parse_obj_as(
                        PackageUploadTarget, self.data.target
                    ),
                ).dict(exclude_unset=True),
                workflow_data=WorkRequestWorkflowData(
                    display_name=f"Package upload {identifier}",
                    step=f"package-upload-{identifier}",
                ),
                task_type=TaskTypes.SERVER,
            )

            self.requires_artifact(uploader, upload)

    def get_label(self) -> str:
        """Return the task label."""
        return "run package uploads"
