# Copyright 2022 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 django usage."""

import asyncio
from datetime import timedelta
from typing import Any, cast
from unittest import TestCase
from unittest.util import safe_repr

from channels.layers import get_channel_layer

from django.contrib.auth import get_user_model
from django.contrib.auth.models import Permission
from django.contrib.contenttypes.models import ContentType
from django.db.models import Model

from debusine.artifacts.models import ArtifactCategory, CollectionCategory
from debusine.db.models import (
    Artifact,
    ArtifactRelation,
    Collection,
    DEFAULT_FILE_STORE_NAME,
    DEFAULT_WORKSPACE_NAME,
    File,
    FileInArtifact,
    FileStore,
    FileUpload,
    Token,
    User,
    WorkRequest,
    Worker,
    WorkflowTemplate,
    Workspace,
    default_workspace,
)
from debusine.db.tests.utils import _calculate_hash_from_data
from debusine.server.file_backend.interface import FileBackendInterface
from debusine.tasks.models import WorkerType
from debusine.test import _BaseTestHelpersMixin
from debusine.test.utils import data_generator


class _DatabaseHelpersMixin(_BaseTestHelpersMixin):
    """
    Collection of database-related methods to help write unit tests.

    This mixin-class provides different methods to set up various database
    entries using the corresponding Django models.

    In unit tests use TestHelpersMixin class.
    """

    @staticmethod
    def get_test_user():
        """Return a test user."""
        user_model = get_user_model()
        try:
            return user_model.objects.get(username="usertest")
        except user_model.DoesNotExist:
            user = get_user_model().objects.create_user(
                username="usertest",
                password="userpassword",
                email="usertest@example.org",
            )
        return user

    @classmethod
    def create_token_enabled(cls, with_user: bool = False, **kwargs):
        """
        Return an enabled Token.

        :param with_user: if True it assigns a User to this token.
        """
        if "enabled" in kwargs:
            raise ValueError('"enabled" cannot be set in create_token_enabled')

        if with_user:
            kwargs["user"] = cls.get_test_user()

        token = Token.objects.create(**kwargs)
        token.enable()
        return token

    def add_user_permission(
        self, user: User, model: type[Model], codename: str
    ) -> None:
        """Add a permission to a user."""
        user.user_permissions.add(
            Permission.objects.get(
                codename=codename,
                content_type=ContentType.objects.get_for_model(model),
            )
        )

    def create_file_in_backend(
        self,
        backend: FileBackendInterface[Any],
        contents: bytes = b"test",
        *,
        remove_file_from_backend=True,
    ) -> File:
        """
        Create a temporary file and adds it in the backend.

        :param backend: file backend to add the file in
        :param contents: contents of the file
        :param remove_file_from_backend: schedule removing file from the backend
        """
        temporary_file = self.create_temporary_file(contents=contents)

        fileobj = backend.add_file(temporary_file)

        if remove_file_from_backend:
            cast(TestCase, self).addCleanup(backend.remove_file, fileobj)

        return fileobj

    @staticmethod
    def create_artifact_relation(
        artifact: Artifact,
        target: Artifact,
        relation_type: ArtifactRelation.Relations = (
            ArtifactRelation.Relations.RELATES_TO
        ),
    ) -> ArtifactRelation:
        """Create an ArtifactRelation."""
        return ArtifactRelation.objects.create(
            artifact=artifact, target=target, type=relation_type
        )

    def create_artifact(
        self,
        paths: list[str] | None = None,
        files_size: int = 100,
        *,
        category: ArtifactCategory = ArtifactCategory.TEST,
        workspace: Workspace | None = None,
        data: dict[str, Any] | None = None,
        expiration_delay: int | None = None,
        work_request: WorkRequest | None = None,
        created_by: User | None = None,
        create_files: bool = False,
        skip_add_files_in_store: bool = False,
    ) -> tuple[Artifact, dict[str, bytes]]:
        """
        Create an artifact and return tuple with the artifact and files data.

        :param paths: list of paths to create (will contain random data)
        :param files_size: size of the test data
        :param category: this artifact's category (see :ref:`artifacts`)
        :param data: key-value data for this artifact (see :ref:`artifacts`)
        :param expiration_delay: set expiration_delay field (in days)
        :param work_request: work request that created this artifact
        :param created_by: set Artifact.created_by to it
        :param create_files: create a file and add it into the LocalFileBackend
        :param skip_add_files_in_store: do not add the files in the store
          (only create the File object in the database)

        This method return a tuple:
        - artifact: Artifact
        - files_contents: Dict[str, bytes] (paths and test data)
        """
        if skip_add_files_in_store and not create_files:
            raise ValueError(
                "skip_add_files_in_store must be False if create_files is False"
            )

        artifact = Artifact.objects.create(
            category=category,
            workspace=workspace or default_workspace(),
            data=data or {},
            expiration_delay=(
                timedelta(expiration_delay)
                if expiration_delay is not None
                else None
            ),
            created_by_work_request=work_request,
            created_by=created_by,
        )

        data_gen = data_generator(files_size)

        files_contents = {}

        if paths is not None:
            for path in paths:
                files_contents[path] = next(data_gen)

        if create_files:
            file_backend = FileStore.default().get_backend_object()

            for path, contents in files_contents.items():
                if skip_add_files_in_store:
                    fileobj, _ = File.get_or_create(
                        hash_digest=_calculate_hash_from_data(contents),
                        size=len(contents),
                    )
                else:
                    fileobj = self.create_file_in_backend(
                        file_backend, contents
                    )

                FileInArtifact.objects.create(
                    artifact=artifact, path=path, file=fileobj
                )

        return artifact, files_contents

    def create_file_upload(self):
        """
        Create a new FileUpload object.

        Create the workspace, artifact, file and file_in_artifact associated
        to the file_upload object.
        """
        artifact, _ = self.create_artifact(
            paths=["README"], create_files=True, skip_add_files_in_store=True
        )

        file_in_artifact = artifact.fileinartifact_set.first()

        file_upload = FileUpload.objects.create(
            file_in_artifact=file_in_artifact,
            path="temp_file_aaaa",
        )

        file_path = file_upload.absolute_file_path()

        self.addCleanup(file_path.unlink, missing_ok=True)

        return file_upload

    def create_file(self, contents: bytes = b"test"):
        """
        Create a File model and return the saved fileobj.

        :param contents: used to compute hash digest and size
        """
        hashed = _calculate_hash_from_data(contents)

        file, _ = File.get_or_create(hash_digest=hashed, size=len(contents))

        cast(TestCase, self).assertEqual(file.hash_digest, hashed)
        cast(TestCase, self).assertEqual(file.size, len(contents))

        return file

    @classmethod
    def create_work_request(
        cls,
        mark_running: bool = False,
        assign_new_worker: bool = False,
        result: WorkRequest.Results | None = None,
        **kwargs,
    ) -> WorkRequest:
        """
        Return a new instance of WorkRequest.

        :param mark_running: if True call mark_running() method
        :param assign_new_worker: if not None assign worker to the work request
        :param result: if not None call mark_completed(result)
        :param kwargs: use them when creating the WorkRequest model
        """
        # FileStore and Workspace are usually created by the migrations
        # When using a TransactionTestCase with async methods
        # the tests don't have access to the created FileStore / Workspace
        # yet (need verification, but they seem to be in a non-committed
        # transaction and the test code is in a different thread because
        # of the implementation of database_sync_to_async
        file_store, _ = FileStore.objects.get_or_create(
            name=DEFAULT_FILE_STORE_NAME
        )

        workspace, _ = Workspace.objects.get_or_create(
            name=DEFAULT_WORKSPACE_NAME, default_file_store=file_store
        )

        defaults = {"workspace": workspace}
        if "created_by" not in kwargs:
            token = cls.create_token_enabled(with_user=True)
            defaults["created_by"] = token.user

        defaults.update(kwargs)

        created_at = kwargs.pop("created_at", None)

        work_request = WorkRequest.objects.create(**defaults)

        # created_at is set to now automatically by create via auto_now_add, so
        # if a specific value was requested we need to set it after creation
        if created_at is not None:
            work_request.created_at = created_at
            work_request.save()

        if assign_new_worker:
            work_request.assign_worker(cls.create_worker())

        if mark_running:
            work_request.mark_running()

        if result is not None:
            work_request.mark_completed(result)

        if completed_at := kwargs.get("completed_at"):
            work_request.completed_at = completed_at
            work_request.save()

        return work_request

    @classmethod
    def create_worker(cls):
        """Return a new Worker."""
        token = cls.create_token_enabled()

        worker = Worker.objects.create_with_fqdn("computer.lan", token)

        worker.set_dynamic_metadata(
            {
                "system:cpu_cores": 4,
                "system:worker_type": WorkerType.EXTERNAL,
                "sbuild:version": 1,
                "sbuild:chroots": "bullseye-amd64",
            }
        )

        return worker

    @staticmethod
    def create_workspace(**kwargs) -> Workspace:
        """Create a Workspace and return it."""
        file_store, _ = FileStore.objects.get_or_create(
            name=DEFAULT_FILE_STORE_NAME
        )

        workspace_name = kwargs.pop("name", DEFAULT_WORKSPACE_NAME)
        workspace, _ = Workspace.objects.get_or_create(
            name=workspace_name,
            defaults={**{"default_file_store": file_store, **kwargs}},
        )

        return workspace

    @staticmethod
    def create_collection(
        name: str,
        category: CollectionCategory,
        *,
        workspace: Workspace | None = None,
        data: dict[str, Any] | None = None,
    ) -> Collection:
        """Create a collection."""
        return Collection.objects.create(
            name=name,
            category=category,
            workspace=workspace or default_workspace(),
            data=data or {},
        )

    @staticmethod
    def create_workflow_template(
        name: str,
        task_name: str,
        *,
        workspace: Workspace | None = None,
        task_data: dict[str, Any] | None = None,
        **kwargs: Any,
    ) -> WorkflowTemplate:
        """Create a workflow template."""
        return WorkflowTemplate.objects.create(
            name=name,
            workspace=workspace or default_workspace(),
            task_name=task_name,
            task_data=task_data or {},
            **kwargs,
        )


