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

"""Unit tests for Debusine Celery configuration."""

import json
from multiprocessing import Lock
from pathlib import Path
from typing import Any
from unittest import mock

from celery import shared_task, signals, states
from celery.contrib.testing.worker import start_worker
from celery.utils.nodenames import gethostname, nodename
from django.conf import settings
from django.test import override_settings
from django_celery_results.models import TaskResult

from debusine.db.models import Token, Worker
from debusine.project.celery import CeleryToken, app, get_celery_token, make_app
from debusine.tasks.models import WorkerType
from debusine.test.django import TestCase, TransactionTestCase

_lock = Lock()


# mypy complains that celery.shared_task is untyped, which is true, but we
# can't fix that here.
@shared_task  # type: ignore[untyped-decorator]
def _locked_noop() -> None:
    """Sample Celery task to do nothing, once a lock is available."""
    with _lock:
        pass


# mypy complains that celery.shared_task is untyped, which is true, but we
# can't fix that here.
@shared_task  # type: ignore[untyped-decorator]
def _locked_reverse_string(argument: str, **kwargs: Any) -> str:
    """Sample Celery task to reverse a string, once a lock is available."""
    with _lock:
        return "".join(reversed(argument))


@override_settings(CELERY_BROKER_URL="memory://")
class TestStorePendingState(TransactionTestCase):
    """Tests for :py:func:`store_pending_state`."""

    def setUp(self) -> None:
        super().setUp()
        self.app = make_app()

    def test_result_states(self) -> None:
        """The task's state can be tracked in the database."""
        with (
            mock.patch("debusine.django.django_utils.test_data_override"),
            start_worker(self.app, perform_ping_check=False) as celery_worker,
        ):
            # Schedule a task, but with a lock held so that we get a chance
            # to look at its result in the database before it finishes.
            with _lock:
                result = _locked_reverse_string.delay(
                    "hello", foo="bar", value=1
                )

                expected_task_name = f"{__name__}._locked_reverse_string"
                expected_task_args = ["hello"]
                expected_task_kwargs = {"foo": "bar", "value": 1}
                task_result = TaskResult.objects.get(task_id=result.id)
                self.assertEqual(task_result.status, states.PENDING)
                self.assertEqual(task_result.task_name, expected_task_name)
                self.assertEqual(
                    json.loads(task_result.task_args), expected_task_args
                )
                self.assertEqual(
                    json.loads(task_result.task_kwargs), expected_task_kwargs
                )

            # Release the lock to allow the task to complete.  Its result in
            # the database is updated with the new status.
            self.assertEqual(result.get(), "olleh")

            task_result.refresh_from_db()
            self.assertEqual(task_result.status, states.SUCCESS)
            self.assertEqual(task_result.task_name, expected_task_name)
            self.assertEqual(
                json.loads(task_result.task_args), expected_task_args
            )
            self.assertEqual(
                json.loads(task_result.task_kwargs), expected_task_kwargs
            )

            signals.worker_shutdown.send(sender=celery_worker)

    def test_result_states_no_arguments(self) -> None:
        """A task with no arguments is tracked correctly in the database."""
        with (
            mock.patch("debusine.django.django_utils.test_data_override"),
            start_worker(self.app, perform_ping_check=False) as celery_worker,
        ):
            # Schedule a task, but with a lock held so that we get a chance
            # to look at its result in the database before it finishes.
            with _lock:
                result = _locked_noop.apply_async()

                expected_task_name = f"{__name__}._locked_noop"
                task_result = TaskResult.objects.get(task_id=result.id)
                self.assertEqual(task_result.status, states.PENDING)
                self.assertEqual(task_result.task_name, expected_task_name)
                self.assertEqual(json.loads(task_result.task_args), [])
                self.assertEqual(json.loads(task_result.task_kwargs), {})

            # Release the lock to allow the task to complete.  Its result in
            # the database is updated with the new status.
            self.assertIsNone(result.get())

            task_result.refresh_from_db()
            self.assertEqual(task_result.status, states.SUCCESS)
            self.assertEqual(task_result.task_name, expected_task_name)
            self.assertEqual(json.loads(task_result.task_args), [])
            self.assertEqual(json.loads(task_result.task_kwargs), {})

            signals.worker_shutdown.send(sender=celery_worker)


