# 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 test-helper code involving workflows."""

from collections.abc import Generator
from contextlib import contextmanager
from typing import Any, override
from unittest import mock

from django.db import connection
from django.test import override_settings

import debusine.tasks.tags as ttags
import debusine.worker.tags as wtags
from debusine.artifacts.models import TaskTypes
from debusine.db.models import Artifact, WorkRequest, WorkflowTemplate
from debusine.server.scheduler import schedule
from debusine.server.tasks.tests.helpers import DBTaskTestCase, SampleBaseDBTask
from debusine.server.workflows import SbuildWorkflow, Workflow, workflow_utils
from debusine.server.workflows.models import BaseWorkflowData
from debusine.tasks import inputs
from debusine.tasks.models import BaseDynamicTaskData
from debusine.test.django import TestCase


class SampleWorkflow[WD: BaseWorkflowData, DTD: BaseDynamicTaskData](
    Workflow[WD, DTD], SampleBaseDBTask[WD, DTD]
):
    """Common test implementation of Workflow methods."""


class WorkflowTestCase[TaskClass: Workflow[Any, Any]](
    DBTaskTestCase[TaskClass], TestCase
):
    """Base class for workflow tests."""

    #: Parent WorkRequest for the workflow being tested
    parent: WorkRequest | None = None
    #: Validate data of the workflow being tested
    validate: bool = True
    #: Internal collection artifacts to pass to create_workflow
    internal_collection_artifacts: (
        list[tuple[Artifact, dict[str, Any]]] | None
    ) = None

    @override
    def setUp(self) -> None:
        super().setUp()
        self.template: WorkflowTemplate | None = None

    def create_template(self, name: str, **kwargs: Any) -> WorkflowTemplate:
        """
        Create a workflow template.

        :param kwargs: if any keyword arguments are provided, they replace the
          corresponding values in ``self.task_data``.
        """
        task_data = self.task_data.copy()
        task_data.update(kwargs)
        return self.playground.create_workflow_template(
            name, self.task_class.name, static_parameters=task_data
        )

    @override
    def work_request_create(self, **kwargs: Any) -> WorkRequest:
        task_data = self.task_data.copy()
        task_data.update(kwargs)
        work_request = self.playground.create_workflow(
            task_name=(
                self.template
                if self.template is not None
                else self.task_class.name
            ),
            task_data=self.task_class.task_data_type(**task_data),
            parent=self.parent,
            validate=self.validate,
            internal_collection_artifacts=self.internal_collection_artifacts,
        )
        return work_request

    @classmethod
    def get_workflow(cls, work_request: WorkRequest) -> TaskClass:
        """Create and return a suitable Workflow for this work request."""
        task = work_request.get_task()
        assert isinstance(task, cls.task_class)
        return task

    @contextmanager
    def schedule_and_run(self) -> Generator[list[WorkRequest]]:
        """
        Run the scheduler and any resulting Celery callbacks.

        Yields the list of scheduled work requests so that the caller can
        assert things about them before they are run.
        """
        start_count = len(connection.run_on_commit)
        with (
            mock.patch("pgtransaction.Atomic.execute_set_transaction_modes"),
            override_settings(PGTRANSACTION_RETRY=0),
        ):
            with self.captureOnCommitCallbacks() as on_commit:
                yield schedule()
            with override_settings(CELERY_TASK_ALWAYS_EAGER=True):
                for callback in on_commit:
                    callback()
                del connection.run_on_commit[
                    start_count : start_count + len(on_commit)
                ]

    def schedule_and_run_workflow_callback(self, parent: WorkRequest) -> None:
        """
        Run the scheduler and a workflow callback.

        :param parent: Expect the workflow callback to be a child of this
          workflow.
        """
        with self.schedule_and_run() as [work_request]:
            self.assertEqual(work_request.task_type, TaskTypes.INTERNAL)
            self.assertEqual(work_request.task_name, "workflow")
            self.assertEqual(work_request.parent, parent)

    def schedule_and_run_workflow(self, workflow: WorkRequest) -> None:
        """
        Run the scheduler and a workflow.

        :param workflow: Expect this workflow to be scheduled.
        """
        with self.schedule_and_run() as work_requests:
            self.assertIn(workflow, work_requests)

    def simulate_sbuild_workflow_completion(
        self,
        sbuild_workflow: WorkRequest,
        only_architecture: str | None = None,
        binary_names: list[str] | None = None,
        section: str = "misc",
        priority: str = "optional",
        use_data: bool = True,
    ) -> tuple[dict[str, list[Artifact]], dict[str, Artifact]]:
        """
        Simulate completion of an sbuild workflow.

        :param sbuild_workflow: The workflow to complete.
        :param only_architecture: If given, only complete the sbuild task
          for this architecture.
        :param binary_names: List of binary package names to create;
          defaults to a single binary named after the source package.
        :param section: ``Section`` field for binary artifacts.
        :param priority: ``Priority`` field for binary artifacts.
        :return: A dictionary mapping architectures to the
          ``debian:binary-package`` artifacts created for them, and a
          dictionary mapping architectures to their
          ``debian:package-build-log`` artifacts.
        """
        sbuild_workflow_task = sbuild_workflow.get_task()
        assert isinstance(sbuild_workflow_task, SbuildWorkflow)
        source_data = workflow_utils.source_package_data(
            sbuild_workflow_task, configuration_key="input.source_artifact"
        )

        binaries: dict[str, list[Artifact]] = {}
        build_logs: dict[str, Artifact] = {}
        for sbuild in sbuild_workflow.children.filter(
            task_type=TaskTypes.WORKER,
            task_name="sbuild",
            workflow_data_json__step__startswith="build-",
        ).order_by("workflow_data_json__step"):
            assert sbuild.workflow_data.step is not None
            arch = sbuild.workflow_data.step.removeprefix("build-")
            if only_architecture is not None and arch != only_architecture:
                continue
            upload = self.playground.create_upload_artifacts(
                src_name=source_data.name,
                version=source_data.version,
                binary=True,
                source=False,
                binaries=[
                    (binary_name, arch)
                    for binary_name in (binary_names or (source_data.name,))
                ],
                section=section,
                priority=priority,
                workspace=sbuild.workspace,
                work_request=sbuild,
            )
            binaries[arch] = upload.binaries
            build_logs[arch] = self.playground.create_build_log_artifact(
                source=source_data.name,
                version=source_data.version,
                build_arch=arch,
                workspace=sbuild.workspace,
                work_request=sbuild,
            )
            self.playground.create_artifact_relation(
                artifact=build_logs[arch], target=upload.upload
            )
            self.playground.advance_work_request(
                sbuild, result=WorkRequest.Results.SUCCESS, use_data=use_data
            )
        return binaries, build_logs

    def compute_scheduler_tags(
        self,
    ) -> tuple[ttags.TaskProvidedTags, wtags.TaskRequiredTags]:
        """Invoke compute_scheduler_tags."""
        self.task.resolve_inputs(inputs.Stage.PENDING)
        return self.task.compute_scheduler_tags()

    @override
    def build_expected_tag_candidates(
        self, **kwargs: set[str]
    ) -> tuple[ttags.TaskProvidedTags, wtags.TaskRequiredTags]:
        provided, required = super().build_expected_tag_candidates(**kwargs)
        required.add(
            wtags.ProvenanceRequired.SYSTEM, {wtags.WORKER_TYPE_NOT_ASSIGNABLE}
        )
        return provided, required
