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

"""Utility functions for workflows."""

import functools
from collections.abc import Collection as AbcCollection
from collections.abc import Sequence
from operator import attrgetter
from typing import Any, TYPE_CHECKING, cast, get_args

from django.db.models import QuerySet

from debusine.artifacts import SourcePackage
from debusine.artifacts.models import (
    ArtifactCategory,
    ArtifactDataWithArch,
    BareDataCategory,
    CollectionCategory,
    DebianBinaryPackage,
    DebianPackageBuildLog,
    DebianSourcePackage,
    DebianSystemTarball,
    DebianUpload,
    DebusinePromise,
    get_architectures,
    get_source_package_name,
)
from debusine.client.models import LookupChildType, LookupResultType
from debusine.db.models import Artifact, ArtifactRelation
from debusine.db.models.lookups import LookupResult
from debusine.tasks import TaskConfigError, ensure_artifact_categories
from debusine.tasks.models import (
    BackendType,
    ExtraExternalRepository,
    ExtraRepository,
    LookupMultiple,
    LookupSingle,
)

if TYPE_CHECKING:
    from debusine.server.workflows import Workflow


@functools.lru_cache(maxsize=100)
def source_package(
    workflow: "Workflow[Any, Any]",
    *,
    configuration_key: str = "source_artifact",
) -> Artifact:
    """
    Retrieve the source package artifact from ``configuration_key``.

    If the source artifact is a :artifact:`debian:upload`, returns its
    :artifact:`debian:source-package`.
    """
    artifact = workflow.work_request.lookup_single(
        attrgetter(configuration_key)(workflow.data),
        expect_type=LookupChildType.ARTIFACT,
    ).artifact
    return locate_debian_source_package(configuration_key, artifact)


@functools.lru_cache(maxsize=100)
def source_package_data(
    workflow: "Workflow[Any, Any]",
    *,
    configuration_key: str = "source_artifact",
) -> DebianSourcePackage:
    """Return source package artifact data for the workflow."""
    return SourcePackage.create_data(
        source_package(workflow, configuration_key=configuration_key).data
    )


def lookup_result_artifact_category(result: LookupResult) -> str:
    """
    Get artifact category from result of looking up an artifact.

    The result may be either an artifact or a promise.
    """
    if (
        result.result_type == LookupResultType.ARTIFACT
        and result.artifact is not None
    ):
        return result.artifact.category
    elif (
        result.result_type == LookupResultType.BARE
        and result.collection_item is not None
        and result.collection_item.category == BareDataCategory.PROMISE
    ):
        return DebusinePromise(**result.collection_item.data).promise_category
    else:
        raise ValueError(
            f"Cannot determine artifact category for lookup result: {result}"
        )


def lookup_result_architectures(result: LookupResult) -> set[str]:
    """Get binary package architectures present within LookupResult."""
    if result.artifact is not None:
        artifact_data = result.artifact.create_data()
        if not isinstance(artifact_data, get_args(ArtifactDataWithArch)):
            raise ValueError(f"Unexpected type: {type(artifact_data).__name__}")
        return get_architectures(artifact_data)
    elif result.collection_item is not None:
        architecture = result.collection_item.data.get("architecture")
        if type(architecture) is not str:
            raise ValueError(
                f"Cannot determine architecture for lookup result: {result}"
            )
        return {architecture}
    else:
        raise ValueError(
            "Unexpected result: must have collection_item or artifact"
        )


class ArtifactHasNoBinaryPackageName(Exception):
    """Raised if it's not possible to determine the artifact's binary name."""


def lookup_result_binary_package_name(result: LookupResult) -> str:
    """Get binary package name from result of looking up an artifact."""
    binary_package_name: str | None

    if result.artifact is not None:
        artifact_data = result.artifact.create_data()
        match artifact_data:
            case DebianBinaryPackage():
                binary_package_name = artifact_data.deb_fields.get("Package")
            case _:
                raise ArtifactHasNoBinaryPackageName(f"{type(artifact_data)}")
    elif result.collection_item is not None:
        binary_package_name = result.collection_item.data.get(
            "binary_package_name"
        )
    else:
        raise ValueError(
            "Unexpected result: must have collection_item or artifact"
        )

    if not isinstance(binary_package_name, str):
        raise ValueError(
            f"Cannot determine binary package name for lookup result: {result}"
        )

    return binary_package_name


