# 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.

"""
Server-side database interaction for worker tasks.

Worker tasks mostly run on external workers, but they also run some
pre-dispatch code on the server which needs to access the database.  To
avoid having to import Django from :py:mod:`debusine.tasks`, we define an
interface that is implemented by server code.
"""

from abc import ABC, abstractmethod
from collections.abc import Collection as AbcCollection
from dataclasses import dataclass
from typing import Any, TYPE_CHECKING, overload

from debusine.artifacts.models import (
    ArtifactCategory,
    ArtifactData,
    CollectionCategory,
)
from debusine.client.models import RelationType
from debusine.tasks.models import (
    BackendType,
    LookupMultiple,
    LookupSingle,
    build_key_value_lookup_segment,
    build_lookup_string_segments,
    parse_key_value_lookup_segment,
    parse_lookup_string_segments,
)

if TYPE_CHECKING:
    from debusine.tasks import BaseTask
    from debusine.tasks.executors import ExecutorImageCategory


@dataclass
class ArtifactInfo:
    """Information about an artifact."""

    id: int
    category: str
    data: ArtifactData


class MultipleArtifactInfo(tuple[ArtifactInfo, ...]):
    """Information about multiple artifacts."""

    def get_ids(self) -> list[int]:
        """Return the ID of each artifact."""
        return [item.id for item in self]


@dataclass
class CollectionInfo:
    """Information about a collection."""

    id: int
    scope_name: str
    workspace_name: str
    category: str
    name: str
    data: dict[str, Any]


