from typing import Any, Dict

from django.contrib.auth.models import AbstractUser

from opentelemetry import trace

from baserow.core.db import transaction_atomic
from baserow.core.registry import (
    CustomFieldsInstanceMixin,
    CustomFieldsRegistryMixin,
    Instance,
    MapAPIExceptionsInstanceMixin,
    ModelInstanceMixin,
    ModelRegistryMixin,
    Registry,
)
from baserow.core.telemetry.utils import baserow_trace_methods
from baserow.core.utils import Progress

from .exceptions import JobTypeAlreadyRegistered, JobTypeDoesNotExist
from .models import Job
from .types import AnyJob

tracer = trace.get_tracer(__name__)


class JobType(
    CustomFieldsInstanceMixin,
    ModelInstanceMixin,
    MapAPIExceptionsInstanceMixin,
    Instance,
    metaclass=baserow_trace_methods(tracer, only="do"),
):

    """
    This abstract class represents a custom job type that can be added to the
    job type registry. It must be extended so customization can be done. Each job
    type will have its own `run` method that will be run asynchronously.
    """

    job_exceptions_map = {}

    """
    A map of exception that can be used to map exceptions to certain task error
    messages.
    """

    def transaction_atomic_context(self, job: Job):
        """
        This method gives the possibility to change the transaction context per request.
        """

        return transaction_atomic()

    def prepare_values(
        self, values: Dict[str, Any], user: AbstractUser
    ) -> Dict[str, Any]:
        """
        The prepare_values hook gives the possibility to change the provided values
        that just before they are going to be used to create or update the instance. For
        example if an ID is provided, it can be converted to a model instance. Or to
        convert a certain date string to a date object. It's also an opportunity to add
        specific validations.

        :param values: The provided values.
        :param user: The user on whose behalf the change is made.
        :return: The updated values.
        """

        return values

    def after_job_creation(self, job: AnyJob, values: Dict[str, Any]):
        """
        This method gives the possibility to change the job just after the
        instance creation. For example, files can be saved, or relationship can be
        added.

        :param job: The created job.
        :param values: The provided values.
        """

    def run(self, job: AnyJob, progress: Progress) -> Any:
        """
        This method is the task of this job type that will be executed asynchronously.

        :param job: the specific instance of the related job instance
        :param progress: A progress object that can be used to track the progress of
          the task.
        """

        raise NotImplementedError("The run method must be implemented.")

    def before_delete(self, job):
        """
        If a job type need to do something before a job deletion, can be done here.
        This method is do nothing by default.
        """

    def on_error(self, job: AnyJob, error: Exception):
        """
        This method gives the possibility to change the job after an exception has
        been raised. This happens after the transaction so database changes are kept.

        :param job: the specific instance of the related job instance
        :param error: the exception raised.
        """


class JobTypeRegistry(
    CustomFieldsRegistryMixin,
    ModelRegistryMixin[Job, JobType],
    Registry[JobType],
):
    """
    The registry that holds all the available job types.
    """

    name = "job_type"

    does_not_exist_exception_class = JobTypeDoesNotExist
    already_registered_exception_class = JobTypeAlreadyRegistered


job_type_registry: JobTypeRegistry = JobTypeRegistry()