def filter_artifact_lookup_by_arch(
    workflow: "Workflow[Any, Any]",
    lookup: LookupMultiple,
    architectures: set[str],
) -> LookupMultiple:
    """Filter an artifact lookup by architecture."""
    results = workflow.work_request.lookup_multiple(
        lookup, expect_type=LookupChildType.ARTIFACT_OR_PROMISE
    )
    relevant: list[LookupSingle] = []
    for result in results:
        arch_in_lookup = lookup_result_architectures(result)
        if arch_in_lookup.issubset(architectures):
            relevant.append(
                result.reconstruct(
                    workflow_root=workflow.work_request.workflow_root
                )
            )
        elif arch_in_lookup.intersection(architectures):
            # Uploads are the only artifacts with multiple architectures at
            # this time.
            assert result.artifact is not None
            assert result.artifact.category == ArtifactCategory.UPLOAD
            for binary_package in follow_artifact_relations(
                result.artifact,
                ArtifactRelation.Relations.EXTENDS,
                ArtifactCategory.BINARY_PACKAGE,
            ):
                pkg_data = cast(
                    DebianBinaryPackage, binary_package.create_data()
                )
                pkg_architectures = get_architectures(pkg_data)
                if pkg_architectures.issubset(architectures):
                    relevant.append(f"{binary_package.id}@artifacts")

    return LookupMultiple.model_validate(sorted(relevant, key=str))


def get_lookup_architectures(
    workflow: "Workflow[Any, Any]", lookup: LookupMultiple
) -> set[str]:
    """
    Return set with all the architectures in the artifacts from the lookup.

    The architectures are extracted from each lookup result using
    :py:func:`lookup_result_architectures`.
    """
    results = workflow.work_request.lookup_multiple(
        lookup, expect_type=LookupChildType.ARTIFACT_OR_PROMISE
    )

    architectures: set[str] = set()
    for result in results:
        architectures |= lookup_result_architectures(result)
    return architectures


def follow_artifact_relations(
    artifact: Artifact,
    relation_type: ArtifactRelation.Relations,
    category: ArtifactCategory,
) -> QuerySet[Artifact]:
    """Follow relations from artifact to find artifacts of category."""
    return Artifact.objects.filter(
        targeted_by__artifact=artifact,
        targeted_by__type=relation_type,
        category=category,
    ).order_by("id")


def locate_debian_source_package(
    configuration_key: str, artifact: Artifact
) -> Artifact:
    """
    Find a :artifact:`debian:source-package` artifact for a workflow.

    :param configuration_key: The key in the workflow's task data from which
      the given artifact was looked up.
    :param artifact: A :artifact:`debian:source-package` or
      :artifact:`debian:upload` artifact.
    :return: If ``artifact`` is a
      :artifact:`debian:source-package`, return it; if it is a
      :artifact:`debian:upload`, return the related
      :artifact:`debian:source-package`.
    """
    ensure_artifact_categories(
        configuration_key=configuration_key,
        category=artifact.category,
        expected=[ArtifactCategory.SOURCE_PACKAGE, ArtifactCategory.UPLOAD],
    )
    match artifact.category:
        case ArtifactCategory.SOURCE_PACKAGE:
            return artifact
        case ArtifactCategory.UPLOAD:
            source_packages = follow_artifact_relations(
                artifact,
                ArtifactRelation.Relations.EXTENDS,
                ArtifactCategory.SOURCE_PACKAGE,
            )
            try:
                return source_packages.get()
            except Artifact.DoesNotExist:
                raise TaskConfigError(
                    f"Unable to find an artifact of category "
                    f"{ArtifactCategory.SOURCE_PACKAGE} with a relationship "
                    f"of type {ArtifactRelation.Relations.EXTENDS} from "
                    f'"{artifact}"'
                )
            except Artifact.MultipleObjectsReturned:
                raise TaskConfigError(
                    f"Multiple artifacts of category "
                    f"{ArtifactCategory.SOURCE_PACKAGE} with a relationship "
                    f"of type {ArtifactRelation.Relations.EXTENDS} from "
                    f'"{artifact}" found'
                )
        case _ as unreachable:  # pragma: no cover
            raise AssertionError(f"Unexpected artifact category: {unreachable}")


def locate_debian_source_package_lookup(
    workflow: "Workflow[Any, Any]", configuration_key: str, lookup: LookupSingle
) -> LookupSingle:
    """
    Return a lookup to a :artifact:`debian:source-package`.

    If the specified lookup returns a :artifact:`debian:source-package`,
    return it.  If it returns a :artifact:`debian:upload`, find the related
    :artifact:`debian:source-package` and return a lookup to it.
    """
    artifact = workflow.work_request.lookup_single(
        lookup, expect_type=LookupChildType.ARTIFACT
    ).artifact
    if artifact.category == ArtifactCategory.UPLOAD:
        source_package = locate_debian_source_package(
            configuration_key, artifact
        )
        return f"{source_package.id}@artifacts"
    ensure_artifact_categories(
        configuration_key=configuration_key,
        category=artifact.category,
        expected=[ArtifactCategory.SOURCE_PACKAGE],
    )
    return lookup


