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

"""debusine celery configuration."""

import logging
import os
import secrets
from collections.abc import Callable, Sequence
from functools import partial
from pathlib import Path
from types import SimpleNamespace
from typing import Any, TYPE_CHECKING
from uuid import uuid4

from celery import Celery, Task, states
from celery.result import AsyncResult
from celery.signals import before_task_publish, worker_init, worker_shutdown
from celery.worker import WorkController
from django.db import connection, transaction
from django.db.models import IntegerField, JSONField, QuerySet
from django.db.models.functions import Cast

from debusine.client.debusine import Debusine

if TYPE_CHECKING:
    from django_celery_results.models import TaskResult

    from debusine.db.models import Token

os.environ.setdefault("DJANGO_SETTINGS_MODULE", "debusine.project.settings")


# mypy complains that Task is untyped, which is true, but we can't fix that
# here.
class TaskWithBetterArgumentRepresentation(Task):  # type: ignore[misc]
    """
    Custom task class to control result argument serialization.

    See https://github.com/celery/django-celery-results/issues/113.
    """

    def apply_async(
        self,
        args: tuple[Any, ...] | None = None,
        kwargs: dict[str, Any] | None = None,
        **options: Any,
    ) -> AsyncResult:
        """
        Override Celery's default argument encoding in the result backend.

        By default, Celery encodes arguments using ``repr()`` when storing
        them in the result backend, which makes it difficult to query them.
        Passing them straight through allows ``django_celery_results`` to
        apply JSON encoding instead.
        """
        options["argsrepr"] = args or ()
        options["kwargsrepr"] = kwargs or {}
        return super().apply_async(args=args, kwargs=kwargs, **options)


def make_app() -> Celery:
    """Make a Celery application."""
    app = Celery("debusine", task_cls=TaskWithBetterArgumentRepresentation)
    app.config_from_object("django.conf:settings", namespace="CELERY")
    app.autodiscover_tasks()
    return app


app = make_app()


# Based loosely on:
#   https://stackoverflow.com/q/9824172
#   https://github.com/celery/django-celery-results/issues/286
# mypy complains that before_task_publish.connect is untyped, which is true,
# but we can't fix that here.
@before_task_publish.connect  # type: ignore[untyped-decorator]
def store_pending_state(
    *, sender: str, headers: dict[str, Any], **kwargs: Any
) -> None:
    """Tell the result backend more details about newly-published tasks."""
    task = app.tasks.get(sender)
    if task is not None:  # pragma: no cover
        # Since we need to tell the result backend about tasks before they
        # are created (in apply_async below), we don't support individual
        # tasks using a different backend.
        assert task.backend == app.backend
    app.backend.store_result(
        headers["id"], None, states.PENDING, request=SimpleNamespace(**headers)
    )


def _apply_async_on_commit_callback(
    task: Task,
    *,
    args: tuple[Any] | None = None,
    kwargs: dict[Any, Any] | None = None,
    task_id: str | None = None,
    **options: Any,
) -> None:
    from django_celery_results.models import TaskResult

    task.apply_async(args=args, kwargs=kwargs, task_id=task_id, **options)
    if app.conf.task_always_eager:
        # In always-eager mode, the TaskResult row won't be automatically
        # updated after the task has completed, but we have no further use
        # for it anyway.  Delete it so that code that inspects the task
        # queue doesn't think the task is still pending.
        TaskResult.objects.filter(task_id=task_id).delete()


def get_current_savepoint_run_on_commit_callbacks() -> list[Callable[[], None]]:
    """Get scheduled ``run_on_commit`` callbacks for the current savepoint."""
    # This uses undocumented Django internals, but if it goes wrong without
    # an obvious unit test failure, then the worst case should be some extra
    # calls to a function that we normally expect to be batched, which is
    # relatively harmless.
    return [
        callback
        for savepoint_ids, callback, *_ in connection.run_on_commit
        if savepoint_ids == set(connection.savepoint_ids)
    ]


def get_apply_async_tasks(
    callbacks: Sequence[Callable[[], None]],
) -> list[Task]:
    """Get Celery tasks scheduled to be run on commit."""
    return [
        callback.args[0]
        for callback in callbacks
        if isinstance(callback, partial)
        and callback.func == _apply_async_on_commit_callback
    ]


def apply_async_on_commit(
    task: Task,
    *,
    args: tuple[Any] | None = None,
    kwargs: dict[Any, Any] | None = None,
    task_id: str | None = None,
    once_per_savepoint: bool = False,
    **options: Any,
) -> str | None:
    """
    Apply a task asynchronously when the transaction is committed.

    Tell the result backend about it even before the transaction is
    committed, so that code in the same transaction can tell that the task
    is pending.

    :return: if a task was created, the Celery task ID; otherwise None.
    """
    # Some tasks should be batched up and only run once per savepoint.
    if once_per_savepoint:
        assert args is None
        assert kwargs is None
        if task in get_apply_async_tasks(
            get_current_savepoint_run_on_commit_callbacks()
        ):
            return None

    if task_id is None:
        task_id = str(uuid4())

    app.backend.store_result(
        task_id,
        None,
        states.PENDING,
        # This is just enough for code that needs to inspect the task queue.
        request=SimpleNamespace(
            task=task.name, argsrepr=args or (), kwargsrepr=kwargs or {}
        ),
    )
    transaction.on_commit(
        partial(
            _apply_async_on_commit_callback,
            task,
            args=args,
            kwargs=kwargs,
            task_id=task_id,
            **options,
        )
    )

    return task_id


