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

"""Common code for workflows that handle regression tracking."""

import datetime as dt
from abc import ABCMeta
from functools import cached_property
from typing import TypeAlias, override

from django.db.models import Max, Window
from django.db.models.fields.json import KT
from django.db.models.functions import Greatest, Rank

from debusine.artifacts.models import ArtifactCategory, TaskTypes
from debusine.client.models import LookupChildType
from debusine.db.models import Artifact, Collection, CollectionItem, WorkRequest
from debusine.server.workflows import workflow_utils
from debusine.server.workflows.base import Workflow, WorkflowValidationError
from debusine.server.workflows.models import RegressionTrackingWorkflowData
from debusine.tasks.models import (
    ActionSkipIfLookupResultChanged,
    BaseDynamicTaskData,
    RegressionAnalysisStatus,
)


class RegressionTrackingWorkflow[
    WD: RegressionTrackingWorkflowData,
    DTD: BaseDynamicTaskData,
](Workflow[WD, DTD], metaclass=ABCMeta):
    """Common code for a workflow that handles regression tracking."""

    @cached_property
    def qa_suite(self) -> Collection | None:
        """Suite that we are tracking regressions against, if configured."""
        if self.data.qa_suite is None:
            return None
        return self.work_request.lookup_single(
            self.data.qa_suite, expect_type=LookupChildType.COLLECTION
        ).collection

    @cached_property
    def qa_suite_changed(self) -> dt.datetime | None:
        """The latest time the ``qa_suite`` collection was changed."""
        # This method is only called when update_qa_results is True, in
        # which case this is checked by a model validator.
        assert self.qa_suite is not None

        changed = self.qa_suite.child_items.aggregate(
            changed=Greatest(Max("created_at"), Max("removed_at"))
        )["changed"]
        assert isinstance(changed, (dt.datetime, type(None)))
        return changed

    @cached_property
    def reference_qa_results(self) -> Collection | None:
        """Collection of reference results of QA tasks, if configured."""
        if self.data.reference_qa_results is None:
            return None
        return self.work_request.lookup_single(
            self.data.reference_qa_results,
            expect_type=LookupChildType.COLLECTION,
        ).collection

    @override
    def validate_input(self) -> None:
        """Thorough validation of input data."""
        # Validate that we can look up self.data.qa_suite and
        # self.data.reference_qa_results.
        try:
            self.qa_suite
            self.reference_qa_results
        except LookupError as e:
            raise WorkflowValidationError(str(e)) from e

    def skip_if_qa_result_changed(
        self,
        work_request: WorkRequest,
        *,
        package: str,
        architecture: str,
        promise_name: str | None = None,
    ) -> None:
        """Skip this work request if another workflow won the race."""
        lookup = (
            f"{self.data.reference_qa_results}/"
            f"latest:{work_request.task_name}_{package}_{architecture}"
        )
        try:
            item = self.work_request.lookup_single(lookup).collection_item
        except KeyError:
            item = None
        work_request.add_event_reaction(
            "on_assignment",
            ActionSkipIfLookupResultChanged(
                lookup=lookup,
                collection_item_id=None if item is None else item.id,
                promise_name=promise_name,
            ),
        )

    def find_reference_artifacts(self, task_name: str) -> dict[str, Artifact]:
        """Find the latest reference result for each architecture."""
        # This should be called from a workflow callback that is only
        # created when enable_regression_tracking is true, in which case
        # this is checked by a model validator.
        assert self.reference_qa_results is not None

        source_data = workflow_utils.source_package_data(self)
        reference_artifacts: dict[str, Artifact] = {}
        for result in (
            self.reference_qa_results.child_items.filter(
                child_type=CollectionItem.Types.ARTIFACT,
                artifact__isnull=False,
                data__task_name=task_name,
                data__package=source_data.name,
                data__has_key="architecture",
            )
            .annotate(
                rank_by_architecture=Window(
                    Rank(),
                    partition_by=KT("data__architecture"),
                    order_by="-data__timestamp",
                )
            )
            .filter(rank_by_architecture=1)
            .select_related("artifact__created_by_work_request")
        ):
            assert result.artifact is not None
            reference_artifacts[result.data["architecture"]] = result.artifact
        return reference_artifacts

    def find_new_artifacts(
        self, task_name: str, category: ArtifactCategory
    ) -> dict[str, Artifact]:
        """Find the new result for each architecture."""
        return {
            artifact.data["architecture"]: artifact
            for artifact in Artifact.objects.filter(
                created_by_work_request__in=(
                    self.work_request.children.unsuperseded()
                    .terminated()
                    .filter(task_type=TaskTypes.WORKER, task_name=task_name)
                ),
                category=category,
                data__has_key="architecture",
            ).select_related("created_by_work_request")
        }

    @staticmethod
    def compare_qa_results(
        reference: WorkRequest | None, new: WorkRequest | None
    ) -> RegressionAnalysisStatus:
        """
        Compare work requests providing two QA results.

        Return a status depending on the two work requests:

        * ``no-result``: when the comparison has not been completed yet
          (usually because we lack one of the two required QA results)
        * ``error``: when the comparison (or one of the required QA tasks)
          errored out
        * ``improvement``: when the new QA result is better than the
          original QA result
        * ``stable``: when the new QA result is neither better nor worse
          than the original QA result
        * ``regression``: when the new QA result is worse than the original
          QA result

        If the status is ``stable``, then the caller may perform a
        finer-grained analysis.
        """
        # Some aliases to make the match statement more readable.
        WR: TypeAlias = WorkRequest
        S: TypeAlias = WorkRequest.Statuses
        R: TypeAlias = WorkRequest.Results

        match (reference, new):
            case (None, _) | (_, None):
                return RegressionAnalysisStatus.NO_RESULT
            case WR(status=S.COMPLETED, result=R.SUCCESS | R.SKIPPED), WR(
                status=S.COMPLETED, result=R.FAILURE | R.ERROR
            ):
                return RegressionAnalysisStatus.REGRESSION
            case WR(status=S.COMPLETED, result=R.FAILURE | R.ERROR), WR(
                status=S.COMPLETED, result=R.SUCCESS | R.SKIPPED
            ):
                return RegressionAnalysisStatus.IMPROVEMENT
            case WR(status=S.COMPLETED, result=R.SUCCESS | R.SKIPPED), WR(
                status=S.COMPLETED, result=R.SUCCESS | R.SKIPPED
            ):
                return RegressionAnalysisStatus.STABLE
            case WR(status=S.COMPLETED, result=reference_result), WR(
                status=S.COMPLETED, result=new_result
            ) if reference_result == new_result:
                return RegressionAnalysisStatus.STABLE
            case _:
                return RegressionAnalysisStatus.ERROR