class TestCeleryToken(TestCase):
    """Tests for :py:class:`CeleryToken`."""

    def test_existing_token(self) -> None:
        tokens_dir = self.create_temporary_directory()
        key = "existing-key"
        Path(tokens_dir, "test").write_text(key)

        with override_settings(
            DEBUSINE_FQDN="debusine.example.org",
            DEBUSINE_TOKENS_DIRECTORY=str(tokens_dir),
        ):
            celery_token = CeleryToken("test")

        self.assertEqual(celery_token.token.hash, Token._generate_hash(key))
        self.assertEqual(
            celery_token.client.base_api_url, "https://debusine.example.org/api"
        )
        self.assertEqual(celery_token.client.token, key)
        self.assertIsNone(celery_token.client.scope)

    def test_new_token(self) -> None:
        tokens_dir = self.create_temporary_directory()

        with override_settings(
            DEBUSINE_FQDN="debusine.example.org",
            DEBUSINE_TOKENS_DIRECTORY=str(tokens_dir),
        ):
            celery_token = CeleryToken("test")

        key = Path(tokens_dir, "test").read_text()
        self.assertEqual(len(key), 64)
        self.assertEqual(celery_token.token.hash, Token._generate_hash(key))
        self.assertEqual(
            celery_token.client.base_api_url, "https://debusine.example.org/api"
        )
        self.assertEqual(celery_token.client.token, key)
        self.assertIsNone(celery_token.client.scope)


