# Source code for debusine.tasks.inputs.inputs

    # Copyright © 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.
    """Input fields for tasks."""

    import abc
    import enum
    from collections.abc import Callable, Generator
    from collections.abc import Collection as ABCCollection
    from operator import attrgetter
    from types import GenericAlias
    from typing import (
        Any,
        ClassVar,
        NotRequired,
        Self,
        TypedDict,
        Unpack,
        cast,
        override,
    )

    from debusine.artifacts.models import ArtifactCategory
    from debusine.client.models import RelationType
    from debusine.tasks.executors import ExecutorImageCategory
    from debusine.tasks.inputs import serializers
    from debusine.tasks.inputs.resolver import FieldResolver
    from debusine.tasks.models import (
        BaseDynamicTaskData,
        BaseTaskInputArtifact,
        ExtraExternalRepository,
        ExtraRepository,
        InputArtifactMultiple,
        InputArtifactSingle,
        InputCollectionSingle,
        LookupMultiple,
        LookupSingle,
    )
    from debusine.utils import extract_generic_type_arguments

    class Stage(enum.IntEnum):
        """Sorted enumeration for stages in preparing tasks for scheduling."""

        #: Inputs needed to prepare tasks for configuration lookup
        PENDING = 1
        #: Inputs needed to prepare tasks for running
        RUNNING = 2

    class TaskFieldContainer:
        """Class that can contain task fields."""

        #: Inputs defined in this task class
        inputs: dict[str, "TaskInput[Any]"] = {}

        @override
        def __init_subclass__(cls, **kwargs: Any) -> None:
            """Merge input field definitions from multiple parents."""
            super().__init_subclass__()
            inputs: dict[str, TaskInput[Any]] = {}
            for parent in reversed(cls.__mro__):
                if not issubclass(parent, TaskFieldContainer):
                    continue
                inputs.update(parent.inputs)
            inputs.update(cls.inputs)
            cls.inputs = inputs

        def __init__(self) -> None:
            """Initialize a TaskFieldContainer."""
            super().__init__()
            #: If set, when fields of this container are resolved, fields of the
            #: child container are also resolved, using the resolver instantiated
            #: for the parent container
            self.child_task_field_container: TaskFieldContainer | None = None

        def iter_all_fields(
            self, stage: Stage | None = None
        ) -> Generator[tuple["TaskFieldContainer", "TaskInput[Any]"]]:
            """Enumerate the fields to resolve, and their container object."""
            if stage is None:
                stage = Stage.RUNNING

            for inp in self.inputs.values():
                if inp.stage > stage:
                    continue
                yield self, inp
            if self.child_task_field_container is not None:
                yield from self.child_task_field_container.iter_all_fields(stage)

        def field_is_resolved(self, name: str) -> bool:
            """Return True if the field has been resolved."""
            return name in self.__dict__

        def save_field_values(self, dynamic_data: BaseDynamicTaskData) -> None:
            """
            Save field values in a BaseDynamicTaskData.

            Fields not defined in this container that are already saved are
            preserved: this allows multiple field containers to serialize to the
            same BaseDynamicTaskData, as long as their names do not conflict.
            """
            task_input_fields = dynamic_data.task_input_fields
            for cont, inp in self.iter_all_fields():
                task_input_fields[inp.name] = inp.serialize_value(
                    getattr(cont, inp.name)
                )

            # We need to perform an assignment to task_input_fields and not just
            # modify it in place, otherwise the pydantic model will consider the
            # field unset and not serialize it
            dynamic_data.task_input_fields = task_input_fields

        def restore_field_values(self, dynamic_data: BaseDynamicTaskData) -> None:
            """
            Restore field values serialized in a BaseDynamicTaskData.

            Stored fields with names not defined in this container are ignored:
            this allows multiple field containers to serialize to the same
            BaseDynamicTaskData, as long as their names do not conflict.
            """
            values = dynamic_data.task_input_fields
            for cont, inp in self.iter_all_fields():
                if inp.name not in values:
                    continue
                setattr(cont, inp.name, inp.deserialize_value(values[inp.name]))

    class TaskInputKwargs(TypedDict):
        """Argument typing for TaskInput constructor."""

        stage: NotRequired[Stage]
        verbose_name: NotRequired[str]
        help_text: NotRequired[str]
        hidden: NotRequired[bool]
        internal: NotRequired[bool]

    class TaskInput[ValueType](abc.ABC):
        """Representation for one or more input artifacts for a task."""

        name: str
        stage: Stage
        serializer: ClassVar[serializers.TaskInputSerializer[ValueType]]

        def get_resolver_method(
            self, resolver: FieldResolver
        ) -> Callable[[Self], ValueType]:
            """Override this to forward to a FieldResolver method."""
            raise NotImplementedError(
                "This input field does not forward resolution to the FieldResolver."
            )

        def resolve(self, resolver: FieldResolver) -> ValueType:
            """Resolve the input from task data."""
            method = self.get_resolver_method(resolver)
            return method(self)

        def get_artifact_ids(self, value: ValueType) -> list[int]:  # noqa: ARG002, U100
            """Return artifact IDs."""
            return []

        def __init__(
            self,
            stage: Stage = Stage.RUNNING,
            verbose_name: str | None = None,
            help_text: str | None = None,
            hidden: bool = False,
            internal: bool = False,
        ) -> None:
            """
            Declare an input field.

            :param stage: stage of task preparation when the field is needed
            :param verbose_name: label to use instead of the autogenerated one
            :param help_text: description of the field
            :param hidden: set to True to hide the field from the UI
            :param internal: show this widget as an internal field in the UI
            """
            self.stage = stage
            self.verbose_name = verbose_name
            self.help_text = help_text
            self.hidden = hidden
            self.internal = internal

        def __set_name__(self, owner: type[TaskFieldContainer], name: str) -> None:
            """Set the field name at class construction."""
            self.name = name
            # Avoid sharing the inputs definition across all subclasses of
            # BaseExternalTask
            if "inputs" not in owner.__dict__:
                owner.inputs = owner.inputs.copy()
            owner.inputs[name] = self

        def __get__(
            self,
            obj: TaskFieldContainer,
            _: type[TaskFieldContainer] | None = None,
        ) -> ValueType:
            """Look up this input in the task object."""
            from debusine.tasks import TaskConfigError

            # If __get__ is called, it means that self.name is not in obj.__dict__,
            # which only happens if the field value has not yet been resolved
            raise TaskConfigError(f"Cannot access unresolved input {self.name!r}")

        def get_value(self, field_container: TaskFieldContainer) -> ValueType:
            """Return the current value of this field in the container."""
            return cast(ValueType, getattr(field_container, self.name))

        def serialize_value(self, value: ValueType) -> Any:
            """Serialize a value for storing in dynamic task data."""
            return self.serializer.serialize(value)

        def deserialize_value(self, value: Any) -> ValueType:
            """Serialize a value for storing in dynamic task data."""
            return self.serializer.deserialize(value)

    class DataFieldInputKwargs(TaskInputKwargs):
        """Argument typing for DataFieldInput constructor."""

        field: NotRequired[str]

    class DataFieldInput[LookupType, ValueType](TaskInput[ValueType]):
        """Input field that resolves values based on a task data field."""

        field: str
        getter: Callable[[str], LookupType]
        # Actual lookup type, filled by __init_subclass__
        lookup_type: type[LookupType]

        @override
        def __init_subclass__(cls, **kwargs: Any) -> None:
            """Fill lookup_type from typing information."""
            super().__init_subclass__(**kwargs)
            [cls.lookup_type, _] = extract_generic_type_arguments(
                cls, DataFieldInput
            )
            # Given that we use lookup_type for isinstance checks, if it is a
            # generic alias we can strip the parameterization and only keep the
            # original type. This allows to type LookupType as ABCCollection[foo]
            # and have isinstance checks against ABCCollection
            if isinstance(cls.lookup_type, GenericAlias):
                cls.lookup_type = cls.lookup_type.__origin__

        def __init__(
            self, field: str | None = None, **kwargs: Unpack[TaskInputKwargs]
        ) -> None:
            """
            Declare an input defined by a task data field.

            :param field: dot-separated list of lookup names in task data. Defaults
              to the input member name
            """
            super().__init__(**kwargs)
            if field is not None:
                self._set_field_name(field)

        @override
        def __set_name__(self, owner: type[TaskFieldContainer], name: str) -> None:
            """Set the field name at class construction."""
            super().__set_name__(owner, name)
            if not hasattr(self, "field"):
                self._set_field_name(name)

        def _set_field_name(self, name: str) -> None:
            """Set the attribute name used to look up the value in task input."""
            self.field = name
            self.getter = attrgetter(name)

    class DataListFieldInput[LookupType, ValueType](
        DataFieldInput[ABCCollection[LookupType], ValueType]
    ):
        """DataFieldInput that resolves to a list of type checked elements."""

        # Type of each list item, filled by __init_subclass__
        item_type: type[LookupType]

        @override
        def __init_subclass__(cls, **kwargs: Any) -> None:
            """Fill lookup_type from typing information."""
            super().__init_subclass__(**kwargs)
            [cls.item_type, _] = extract_generic_type_arguments(
                cls, DataListFieldInput
            )

    class BaseArtifactInputKwargs(DataFieldInputKwargs):
        """Argument typing for BaseArtifactInput constructor."""

        categories: NotRequired[ABCCollection[ArtifactCategory]]

    class BaseArtifactInput[LookupType, ValueType](
        DataFieldInput[LookupType, ValueType]
    ):
        """Look up artifacts enforcing their types."""

        categories: tuple[ArtifactCategory, ...] | None

        def __init__(
            self,
            *,
            categories: ABCCollection[ArtifactCategory] | None = None,
            **kwargs: Unpack[DataFieldInputKwargs],
        ) -> None:
            """
            Look up artifacts enforcing their type.

            :param field: dot-separated list of lookup names in task data
            :param categories: list of acceptable artifact categories
            """
            super().__init__(**kwargs)
            self.categories = tuple(categories) if categories is not None else None

        def check_category(
            self, info: BaseTaskInputArtifact, field_name: str | None = None
        ) -> None:
            """Check the category of a single input artifact."""
            # Prevent import loop
            from debusine.tasks import ensure_artifact_categories

            if self.categories is None:
                return

            # TODO: this should not happen, as we are checking the results of
            # database lookups which always have a category. However, info.category
            # can be None for compatibility during a transition time of merging
            # with InputArtifactMultiple, which needs to be able to deserialize
            # from only a list of artifact IDs
            assert info.category is not None

            ensure_artifact_categories(
                configuration_key=field_name or self.field,
                category=info.category,
                expected=self.categories,
            )

    class ArtifactListInput[LookupType, ValueType](
        BaseArtifactInput[ABCCollection[LookupType], ValueType],
        DataListFieldInput[LookupType, ValueType],
    ):
        """Input taking a list of artifact lookups, enforcing their types."""

    class SingleInput(BaseArtifactInput[LookupSingle, InputArtifactSingle]):
        """One single input artifact specified in a task data field."""

        serializer = serializers.PydanticSerializer(InputArtifactSingle)

        @override
        def resolve(self, resolver: FieldResolver) -> InputArtifactSingle:
            lookup = resolver.get_field(self)
            assert lookup is not None
            info = resolver.lookup_single_artifact(lookup, label=self.field)
            self.check_category(info)
            return info

        @override
        def get_artifact_ids(self, value: InputArtifactSingle) -> list[int]:
            return value.resolved_ids()

    class OptionalSingleInput(
        BaseArtifactInput[LookupSingle, InputArtifactSingle | None]
    ):
        """One optional single input artifact specified in a task data field."""

        serializer = serializers.OptionalPydanticSerializer(InputArtifactSingle)

        @override
        def resolve(self, resolver: FieldResolver) -> InputArtifactSingle | None:
            lookup = resolver.get_field(self)
            if lookup is None:
                return None
            info = resolver.lookup_single_artifact(lookup, label=self.field)
            self.check_category(info)
            return info

        @override
        def get_artifact_ids(self, value: InputArtifactSingle | None) -> list[int]:
            if value is None:
                return []

            return value.resolved_ids()

    class SingleInputList(
        ArtifactListInput[LookupSingle, list[InputArtifactSingle]]
    ):
        """A list of single input artifact lookups in a task data field."""

        serializer = serializers.PydanticListSerializer(InputArtifactSingle)

        @override
        def resolve(self, resolver: FieldResolver) -> list[InputArtifactSingle]:
            infos: list[InputArtifactSingle] = []
            for lookup in resolver.get_field_collection(self):
                assert isinstance(lookup, LookupSingle)
                info = resolver.lookup_single_artifact(lookup, label=self.field)
                self.check_category(info)
                infos.append(info)
            return infos

        @override
        def get_artifact_ids(self, value: list[InputArtifactSingle]) -> list[int]:
            return [
                artifact.artifact_id
                for artifact in value
                if artifact.artifact_id is not None
            ]

    class MultiInput(BaseArtifactInput[LookupMultiple, InputArtifactMultiple]):
        """A list of multiple artifacts lookups specified in a task data field."""

        serializer = serializers.PydanticSerializer(InputArtifactMultiple)

        @override
        def resolve(self, resolver: FieldResolver) -> InputArtifactMultiple:
            lookup = resolver.get_field(self)
            assert lookup is not None
            result = resolver.lookup_multiple_artifacts(lookup, label=self.field)
            if self.categories is not None:
                for i, artifact in enumerate(result.artifacts):
                    self.check_category(artifact, f"{self.field}[{i}]")
            return result

        @override
        def get_artifact_ids(self, value: InputArtifactMultiple) -> list[int]:
            return [
                artifact.artifact_id
                for artifact in value.artifacts
                if artifact.artifact_id is not None
            ]

    class MultiInputList(
        ArtifactListInput[LookupMultiple, list[InputArtifactMultiple]]
    ):
        """A multiple artifacts lookup specified in a task data field."""

        serializer = serializers.MultiInputListSerializer()

        @override
        def resolve(self, resolver: FieldResolver) -> list[InputArtifactMultiple]:
            infos: list[InputArtifactMultiple] = []
            for lookup_idx, lookup in enumerate(
                resolver.get_field_collection(self)
            ):
                result = resolver.lookup_multiple_artifacts(
                    lookup, label=self.field
                )
                if self.categories is not None:
                    for result_idx, artifact in enumerate(result.artifacts):
                        self.check_category(
                            artifact, f"{self.field}[{lookup_idx}][{result_idx}]"
                        )
                infos.append(result)
            return infos

        @override
        def get_artifact_ids(self, value: list[InputArtifactMultiple]) -> list[int]:
            ids = []

            for input_artifact_multiple in value:
                ids += [
                    artifact.artifact_id
                    for artifact in input_artifact_multiple.artifacts
                    if artifact.artifact_id is not None
                ]
            return ids

    class UploadArtifactsInput(
        BaseArtifactInput[LookupMultiple, InputArtifactMultiple]
    ):
        """A multiple artifact lookup finding all related binaries in an upload."""

        serializer = serializers.PydanticSerializer(InputArtifactMultiple)

        @override
        def __init__(self, **kwargs: Unpack[BaseArtifactInputKwargs]) -> None:
            categories_with_upload = [ArtifactCategory.UPLOAD]
            categories_with_upload.extend(kwargs.get("categories") or ())
            kwargs["categories"] = categories_with_upload
            super().__init__(**kwargs)

        @override
        def resolve(self, resolver: FieldResolver) -> InputArtifactMultiple:
            lookup = resolver.get_field(self)
            assert lookup is not None

            # Incrementally build the result set
            result: list[BaseTaskInputArtifact] = []

            # Lookup the initial set of packages
            upload_ids: list[int] = []
            initial = resolver.lookup_multiple_artifacts(
                lookup, label=self.field
            ).artifacts
            for i, artifact in enumerate(initial):
                if artifact.category == ArtifactCategory.UPLOAD:
                    # File uploads for later expansion
                    assert artifact.artifact_id is not None
                    upload_ids.append(artifact.artifact_id)
                else:
                    self.check_category(artifact, f"{self.field}[{i}]")
                    result.append(artifact)

            # Expand with binaries related to uploads
            if upload_ids:
                result.extend(
                    resolver.find_related_artifacts(
                        upload_ids,
                        target_category=ArtifactCategory.BINARY_PACKAGE,
                        relation_type=RelationType.EXTENDS,
                        label=self.field,
                    ).artifacts
                )

            return InputArtifactMultiple(
                lookup=lookup, label=self.field, artifacts=result
            )

    class EnvironmentInputKwargs(BaseArtifactInputKwargs):
        """Argument typing for EnvironmentInput constructor."""

        image_category: NotRequired[ExecutorImageCategory | None]
        set_backend: NotRequired[bool]
        try_variant: NotRequired[bool]

    class EnvironmentInput(BaseArtifactInput[LookupSingle, InputArtifactSingle]):
        """One environment artifact specified in a task data field."""

        serializer = serializers.PydanticSerializer(InputArtifactSingle)

        def __init__(
            self,
            *,
            image_category: ExecutorImageCategory | None = None,
            set_backend: bool = True,
            try_variant: bool = True,
            **kwargs: Unpack[BaseArtifactInputKwargs],
        ) -> None:
            """
            Look up artifacts enforcing their type.

            :param field: dot-separated list of lookup names in task data
            :param categories: list of acceptable artifact categories
            :param image_category: try to use an environment with this image
              category; defaults to the image category needed by the executor
              for `self.backend`
            :param set_backend: if True (default), try to use an environment
              matching `self.backend`
            :param try_variant: if True (default), try to use an environment
              whose variant is `self.name`, but fall back to looking up an
              environment without a variant if the first lookup fails
            """
            super().__init__(**kwargs)
            self.image_category = image_category
            self.set_backend = set_backend
            self.try_variant = try_variant

        @override
        def get_resolver_method(
            self, resolver: FieldResolver
        ) -> Callable[[Self], InputArtifactSingle]:
            return resolver.resolve_environmentinput

        @override
        def get_artifact_ids(self, value: InputArtifactSingle) -> list[int]:
            return value.resolved_ids()

    class OptionalEnvironmentInput(
        BaseArtifactInput[LookupSingle, InputArtifactSingle | None]
    ):
        """One optional environment artifact specified in a task data field."""

        serializer = serializers.OptionalPydanticSerializer(InputArtifactSingle)

        @override
        def get_resolver_method(
            self, resolver: FieldResolver
        ) -> Callable[[Self], InputArtifactSingle | None]:
            return resolver.resolve_optionalenvironmentinput

    class SuiteArchiveInput(
        DataFieldInput[
            LookupSingle, tuple[InputCollectionSingle, InputCollectionSingle | None]
        ]
    ):
        """Look up a suite and an archive based on a task data field."""

        serializer = serializers.SuiteArchiveInputSerializer()

        @override
        def get_resolver_method(
            self, resolver: FieldResolver
        ) -> Callable[
            [Self], tuple[InputCollectionSingle, InputCollectionSingle | None]
        ]:
            return resolver.resolve_suitearchiveinput

    class ExtraRepositoriesInput(
        DataFieldInput[list[ExtraRepository], list[ExtraExternalRepository] | None]
    ):
        """Look up extra repositories."""

        serializer = serializers.ExtraRepositoriesSerializer()

        @override
        def get_resolver_method(
            self, resolver: FieldResolver
        ) -> Callable[[Self], list[ExtraExternalRepository] | None]:
            return resolver.resolve_extrarepositoriesinput

    class DebusineFQDNInput(TaskInput[str]):
        """Resolve to the Debusine FQDN."""

        serializer = serializers.ScalarSerializer(str)

        @override
        def resolve(self, resolver: FieldResolver) -> str:
            return resolver.get_server_setting("DEBUSINE_FQDN")

    class BuildArchitectureInput(DataFieldInput[str, str]):
        """
        Build architecture for the task.

        This currently looks up a value in the task data, and uses a fallback value
        if it is not set.
        """

        serializer = serializers.ScalarSerializer(str)

        @override
        def resolve(self, resolver: FieldResolver) -> str:
            if arch := resolver.get_field(self):
                return arch
            return resolver.get_fallback_architecture()

    class UnsignedMultiInput(
        BaseArtifactInput[LookupMultiple, InputArtifactMultiple]
    ):
        """A multiple artifact lookup finding all related artifacts."""

        serializer = serializers.PydanticSerializer(InputArtifactMultiple)

        @override
        def resolve(self, resolver: FieldResolver) -> InputArtifactMultiple:
            lookup = resolver.get_field(self)
            assert lookup is not None

            result: list[BaseTaskInputArtifact] = []

            signed_artifacts: set[int] = set()

            for signed in resolver.lookup_multiple_artifacts(
                lookup, label=self.field
            ).artifacts:
                assert signed.artifact_id
                signed_artifacts.add(signed.artifact_id)

            signing_inputs: set[int] = set()

            for signing_input in resolver.find_related_artifacts(
                signed_artifacts,
                target_category=ArtifactCategory.SIGNING_INPUT,
                relation_type=RelationType.RELATES_TO,
                label=self.field,
            ).artifacts:
                assert signing_input.artifact_id
                result.append(signing_input)
                signing_inputs.add(signing_input.artifact_id)

            result += resolver.find_related_artifacts(
                signing_inputs,
                target_category=ArtifactCategory.BINARY_PACKAGE,
                relation_type=RelationType.RELATES_TO,
                label=self.field,
            ).artifacts

            def artifact_id(result: BaseTaskInputArtifact) -> int:
                assert result.artifact_id is not None
                return result.artifact_id

            result.sort(key=artifact_id)

            return InputArtifactMultiple(
                lookup=lookup, label=self.field, artifacts=result
            )