import datetime
import logging
import uuid
from typing import TYPE_CHECKING, Any, Generic, Optional, TypeVar

import django
from django.conf import settings
from django.core.exceptions import SuspiciousOperation
from django.db import models
from django.db.models import F, Q
from django.db.models.constraints import CheckConstraint
from django.utils import timezone
from django.utils.module_loading import import_string
from django.utils.translation import gettext_lazy as _
from django_tasks.base import (
    DEFAULT_TASK_PRIORITY,
    DEFAULT_TASK_QUEUE_NAME,
    TASK_MAX_PRIORITY,
    TASK_MIN_PRIORITY,
    Task,
    TaskError,
    TaskResultStatus,
)
from django_tasks.utils import (
    get_exception_traceback,
    get_module_path,
)
from typing_extensions import ParamSpec

from .compat import TASK_CLASSES
from .utils import normalize_uuid, retry

logger = logging.getLogger("django_tasks_db")

T = TypeVar("T")
P = ParamSpec("P")

if TYPE_CHECKING:
    from .backend import TaskResult

    class GenericBase(Generic[P, T]):
        pass

else:

    class GenericBase:
        """
        https://code.djangoproject.com/ticket/33174
        """

        def __class_getitem__(cls, _):
            return cls


def get_date_max() -> datetime.datetime:
    return datetime.datetime(
        9999, 1, 1, tzinfo=datetime.timezone.utc if settings.USE_TZ else None
    )


class DBTaskResultQuerySet(models.QuerySet):
    def ready(self) -> "DBTaskResultQuerySet":
        """
        Return tasks which are ready to be processed.
        """
        return self.filter(
            status=TaskResultStatus.READY,
        ).filter(
            models.Q(run_after=get_date_max()) | models.Q(run_after__lte=timezone.now())
        )

    def successful(self) -> "DBTaskResultQuerySet":
        return self.filter(status=TaskResultStatus.SUCCESSFUL)

    def failed(self) -> "DBTaskResultQuerySet":
        return self.filter(status=TaskResultStatus.FAILED)

    def running(self) -> "DBTaskResultQuerySet":
        return self.filter(status=TaskResultStatus.RUNNING)

    def finished(self) -> "DBTaskResultQuerySet":
        return self.failed() | self.successful()

    @retry()
    def get_locked(self) -> Optional["DBTaskResult"]:
        """
        Get a job, locking the row and accounting for deadlocks.
        """
        return self.select_for_update(skip_locked=True).first()


class DBTaskResult(GenericBase[P, T], models.Model):
    id = models.UUIDField(primary_key=True, default=uuid.uuid4, editable=False)

    status = models.CharField(
        _("status"),
        choices=TaskResultStatus.choices,
        default=TaskResultStatus.READY,
        max_length=max(len(value) for value in TaskResultStatus.values),
    )

    enqueued_at = models.DateTimeField(_("enqueued at"), auto_now_add=True)
    started_at = models.DateTimeField(_("started at"), null=True)
    finished_at = models.DateTimeField(_("finished at"), null=True)

    args_kwargs = models.JSONField(_("args kwargs"))

    priority = models.IntegerField(_("priority"), default=DEFAULT_TASK_PRIORITY)

    task_path = models.TextField(_("task path"))
    worker_ids = models.JSONField(_("worker id"), default=list)

    queue_name = models.CharField(
        _("queue name"), default=DEFAULT_TASK_QUEUE_NAME, max_length=32
    )
    backend_name = models.CharField(_("backend name"), max_length=32)

    run_after = models.DateTimeField(_("run after"))

    return_value = models.JSONField(_("return value"), default=None, null=True)

    exception_class_path = models.TextField(_("exception class path"))
    traceback = models.TextField(_("traceback"))

    objects = DBTaskResultQuerySet.as_manager()

    class Meta:
        ordering = [F("priority").desc(), F("run_after").asc()]
        verbose_name = _("Task Result")
        verbose_name_plural = _("Task Results")
        indexes = [
            models.Index(
                "status",
                *ordering,
                name="tasks_db_new_ordering_idx",
                condition=Q(status=TaskResultStatus.READY),
            ),
            models.Index(fields=["queue_name"]),
            models.Index(fields=["backend_name"]),
        ]

        if django.VERSION >= (5, 1):
            constraints = [
                CheckConstraint(
                    condition=Q(priority__range=(TASK_MIN_PRIORITY, TASK_MAX_PRIORITY)),
                    name="priority_range",
                )
            ]
        else:
            constraints = [
                CheckConstraint(
                    check=Q(priority__range=(TASK_MIN_PRIORITY, TASK_MAX_PRIORITY)),
                    name="priority_range",
                )
            ]

    @property
    def task(self) -> Task[P, T]:
        task = import_string(self.task_path)

        if not isinstance(task, TASK_CLASSES):
            raise SuspiciousOperation(
                f"Task {self.id} does not point to a Task ({self.task_path})"
            )

        return task.using(  # type: ignore[no-any-return]
            priority=self.priority,
            queue_name=self.queue_name,
            run_after=None if self.run_after == get_date_max() else self.run_after,
            backend=self.backend_name,
        )

    @property
    def task_result(self) -> "TaskResult[T]":
        from .backend import TaskResult

        task_result: TaskResult[T] = TaskResult(
            db_result=self,
            task=self.task,
            id=normalize_uuid(self.id),
            status=TaskResultStatus[self.status],
            enqueued_at=self.enqueued_at,
            started_at=self.started_at,
            last_attempted_at=self.started_at,
            finished_at=self.finished_at,
            args=self.args_kwargs["args"],
            kwargs=self.args_kwargs["kwargs"],
            backend=self.backend_name,
            errors=[],
            worker_ids=self.worker_ids,
        )

        if self.status == TaskResultStatus.FAILED:
            task_result.errors.append(
                TaskError(
                    exception_class_path=self.exception_class_path,
                    traceback=self.traceback,
                )
            )

        object.__setattr__(task_result, "_return_value", self.return_value)

        return task_result

    @property
    def task_name(self) -> str:
        # If the function for an existing task is no longer available, it'll either raise an
        # ImportError or ModuleNotFoundError (a subclass of ImportError).
        try:
            return self.task.name
        except ImportError:
            pass

        try:
            return self.task_path.rsplit(".", 1)[1]
        except IndexError:
            return self.task_path

    @retry(backoff_delay=0)
    def claim(self, worker_id: str) -> None:
        """
        Mark as job as being run
        """
        self.status = TaskResultStatus.RUNNING
        self.started_at = timezone.now()
        self.worker_ids = [*self.worker_ids, worker_id]
        self.save(update_fields=["status", "started_at", "worker_ids"])

    @retry()
    def set_successful(self, return_value: Any) -> None:
        self.status = TaskResultStatus.SUCCESSFUL
        self.finished_at = timezone.now()
        self.return_value = return_value
        self.exception_class_path = ""
        self.traceback = ""

        self.save(
            update_fields=[
                "status",
                "return_value",
                "finished_at",
                "exception_class_path",
                "traceback",
            ]
        )

    @retry()
    def set_failed(self, exc: BaseException) -> None:
        self.status = TaskResultStatus.FAILED
        self.finished_at = timezone.now()
        self.exception_class_path = get_module_path(type(exc))
        self.traceback = get_exception_traceback(exc)
        self.return_value = None

        self.save(
            update_fields=[
                "status",
                "return_value",
                "finished_at",
                "exception_class_path",
                "traceback",
            ]
        )