@override_settings(CELERY_BROKER_URL="memory://")
class TestConnectWorker(TransactionTestCase):
    """Tests for :py:func:`connect_worker`."""

    def setUp(self) -> None:
        super().setUp()
        self.app = make_app()
        Worker.objects.filter(
            worker_type__in={WorkerType.CELERY, WorkerType.SCHEDULER}
        ).delete()

    def test_server_task_worker(self) -> None:
        """The main worker creates a Worker instance that runs server tasks."""
        with (
            mock.patch("debusine.django.django_utils.test_data_override"),
            # Set a concurrency that's unlikely to be the actual number of
            # CPUs we have.
            start_worker(
                self.app,
                concurrency=5,
                perform_ping_check=False,
                hostname=nodename("celery", gethostname()),
            ) as celery_worker,
        ):
            worker = Worker.objects.get(worker_type=WorkerType.CELERY)
            self.assertEqual(worker.concurrency, 5)
            self.assertEqual(
                worker.dynamic_metadata["system:worker_type"], WorkerType.CELERY
            )
            self.assertIsNotNone(worker.connected_at)
            self.assertFalse(
                Worker.objects.filter(worker_type=WorkerType.SCHEDULER).exists()
            )

            token = Token.objects.get_token_or_none(
                Path(settings.DEBUSINE_TOKENS_DIRECTORY, "celery").read_text()
            )
            assert token is not None
            self.assertEqual(token.worker, worker)
            self.assertTrue(token.enabled)
            celery_token = get_celery_token()
            assert celery_token is not None
            self.assertEqual(celery_token.token, token)

            signals.worker_shutdown.send(sender=celery_worker)

            worker.refresh_from_db()
            self.assertIsNone(worker.connected_at)
            self.assertFalse(
                Worker.objects.filter(worker_type=WorkerType.SCHEDULER).exists()
            )

    def test_server_task_worker_token_already_exists(self) -> None:
        existing_worker = Worker.objects.get_or_create_celery()
        existing_worker.token = CeleryToken("celery").token
        existing_worker.save()

        with (
            mock.patch("debusine.django.django_utils.test_data_override"),
            # Set a concurrency that's unlikely to be the actual number of
            # CPUs we have.
            start_worker(
                self.app,
                concurrency=5,
                perform_ping_check=False,
                hostname=nodename("celery", gethostname()),
            ) as celery_worker,
        ):
            worker = Worker.objects.get(worker_type=WorkerType.CELERY)
            self.assertEqual(worker, existing_worker)

            token = Token.objects.get_token_or_none(
                Path(settings.DEBUSINE_TOKENS_DIRECTORY, "celery").read_text()
            )
            assert token is not None
            self.assertEqual(token.worker, worker)
            self.assertTrue(token.enabled)
            self.assertEqual(token, existing_worker.token)
            celery_token = get_celery_token()
            assert celery_token is not None
            self.assertEqual(celery_token.token, token)

            signals.worker_shutdown.send(sender=celery_worker)

    def test_scheduler_worker(self) -> None:
        """The scheduler worker creates a Worker instance."""
        with (
            mock.patch("debusine.django.django_utils.test_data_override"),
            start_worker(
                app,
                concurrency=1,
                perform_ping_check=False,
                hostname=nodename("scheduler", gethostname()),
            ) as celery_worker,
        ):
            worker = Worker.objects.get(worker_type=WorkerType.SCHEDULER)
            self.assertEqual(worker.concurrency, 1)
            self.assertEqual(
                worker.dynamic_metadata["system:worker_type"],
                WorkerType.SCHEDULER,
            )
            self.assertIsNotNone(worker.connected_at)
            self.assertFalse(
                Worker.objects.filter(worker_type=WorkerType.CELERY).exists()
            )

            token = Token.objects.get_token_or_none(
                Path(
                    settings.DEBUSINE_TOKENS_DIRECTORY, "scheduler"
                ).read_text()
            )
            assert token is not None
            self.assertEqual(token.worker, worker)
            self.assertTrue(token.enabled)
            celery_token = get_celery_token()
            assert celery_token is not None
            self.assertEqual(celery_token.token, token)

            signals.worker_shutdown.send(sender=celery_worker)

            worker.refresh_from_db()
            self.assertIsNone(worker.connected_at)
            self.assertFalse(
                Worker.objects.filter(worker_type=WorkerType.CELERY).exists()
            )

    def test_scheduler_worker_token_already_exists(self) -> None:
        existing_worker = Worker.objects.get_or_create_scheduler()
        existing_worker.token = CeleryToken("scheduler").token
        existing_worker.save()

        with (
            mock.patch("debusine.django.django_utils.test_data_override"),
            start_worker(
                app,
                concurrency=1,
                perform_ping_check=False,
                hostname=nodename("scheduler", gethostname()),
            ) as celery_worker,
        ):
            worker = Worker.objects.get(worker_type=WorkerType.SCHEDULER)
            self.assertEqual(worker, existing_worker)

            token = Token.objects.get_token_or_none(
                Path(
                    settings.DEBUSINE_TOKENS_DIRECTORY, "scheduler"
                ).read_text()
            )
            assert token is not None
            self.assertEqual(token.worker, worker)
            self.assertTrue(token.enabled)
            self.assertEqual(token, existing_worker.token)
            celery_token = get_celery_token()
            assert celery_token is not None
            self.assertEqual(celery_token.token, token)

            signals.worker_shutdown.send(sender=celery_worker)

    def test_provisioner_worker(self) -> None:
        """The provisioner worker does not create a Worker instance."""
        with (
            mock.patch("debusine.django.django_utils.test_data_override"),
            start_worker(
                app,
                concurrency=1,
                perform_ping_check=False,
                hostname=nodename("provisioner", gethostname()),
            ) as celery_worker,
        ):
            self.assertFalse(
                Worker.objects.filter(
                    worker_type__in={WorkerType.CELERY, WorkerType.SCHEDULER}
                ).exists()
            )
            self.assertIsNone(get_celery_token())

            signals.worker_shutdown.send(sender=celery_worker)

            self.assertFalse(
                Worker.objects.filter(
                    worker_type__in={WorkerType.CELERY, WorkerType.SCHEDULER}
                ).exists()
            )