class TaskDatabaseInterface(ABC):
    """Interface for interacting with the database from worker tasks."""

    @overload
    @abstractmethod
    def lookup_single_artifact(
        self,
        lookup: LookupSingle,
        default_category: CollectionCategory | None = None,
    ) -> ArtifactInfo: ...

    @overload
    @abstractmethod
    def lookup_single_artifact(
        self,
        lookup: None,
        default_category: CollectionCategory | None = None,
    ) -> None: ...

    @abstractmethod
    def lookup_single_artifact(
        self,
        lookup: LookupSingle | None,
        default_category: CollectionCategory | None = None,
    ) -> ArtifactInfo | None:
        """
        Look up a single artifact.

        :param lookup: A :ref:`lookup-single`.
        :param default_category: If the first segment of a string lookup
          (which normally identifies a collection) does not specify a
          category, use this as the default category.
        :return: Information about the artifact, or None if the provided
          lookup is None (for convenience in some call sites).
        :raises KeyError: if the lookup does not resolve to an item.
        :raises LookupError: if the lookup is invalid in some way, or does
          not resolve to an artifact.
        """

    @abstractmethod
    def lookup_multiple_artifacts(
        self,
        lookup: LookupMultiple | None,
        default_category: CollectionCategory | None = None,
    ) -> MultipleArtifactInfo:
        """
        Look up multiple artifacts.

        :param lookup: A :ref:`lookup-multiple`.
        :param default_category: If the first segment of a string lookup
          (which normally identifies a collection) does not specify a
          category, use this as the default category.
        :return: Information about each artifact.
        :raises KeyError: if any of the lookups does not resolve to an item.
        :raises LookupError: if any of the lookups is invalid in some way,
          or does not resolve to an artifact.
        """

    @abstractmethod
    def find_related_artifacts(
        self,
        artifact_ids: AbcCollection[int],
        target_category: ArtifactCategory,
        relation_type: RelationType = RelationType.RELATES_TO,
    ) -> MultipleArtifactInfo:
        """
        Find artifacts via relations.

        For each of the artifacts in ``artifact_ids``, find all artifacts of
        ``target_category`` that are related to the initial artifacts with a
        relation of type ``relation_type``.

        :param artifact_ids: The IDs of artifacts to start from.
        :param target_category: The category of artifacts to find.
        :param relation_type: The type of relation to follow.
        :return: Information about each artifact.
        """

    @overload
    @abstractmethod
    def lookup_single_collection(
        self,
        lookup: LookupSingle,
        default_category: CollectionCategory | None = None,
    ) -> CollectionInfo: ...

    @overload
    @abstractmethod
    def lookup_single_collection(
        self,
        lookup: None,
        default_category: CollectionCategory | None = None,
    ) -> None: ...

    @abstractmethod
    def lookup_single_collection(
        self,
        lookup: LookupSingle | None,
        default_category: CollectionCategory | None = None,
    ) -> CollectionInfo | None:
        """
        Look up a single collection.

        :param lookup: A :ref:`lookup-single`.
        :param default_category: If the first segment of a string lookup
          (which normally identifies a collection) does not specify a
          category, use this as the default category.
        :return: Information about the collection, or None if the provided
          lookup is None (for convenience in some call sites).
        :raises KeyError: if the lookup does not resolve to an item.
        :raises LookupError: if the lookup is invalid in some way, or does
          not resolve to a collection.
        """

    def lookup_singleton_collection(
        self, category: CollectionCategory
    ) -> CollectionInfo | None:
        """Look up a singleton collection for `category`, if it exists."""
        try:
            return self.lookup_single_collection(f"_@{category}")
        except KeyError:
            return None

    @abstractmethod
    def get_server_setting(self, setting: str) -> str:
        """Look up a Django setting (strings only)."""

    @abstractmethod
    def export_suite_signing_keys(self, suite_id: int) -> str | None:
        """Export signing keys for a suite as an ASCII-armored keyring."""

    def resolve_inputs(self, task: "BaseTask[Any, Any]") -> None:
        """
        Resolve inputs for this task.

        This will look up information required by the input using this task
        database, and make it available to task methods.
        """
        from debusine.tasks.inputs import TaskInput

        TaskInput.resolve_inputs(task, self)

    def get_environment(
        self,
        lookup: LookupSingle,
        *,
        architecture: str | None = None,
        backend: BackendType | None = None,
        default_category: CollectionCategory | None = None,
        image_category: "ExecutorImageCategory | None" = None,
        try_variant: str | None = None,
    ) -> ArtifactInfo:
        """
        Get an environment.

        This automatically fills in some additional constraints if needed.

        :param lookup: the base lookup provided by the task data
        :param architecture: the task's build architecture, if available
        :param backend: the task's backend, or None if the environment lookup
          does not need to be constrained to a particular backend
        :param default_category: the default category to use for the first
          segment of the lookup
        :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 try_variant: None if the environment lookup does not need to be
          constrained to a particular variant; otherwise, and if `lookup` does
          not already specify a variant, then try looking up an environment with
          this variant first, and fall back to looking for an environment with
          no variant
        :return: the ArtifactInfo of a suitable environment artifact
        """
        # Prevent import loop
        from debusine.tasks.executors import (
            ExecutorImageCategory,
            executor_class,
        )

        lookups: list[LookupSingle] = []

        if (
            isinstance(lookup, str)
            and len(segments := parse_lookup_string_segments(lookup)) == 2
            and (parsed := parse_key_value_lookup_segment(segments[1]))[0]
            == "match"
        ):
            # Supplement the environment lookup with the task architecture,
            # if required.
            lookup_type, filters = parsed
            assert lookup_type == "match"

            # TODO: If filters already have values for architecture/format, then
            # check their consistency.

            if architecture is not None:
                filters.setdefault("architecture", architecture)

            if image_category is None and backend is not None:
                image_category = executor_class(backend).image_category
            match image_category:
                case ExecutorImageCategory.TARBALL:
                    filters.setdefault("format", "tarball")
                case ExecutorImageCategory.IMAGE:
                    filters.setdefault("format", "image")
                case _ as unreachable:
                    raise AssertionError(
                        f"Unexpected image category: {unreachable}"
                    )

            if backend is not None:
                filters.setdefault("backend", backend)

            if try_variant is not None and "variant" not in filters:
                for variant in (try_variant, ""):
                    lookups.append(
                        build_lookup_string_segments(
                            segments[0],
                            build_key_value_lookup_segment(
                                lookup_type, {**filters, "variant": variant}
                            ),
                        )
                    )
            else:
                lookups.append(
                    build_lookup_string_segments(
                        segments[0],
                        build_key_value_lookup_segment(lookup_type, filters),
                    )
                )
        else:
            lookups.append(lookup)

        for try_lookup in lookups[:-1]:
            try:
                return self.lookup_single_artifact(
                    try_lookup, default_category=default_category
                )
            except KeyError:
                pass

        return self.lookup_single_artifact(
            lookups[-1], default_category=default_category
        )
