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

"""debdiff workflow."""

from collections import defaultdict
from collections.abc import Sequence
from typing import NewType, override

from debian.debian_support import Version

from debusine.artifacts import SourcePackage
from debusine.artifacts.models import (
    ArtifactCategory,
    DebianBinaryPackage,
    DebianSourcePackage,
    get_source_package_name,
)
from debusine.client.models import LookupChildType
from debusine.db.models import Artifact, ArtifactRelation, WorkRequest
from debusine.db.models.lookups import LookupResult
from debusine.server.collections import DebianSuiteManager
from debusine.server.collections.lookup import artifact_ids
from debusine.server.workflows import Workflow, workflow_utils
from debusine.server.workflows.models import (
    DebDiffWorkflowData,
    DebDiffWorkflowDynamicData,
    WorkRequestWorkflowData,
)
from debusine.server.workflows.workflow_utils import (
    lookup_result_architectures,
    lookup_result_binary_package_name,
    lookup_result_binary_package_version,
)
from debusine.tasks.models import (
    DebDiffData,
    DebDiffFlags,
    DebDiffInput,
    InputArtifact,
    InputArtifactMultiple,
    InputArtifactSingle,
    LookupMultiple,
    LookupSingle,
)

PackageName = NewType("PackageName", str)
Architecture = NewType("Architecture", str)
PackageArchLookupMap = dict[PackageName, dict[Architecture, LookupSingle]]