class ChannelsHelpersMixin:
    """
    Channel-related methods to help writing unit tests.

    Provides methods to setup a channel and assert messages or lack of messages.
    """

    async def create_channel(self, group_name) -> dict[str, Any]:
        """
        Create a channel and add it to the group named ``group_name``.

        Return dict with layer and name.
        """
        channel_layer = get_channel_layer()
        channel_name = await channel_layer.new_channel()
        await channel_layer.group_add(group_name, channel_name)

        return {"layer": channel_layer, "name": channel_name}

    async def assert_channel_nothing_received(self, channel: dict[str, Any]):
        """Assert that nothing is received in channel."""
        try:
            received = await asyncio.wait_for(
                channel["layer"].receive(channel["name"]), timeout=0.1
            )
        except asyncio.exceptions.TimeoutError:
            pass
        else:
            cast(TestCase, self).fail(
                "Expected nothing. Received: '%s'" % safe_repr(received)
            )

    async def assert_channel_received(self, channel: dict[str, Any], data):
        """Assert that data is received in channel_layer, channel_name."""
        try:
            received = await asyncio.wait_for(
                channel["layer"].receive(channel["name"]), timeout=0.1
            )
            cast(TestCase, self).assertEqual(received, data)
        except asyncio.exceptions.TimeoutError:
            cast(TestCase, self).fail(
                "Expected '%s' received nothing" % safe_repr(data)
            )