def celery_tasks_in_progress(task_name: str) -> QuerySet["TaskResult"]:
    """Return a QuerySet of Celery tasks in progress with a given task name."""
    from django_celery_results.models import TaskResult

    task_results = TaskResult.objects.filter(
        status__in={
            states.PENDING,
            states.RECEIVED,
            states.STARTED,
            states.RETRY,
        },
        task_name=task_name,
    ).annotate(
        task_args_array=Cast("task_args", JSONField()),
        first_task_arg=Cast("task_args_array__0", IntegerField()),
    )
    assert isinstance(task_results, QuerySet)
    return task_results


class CeleryToken:
    """Manage a Celery worker's token."""

    token: "Token"
    client: Debusine

    def __init__(self, name: str) -> None:
        """Read an existing token key, or generate one if there is none."""
        from django.conf import settings

        from debusine.db.models import Token
        from debusine.utils import atomic_writer

        path = Path(settings.DEBUSINE_TOKENS_DIRECTORY, name)
        try:
            key = path.read_text()
        except FileNotFoundError:
            key = secrets.token_hex(32)
            with atomic_writer(path, mode="w", chmod=0o600) as fd:
                fd.write(key)

        self.token, _ = Token.objects.get_or_create(
            token_type=Token.Types.WORKER, hash=Token._generate_hash(key)
        )
        self.client = Debusine(
            f"https://{settings.DEBUSINE_FQDN}/api",
            key,
            logger=logging.getLogger(),
        )


_celery_token: CeleryToken | None = None


def get_celery_token() -> CeleryToken | None:
    """Get information about the current Celery worker token."""
    return _celery_token


def get_celery_worker_provided_tags() -> set[str]:
    """Return a set with tags provided by the Celery worker."""
    from debusine.server.tasks import get_provided_celery_worker_tags
    from debusine.tasks.models import WorkerType
    from debusine.worker.system_information import system_provided_tags

    tags: set[str] = set()
    tags |= system_provided_tags(WorkerType.CELERY)
    tags |= get_provided_celery_worker_tags()
    return tags


def get_celery_worker_metadata() -> dict[str, Any]:
    """Compute the dynamic metadata sent by the Celery worker."""
    from debusine.tasks.models import WorkerType
    from debusine.worker.system_information import system_metadata

    return {
        **system_metadata(WorkerType.CELERY),
        "provided_tags": sorted(get_celery_worker_provided_tags()),
    }


def get_scheduler_worker_metadata() -> dict[str, Any]:
    """Compute the dynamic metadata sent by the scheduler worker."""
    from debusine.tasks.models import WorkerType
    from debusine.worker.system_information import system_metadata

    return {**system_metadata(WorkerType.SCHEDULER)}


# mypy complains that worker_init.connect is untyped, which is true, but we
# can't fix that here.
@worker_init.connect  # type: ignore[untyped-decorator]
@transaction.atomic
def connect_worker(*, sender: WorkController, **kwargs: Any) -> None:
    """Set up and connect a :py:class:`Worker` when Celery starts."""
    from debusine.db.models import Worker

    worker: Worker | None = None

    if sender.hostname.startswith("celery@"):
        worker = Worker.objects.get_or_create_celery()
        worker.concurrency = sender.concurrency
        worker.save()
        metadata = get_celery_worker_metadata()
    elif sender.hostname.startswith("scheduler@"):
        worker = Worker.objects.get_or_create_scheduler()
        metadata = get_scheduler_worker_metadata()

    if worker is not None:
        global _celery_token
        _celery_token = CeleryToken(sender.hostname.split("@")[0])
        if worker.token is None:
            worker.token = _celery_token.token
        worker.save()
        worker.token.enable()

        worker.set_dynamic_metadata(metadata)
        worker.mark_connected()


# mypy complains that worker_shutdown.connect is untyped, which is true, but
# we can't fix that here.
@worker_shutdown.connect  # type: ignore[untyped-decorator]
def disconnect_worker(*, sender: WorkController, **kwargs: Any) -> None:
    """Disconnect the appropriate :py:class:`Worker` when Celery exits."""
    from debusine.db.models import Worker

    if sender.hostname.startswith("celery@"):
        Worker.objects.get_or_create_celery().mark_disconnected()
    elif sender.hostname.startswith("scheduler@"):
        Worker.objects.get_or_create_scheduler().mark_disconnected()

    global _celery_token
    _celery_token = None
