# 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 server tasks."""

import abc
from typing import Any, override

import debusine.tasks.tags as ttags
import debusine.worker.tags as wtags
from debusine.db.models import WorkRequest
from debusine.db.tasks import DBTask
from debusine.db.tasks.mock_resolver import MockFieldResolver
from debusine.server.tasks.base import BaseServerTask
from debusine.server.tasks.wait import BaseWaitTask
from debusine.tasks import inputs
from debusine.tasks.models import BaseDynamicTaskData, BaseTaskData
from debusine.tasks.tests.helpers import CurrentTaskAccessor, TaskTestCase
from debusine.test.django import TestCase


class DBTaskTestCase[TaskClass: DBTask[Any, Any]](
    TaskTestCase[TaskClass], TestCase, abc.ABC
):
    """Common base for DBTask tests."""

    #: Field resolver for the test
    resolver: MockFieldResolver

    #: Shortcut to self.task.work_request
    work_request = CurrentTaskAccessor[WorkRequest]()

    @override
    def setUp(self) -> None:
        super().setUp()
        self.resolver = MockFieldResolver()

    @abc.abstractmethod
    def work_request_create(self, **kwargs: Any) -> WorkRequest:
        """Set up the WorkRequest for the current task."""

    @override
    def task_create(self, **kwargs: Any) -> TaskClass:
        work_request = self.work_request_create(**kwargs)
        task = work_request.get_task()
        assert isinstance(task, self.task_class)
        return task

    @override
    def task_set_current(self, task: TaskClass) -> None:
        super().task_set_current(task)
        self.__dict__["work_request"] = task.work_request

    @override
    def task_teardown(self, task: TaskClass) -> None:
        if task._debug_log_files_directory:
            task._debug_log_files_directory.cleanup()
        super().task_teardown(task)

    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()

    # TODO: this currently cannot use self.task because it cannot be typed to
    # return the task's dynamic data type. Passing self.task is shorter to type
    # than casting or passing the dynamic data type
    def compute_dynamic_data[DTD: BaseDynamicTaskData](
        self,
        task: DBTask[Any, DTD] | None = None,
        strip_task_inputs: bool = False,
    ) -> DTD:
        """Compute dynamic data for this task."""
        task = task or self.task
        task.resolve_inputs(inputs.Stage.RUNNING)
        ddata = task.compute_dynamic_data()
        if strip_task_inputs:
            ddata.task_input_fields = {}
        return ddata

    @override
    def build_expected_tag_candidates(
        self, **kwargs: set[str]
    ) -> tuple[ttags.TaskProvidedTags, wtags.TaskRequiredTags]:
        workspace = self.task.workspace
        provided, required = super().build_expected_tag_candidates(**kwargs)
        provided.add(
            ttags.ProvenanceProvided.SYSTEM,
            {
                ttags.SCOPE_PREFIX + workspace.scope.name,
                ttags.WORKSPACE_PREFIX
                + f"{workspace.scope.name}:{workspace.name}",
            },
        )
        return provided, required


class ServerTaskTestCase[TaskClass: BaseServerTask[Any, Any]](
    DBTaskTestCase[TaskClass], TestCase
):
    """Common base for server task tests."""

    #: Set to True for the task to be marked running
    mark_running: bool

    @override
    def setUp(self) -> None:
        super().setUp()
        self.mark_running = True

    @override
    def work_request_create(self, **kwargs: Any) -> WorkRequest:
        task_data = self.task_data.copy()
        task_data.update(kwargs)
        work_request = self.playground.create_server_task(
            task_name=self.task_class.name,
            task_data=self.task_class.task_data_type(**task_data),
        )
        self.playground.advance_work_request(
            work_request, mark_running=self.mark_running
        )
        return work_request

    @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_CELERY,
                wtags.TASK_PREFIX
                + f"server:{self.task.name}:version:{self.task.TASK_VERSION}",
            },
        )
        return provided, required


class SampleBaseDBTask[TD: BaseTaskData, DTD: BaseDynamicTaskData](
    DBTask[TD, DTD]
):
    """Common test implementation of abstract task methods."""

    @override
    def compute_dynamic_data(self) -> DTD:
        """Resolve artifact lookups for this task."""
        return self.dynamic_task_data_type()


class SampleBaseServerTask[TD: BaseTaskData, DTD: BaseDynamicTaskData](
    SampleBaseDBTask[TD, DTD], BaseServerTask[TD, DTD]
):
    """Common test implementation of BaseServerTask methods."""


class SampleBaseWaitTask[TD: BaseTaskData, DTD: BaseDynamicTaskData](
    SampleBaseDBTask[TD, DTD], BaseWaitTask[TD, DTD]
):
    """Common test implementation of BaseWaitTask methods."""