class DebDiffWorkflow(
    Workflow[DebDiffWorkflowData, DebDiffWorkflowDynamicData]
):
    """DebDiff workflow."""

    TASK_NAME = "debdiff"

    @override
    def compute_dynamic_data(self) -> DebDiffWorkflowDynamicData:
        """
        Compute dynamic data for this workflow.

        :subject: package name of ``source_artifact``
        :parameter_summary: name of debian suite: old_version → new_version
        :source_artifact_id: id of the source artifact
        :binary_artifacts_ids: ids of the binary artifacts
        :original_id: id of the original artifact
        """
        source_artifact = workflow_utils.source_package(self)
        new_source_data = SourcePackage.create_data(source_artifact.data)
        new_source_package_name = get_source_package_name(new_source_data)
        version_new = new_source_data.version

        debian_suite = self.work_request.lookup_single(
            self.data.original, expect_type=LookupChildType.COLLECTION
        ).collection

        suite_manager = debian_suite.manager
        assert isinstance(suite_manager, DebianSuiteManager)

        original_source_package = suite_manager.lookup_newest_package_recursive(
            f"source:{new_source_package_name}"
        )

        version_original = None
        original_source_package_id = None
        if original_source_package is not None:
            assert original_source_package.artifact is not None
            version_original = Version(
                DebianSourcePackage(
                    **original_source_package.artifact.data
                ).version
            )
            original_source_package_id = original_source_package.artifact.id
        else:
            original_binary_packages = (
                self._get_binary_packages_from_source_package(
                    debian_suite=self.data.original,
                    source_package_name=new_source_package_name,
                )
            )

            if original_binary_packages:
                for original_binary_package in original_binary_packages:
                    assert original_binary_package.artifact is not None
                    version = Version(
                        DebianBinaryPackage(
                            **original_binary_package.artifact.data
                        ).srcpkg_version
                    )
                    if version_original is None or version > version_original:
                        version_original = version

        if version_original is not None:
            parameter_summary = (
                f"{debian_suite.name}: {version_original} → {version_new}"
            )
        else:
            parameter_summary = f"{debian_suite.name}: missing → {version_new}"

        binary_artifacts_results = self.work_request.lookup_multiple(
            self.data.binary_artifacts,
            expect_type=LookupChildType.ARTIFACT_OR_PROMISE,
        )

        return DebDiffWorkflowDynamicData(
            subject=new_source_package_name,
            parameter_summary=parameter_summary,
            source_artifact_id=source_artifact.id,
            binary_artifacts_ids=artifact_ids(binary_artifacts_results),
            original_id=original_source_package_id,
        )

    @override
    def get_input_artifacts(self) -> list[InputArtifact]:
        """Return input artifacts."""
        return [
            InputArtifactSingle(
                lookup=self.data.source_artifact,
                label="source_artifact",
                artifact_id=(
                    self.dynamic_data.source_artifact_id
                    if self.dynamic_data is not None
                    else None
                ),
            ),
            InputArtifactSingle(
                lookup=self.data.original,
                label="original",
                artifact_id=(
                    self.dynamic_data.original_id
                    if self.dynamic_data is not None
                    else None
                ),
            ),
            InputArtifactMultiple.from_ids(
                lookup=self.data.binary_artifacts,
                label="binary_artifacts",
                artifact_ids=(
                    self.dynamic_data.binary_artifacts_ids
                    if self.dynamic_data is not None
                    else None
                ),
            ),
        ]

    @override
    def populate(self) -> None:
        """Create work requests."""
        new_source_data = workflow_utils.source_package_data(self)

        environment = f"{self.data.vendor}/match:codename={self.data.codename}"

        # Populate DebDiff tasks for the source package
        debian_suite = self.work_request.lookup_single(
            self.data.original, expect_type=LookupChildType.COLLECTION
        ).collection

        suite_manager = debian_suite.manager
        assert isinstance(suite_manager, DebianSuiteManager)

        original_source = suite_manager.lookup_newest_package_recursive(
            f"source:{new_source_data.name}"
        )
        original_source_version: str | None = None
        if original_source is not None:
            original_source_version = original_source.data["version"]
            original_source_lookup = (
                f"{self.data.original}/name:{original_source.name}"
            )
            self._populate_source_artifact_debdiff(
                original=original_source_lookup,
                new=self.data.source_artifact,
                build_architecture=self.data.arch_all_build_architecture,
                environment=environment,
                extra_flags=self.data.extra_flags,
            )

        # Populate DebDiff tasks for the binary packages
        original_binaries_arch_lookup = (
            self._get_original_packages_by_name_version_arch(
                new_source_data.name, original_source_version
            )
        )
        archs_old: set[Architecture] = set()
        for archs_lookup in original_binaries_arch_lookup.values():
            archs_old.update(archs_lookup.keys())

        concrete_archs_old = archs_old - {Architecture("all")}

        new_binaries_arch_lookup = self._get_new_packages_by_name_arch()

        archs_new = self._get_architectures(new_binaries_arch_lookup)
        concrete_archs_new = archs_new - {Architecture("all")}

        if concrete_archs_new:
            # New binaries provide concrete archs: compare per concrete arch
            archs_to_compare = concrete_archs_new
        elif concrete_archs_old and Architecture("all") not in archs_old:
            # Transition any → all: new provides only "all" (because no
            # concrete_archs_new), old provides only concretes
            archs_to_compare = concrete_archs_old
        else:
            # At least one side provides "all" (and no concrete archs on new)
            archs_to_compare = {Architecture("all")}

        for arch in sorted(archs_to_compare):
            self._populate_binary_for_arch(
                architecture=arch,
                orig_map=original_binaries_arch_lookup,
                new_map=new_binaries_arch_lookup,
                environment=environment,
                extra_flags=self.data.extra_flags,
            )

    def _populate_binary_for_arch(
        self,
        *,
        architecture: Architecture,
        orig_map: PackageArchLookupMap,
        new_map: PackageArchLookupMap,
        environment: LookupSingle,
        extra_flags: list[DebDiffFlags],
    ) -> None:
        arch_all = Architecture("all")
        new_relevants = set()
        for arch_lookup in new_map.values():
            for arch in (architecture, arch_all):
                if arch in arch_lookup:
                    new_relevants.add(arch_lookup[arch])

        original_relevants: list[LookupSingle] = []
        for old_pkg_name, old_arch_map in orig_map.items():
            # Architectures are considered compatible in two cases:
            # - if comparing a concrete architecture: include the architecture
            #   and "all" packages
            # - if comparing "all": include only "all" packages
            original_relevants += [
                old_lookup
                for old_arch, old_lookup in old_arch_map.items()
                if old_pkg_name in new_map
                and (
                    (architecture == arch_all and old_arch == arch_all)
                    or (
                        architecture != arch_all
                        and old_arch in (arch_all, architecture)
                    )
                )
            ]

        if original_relevants and new_relevants:
            build_arch = (
                self.data.arch_all_build_architecture
                if architecture == Architecture("all")
                else architecture
            )

            self._populate_binary_artifact_debdiff(
                originals=LookupMultiple.model_validate(
                    sorted(original_relevants)
                ),
                news=LookupMultiple.model_validate(sorted(new_relevants)),
                architecture=architecture,
                build_architecture=build_arch,
                environment=environment,
                extra_flags=extra_flags,
            )

    def _get_binary_packages_from_source_package(
        self,
        debian_suite: LookupSingle,
        source_package_name: str,
        source_package_version: str | None = None,
    ) -> Sequence[LookupResult]:
        collection = self.work_request.lookup_single(
            debian_suite, expect_type=LookupChildType.COLLECTION
        ).collection
        assert isinstance(collection.manager, DebianSuiteManager)

        newest_binary_packages: dict[
            tuple[PackageName, Architecture], tuple[Version, LookupResult]
        ] = {}

        for required_suite in [
            collection,
            *collection.manager.required_suites(),
        ]:
            lookup_kwargs = {
                "collection": str(required_suite),
                "child_type": LookupChildType.ARTIFACT,
                "data__srcpkg_name": source_package_name,
                "category": ArtifactCategory.BINARY_PACKAGE,
            }

            if source_package_version is not None:
                lookup_kwargs["data__srcpkg_version"] = source_package_version

            for result in self.work_request.lookup_multiple(
                LookupMultiple.model_validate(lookup_kwargs),
                expect_type=LookupChildType.ARTIFACT,
            ):
                binary_package_name = PackageName(
                    lookup_result_binary_package_name(result)
                )

                # DebianBinaryPackages always have one and only one architecture
                binary_package_arch = Architecture(
                    lookup_result_architectures(result).pop()
                )

                binary_package_version = lookup_result_binary_package_version(
                    result
                )

                name_arch = (binary_package_name, binary_package_arch)

                if name_arch not in newest_binary_packages:
                    # First time found
                    newest_binary_packages[name_arch] = (
                        binary_package_version,
                        result,
                    )
                elif (
                    newest_binary_packages[name_arch][0]
                    < binary_package_version
                ):
                    # A newer version was found in another suite.
                    newest_binary_packages[name_arch] = (
                        binary_package_version,
                        result,
                    )

        return [value[1] for value in newest_binary_packages.values()]

    def _get_original_packages_by_name_version_arch(
        self,
        srcpkg_name: str,
        srcpkg_version: str | None,
    ) -> PackageArchLookupMap:
        """
        Return mapping of original binary packages by name and arch.

        Only packages with the highest Version. E.g. if there are the following
        packages: hello/2.0.0, libhello2/2.0.0 and libhello1/1.0.0 the last
        one will not be returned.
        """
        original_binary_results = self._get_binary_packages_from_source_package(
            debian_suite=self.data.original,
            source_package_name=srcpkg_name,
            source_package_version=srcpkg_version,
        )

        packages_by_srcpkg_version: defaultdict[
            Version, list[tuple[PackageName, Architecture, LookupSingle]]
        ] = defaultdict(list)

        for lookup_result in original_binary_results:
            architectures = lookup_result_architectures(lookup_result)
            assert len(architectures) == 1
            arch = Architecture(architectures.pop())
            pkg_name = PackageName(
                lookup_result_binary_package_name(lookup_result)
            )
            assert lookup_result.artifact is not None
            version = Version(
                DebianBinaryPackage(
                    **lookup_result.artifact.data
                ).srcpkg_version
            )

            assert lookup_result.collection_item is not None
            packages_by_srcpkg_version[version].append(
                (pkg_name, arch, lookup_result.reconstruct(self.work_request))
            )

        # Return packages with only the highest version
        highest_srcpkg_version = max(packages_by_srcpkg_version, default=None)
        result: PackageArchLookupMap = defaultdict(dict)
        if highest_srcpkg_version is not None:
            for pkg_name, arch, lookup in packages_by_srcpkg_version[
                highest_srcpkg_version
            ]:
                result[pkg_name][arch] = lookup
        return result

    @staticmethod
    def _extract_package_name_and_architecture(
        artifact: Artifact,
    ) -> tuple[PackageName, Architecture]:
        """Expect artifact of type BINARY_PACKAGE."""
        data = artifact.create_data()
        assert isinstance(data, DebianBinaryPackage)

        architecture = Architecture(data.deb_fields.get("Architecture", ""))

        package_name = PackageName(data.deb_fields.get("Package", ""))

        return package_name, architecture

    @classmethod
    def _package_arch_pairs(
        cls,
        lookup_result: LookupResult,
    ) -> list[tuple[PackageName, Architecture]]:
        """Return pairs of package names and architectures for a lookup."""
        if lookup_result.artifact is None:
            raise ValueError(
                "DebDiffWorkflow requires binary artifacts (not promises)"
            )

        match lookup_result.artifact.category:
            case ArtifactCategory.BINARY_PACKAGE:
                architectures = lookup_result_architectures(lookup_result)
                assert len(architectures) == 1
                architecture = Architecture(architectures.pop())
                pkg_name = PackageName(
                    lookup_result_binary_package_name(lookup_result)
                )
                return [(pkg_name, architecture)]

            case ArtifactCategory.UPLOAD:
                pairs: list[tuple[PackageName, Architecture]] = []
                for relation in lookup_result.artifact.relations.filter(
                    type=ArtifactRelation.Relations.EXTENDS,
                    target__category=ArtifactCategory.BINARY_PACKAGE,
                ).select_related("target"):
                    package_name, architecture = (
                        cls._extract_package_name_and_architecture(
                            relation.target
                        )
                    )
                    pairs.append((package_name, architecture))

                return sorted(pairs)

            # https://github.com/nedbat/coveragepy/issues/1860 (fixed in
            # coverage 7.6.2)
            case _ as artifact_category:  # pragma: no cover
                raise ValueError(f"Unexpected category {artifact_category}")

    def _get_new_packages_by_name_arch(
        self,
    ) -> PackageArchLookupMap:
        """Return mapping of new binary packages by name and arch."""
        new_binary_results = self.work_request.lookup_multiple(
            self.data.binary_artifacts,
            expect_type=LookupChildType.ARTIFACT_OR_PROMISE,
        )
        packages_by_name_arch: defaultdict[
            PackageName, dict[Architecture, LookupSingle]
        ] = defaultdict(dict)

        for result in new_binary_results:
            # DebDiffWorkflow gets populated once the build is ready
            assert result.artifact is not None

            artifact = result.artifact

            match artifact.category:
                case ArtifactCategory.BINARY_PACKAGE:
                    for pkg_name, arch in self._package_arch_pairs(result):
                        packages_by_name_arch[pkg_name][arch] = (
                            result.reconstruct(
                                workflow_root=self.work_request.workflow_root,
                            )
                        )

                case ArtifactCategory.UPLOAD:
                    for relation in artifact.relations.filter(
                        type=ArtifactRelation.Relations.EXTENDS,
                        target__category=ArtifactCategory.BINARY_PACKAGE,
                    ).select_related("target"):
                        pkg_name, arch = (
                            self._extract_package_name_and_architecture(
                                relation.target
                            )
                        )
                        packages_by_name_arch[pkg_name][arch] = (
                            relation.target.id
                        )

                # https://github.com/nedbat/coveragepy/issues/1860 (fixed in
                # coverage 7.6.2)
                case _ as artifact_category:  # pragma: no cover
                    raise ValueError(
                        f"Unexpected category {artifact_category!r}"
                    )

        return packages_by_name_arch

    def _get_architectures(
        self,
        packages_by_name_arch: dict[
            PackageName, dict[Architecture, LookupSingle]
        ],
    ) -> set[Architecture]:
        """Return all architectures present in package mapping."""
        archs: set[Architecture] = set()
        for arch_map in packages_by_name_arch.values():
            archs.update(arch_map.keys())
        return archs

    def _populate_source_artifact_debdiff(
        self,
        *,
        original: LookupSingle,
        new: LookupSingle,
        build_architecture: str,
        environment: LookupSingle,
        extra_flags: list[DebDiffFlags],
    ) -> WorkRequest:
        new = workflow_utils.locate_debian_source_package_lookup(
            self, "source_artifact", new
        )

        wr, _ = self.work_request_ensure_child_worker(
            task_name="debdiff",
            task_data=DebDiffData(
                input=DebDiffInput(
                    source_artifacts=[original, new],
                ),
                environment=environment,
                build_architecture=build_architecture,
                extra_flags=extra_flags,
            ),
            workflow_data=WorkRequestWorkflowData(
                display_name="DebDiff for source package",
                step="debdiff-source",
            ),
        )
        self.requires_artifact(wr, original)
        self.requires_artifact(wr, new)

        return wr

    def _populate_binary_artifact_debdiff(
        self,
        *,
        originals: LookupMultiple,
        news: LookupMultiple,
        architecture: str,
        build_architecture: str,
        environment: LookupSingle,
        extra_flags: list[DebDiffFlags],
    ) -> WorkRequest:
        wr, _ = self.work_request_ensure_child_worker(
            task_name="debdiff",
            task_data=DebDiffData(
                input=DebDiffInput(
                    binary_artifacts=[originals, news],
                ),
                environment=environment,
                build_architecture=build_architecture,
                extra_flags=extra_flags,
            ),
            workflow_data=WorkRequestWorkflowData(
                display_name=f"DebDiff for binary packages ({architecture})",
                step=f"debdiff-binaries-{architecture}",
            ),
        )
        self.requires_artifact(wr, originals)
        self.requires_artifact(wr, news)

        return wr