def locate_debian_binary_packages(
    configuration_key: str, artifacts: Sequence[Artifact]
) -> list[Artifact]:
    """
    Find :artifact:`debian:binary-package` artifacts for a workflow.

    :param configuration_key: The key in the workflow's task data from which
      the given artifacts were looked up.
    :param artifacts: A sequence of :artifact:`debian:binary-package` or
      :artifact:`debian:upload` artifacts.
    :return: A list of artifacts: for each element of ``artifacts``, if it
      is a :artifact:`debian:binary-package`, return it; if it is a
      :artifact:`debian:upload`, return all the related
      :artifact:`debian:binary-package` artifacts.
    """
    binary_packages: list[Artifact] = []
    for i, artifact in enumerate(artifacts):
        ensure_artifact_categories(
            configuration_key=f"{configuration_key}[{i}]",
            category=artifact.category,
            expected=[ArtifactCategory.BINARY_PACKAGE, ArtifactCategory.UPLOAD],
        )
        match artifact.category:
            case ArtifactCategory.BINARY_PACKAGE:
                binary_packages.append(artifact)
            case ArtifactCategory.UPLOAD:
                binary_packages += follow_artifact_relations(
                    artifact,
                    ArtifactRelation.Relations.EXTENDS,
                    ArtifactCategory.BINARY_PACKAGE,
                )
            case _ as unreachable:  # pragma: no cover
                raise AssertionError(
                    f"Unexpected artifact category: {unreachable}"
                )
    return binary_packages


def get_source_package_names(
    results: Sequence[LookupResult],
    *,
    configuration_key: str,
    artifact_expected_categories: AbcCollection[ArtifactCategory],
) -> list[str]:
    """
    Return a sorted list of source package names from results.

    It ensures that:

    - The :py:class:`LookupResult` objects contain either an artifact or
      promise.
    - Artifacts belong to the artifact_expected_categories.
    - If :py:class:`LookupResult` is a promise: extracts the name from the
      promise data ``source_package_name``.

    :param results: A sequence of :py:class:`LookupResult` objects
      representing artifacts to be processed. Each entry is expected to be
      either an artifact or a promise.
    :param configuration_key: A string used by
      :py:meth:`ensure_artifact_categories` for the exception
      message.
    :param artifact_expected_categories: valid :py:class:`ArtifactCategory`
      that artifacts must belong to.
    :return: A sorted list of source package names.
    """
    source_package_names = set()

    for result in results:
        # lookup_multiple expect_type: only artifacts or promises
        match result.result_type:
            case LookupResultType.ARTIFACT:
                assert result.artifact is not None
                category = result.artifact.category

                ensure_artifact_categories(
                    configuration_key=configuration_key,
                    category=category,
                    expected=artifact_expected_categories,
                )
                artifact_data = result.artifact.create_data()

                assert isinstance(
                    artifact_data,
                    (
                        DebianSourcePackage,
                        DebianUpload,
                        DebianBinaryPackage,
                        DebianPackageBuildLog,
                    ),
                )
                source_package_names.add(get_source_package_name(artifact_data))

            case _:
                # Makes coverage happy
                # It's a promise.
                assert result.result_type == LookupResultType.BARE
                assert result.collection_item is not None

                ensure_artifact_categories(
                    configuration_key=configuration_key,
                    category=result.collection_item.data["promise_category"],
                    expected=artifact_expected_categories,
                )

                if (
                    package_name := result.collection_item.data.get(
                        "source_package_name"
                    )
                ) is not None:
                    source_package_names.add(package_name)

    return sorted(source_package_names)


def get_available_architectures(
    workflow: "Workflow[Any, Any]", *, vendor: str, codename: str
) -> set[str]:
    """Get architectures available for use with this vendor/codename."""
    architectures = set()

    for result in workflow.work_request.lookup_multiple(
        LookupMultiple.model_validate(
            {"collection": vendor, "data__codename": codename}
        ),
        default_category=CollectionCategory.ENVIRONMENTS,
        expect_type=LookupChildType.ARTIFACT,
    ):
        architectures.add(result.artifact.data.get("architecture"))

    if not architectures:
        raise TaskConfigError(
            f"Unable to find any environments for {vendor}:{codename}"
        )

    architectures.add("all")

    return architectures


def configure_for_overlay_suite(
    workflow: "Workflow[Any, Any]",
    *,
    extra_repositories: Sequence[ExtraRepository] | None,
    vendor: str,
    codename: str,
    environment: LookupSingle,
    backend: BackendType,
    architecture: str,
    try_variant: str,
) -> list[ExtraRepository] | None:
    """Return any needed extra repository to use an overlay suite."""
    match (vendor, codename):
        case ("debian", "experimental"):
            pass
        case _:
            return (
                None if extra_repositories is None else list(extra_repositories)
            )

    if extra_repositories is None:
        extra_repositories = []

    environment_artifact = workflow.work_request.lookup_environment(
        environment,
        architecture=architecture,
        backend=backend,
        default_category=CollectionCategory.ENVIRONMENTS,
        try_variant=try_variant,
    ).artifact.create_data()
    assert isinstance(environment_artifact, DebianSystemTarball)
    mirror = environment_artifact.mirror
    components = environment_artifact.components

    return [
        *extra_repositories,
        ExtraExternalRepository(
            url=mirror, suite=codename, components=components
        ),
    ]
