# Copyright 2022-2024 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.

"""Infrastructure to create test scenarios in the database."""

import tempfile
from datetime import timedelta
from pathlib import Path
from typing import Any

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.artifacts.playground import ArtifactPlayground
from debusine.client.models import model_to_json_serializable_dict
from debusine.db.models import (
    Artifact,
    ArtifactRelation,
    Collection,
    CollectionItem,
    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.collections.lookup import lookup_single
from debusine.server.file_backend.interface import FileBackendInterface
from debusine.server.workflows.models import WorkRequestWorkflowData
from debusine.tasks.models import (
    BackendType,
    SbuildBuildComponent,
    SbuildData,
    SbuildInput,
    WorkerType,
)


class Playground(ArtifactPlayground):
    """Generate data scenarios for the database."""

    def __init__(
        self,
        default_file_store_name: str = DEFAULT_FILE_STORE_NAME,
        default_workspace_name: str = DEFAULT_WORKSPACE_NAME,
        default_username: str = "playground",
        default_user_password: str | None = None,
        default_user_email: str = "playground@example.org",
    ):
        """Set default values."""
        super().__init__()
        self.default_file_store_name = default_file_store_name
        self.default_workspace_name = default_workspace_name
        self.default_username = default_username
        self.default_user_password = default_user_password
        self.default_user_email = default_user_email

    def get_default_user(self) -> User:
        """Return the default user for playground methods."""
        try:
            return User.objects.get(username=self.default_username)
        except User.DoesNotExist:
            user = User.objects.create_user(
                username=self.default_username,
                email=self.default_user_email,
            )
            if self.default_user_password is not None:
                user.set_password(self.default_user_password)
                user.save()
        return user

    def get_default_file_store(self) -> FileStore:
        """Get the default file store used."""
        try:
            return FileStore.objects.get(name=self.default_file_store_name)
        except FileStore.DoesNotExist:
            return FileStore.objects.create(
                name=self.default_file_store_name,
                backend=FileStore.BackendChoices.LOCAL,
                configuration={},
            )

    def get_default_workspace(self) -> Workspace:
        """Get the default workspace used by create methods."""
        return self.create_workspace()

    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_token_enabled(
        self, with_user: bool = False, **kwargs: Any
    ) -> Token:
        """
        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"] = self.get_default_user()

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

    def create_workspace(
        self, *, file_store: FileStore | None = None, **kwargs: Any
    ) -> Workspace:
        """Create a Workspace and return it."""
        if file_store is None:
            file_store = self.get_default_file_store()
        workspace_name = kwargs.pop("name", self.default_workspace_name)
        workspace, _ = Workspace.objects.get_or_create(
            name=workspace_name,
            defaults={**{"default_file_store": file_store, **kwargs}},
        )

        return workspace

    def create_collection(
        self,
        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 {},
        )

    def create_workflow_template(
        self,
        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,
        )

    def create_work_request(
        self,
        mark_running: bool = False,
        assign_new_worker: bool = False,
        result: WorkRequest.Results | None = None,
        **kwargs: Any,
    ) -> 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
        workspace = self.create_workspace()

        defaults: dict[str, Any] = {"workspace": workspace}
        if "created_by" not in kwargs:
            token = self.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(self.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

    def create_worker(self, worker_type: WorkerType = WorkerType.EXTERNAL):
        """Return a new Worker."""
        token = self.create_token_enabled()

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

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

        return worker

    def create_file(self, contents: bytes = b"test") -> File:
        """
        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))
        return file

    def create_file_in_backend(
        self,
        backend: FileBackendInterface[Any] | None = None,
        contents: bytes = b"test",
    ) -> 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
        """
        if backend is None:
            backend = self.get_default_file_store().get_backend_object()
        with tempfile.NamedTemporaryFile("w+b") as fd:
            fd.write(contents)
            fd.flush()

            return backend.add_file(Path(fd.name))

    def create_artifact(
        self,
        paths: list[str] | dict[str, bytes] | 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)
        """
        # Import here to avoid a circular loop
        from debusine.test.utils import data_generator

        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"
            )

        if workspace is None:
            workspace = self.get_default_workspace()

        artifact = Artifact.objects.create(
            category=category,
            workspace=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 isinstance(paths, dict):
            files_contents.update(paths)
        elif paths is None:
            pass
        else:
            for path in paths:
                files_contents[path] = next(data_gen)

        if create_files:
            file_backend = workspace.default_file_store.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) -> FileUpload:
        """
        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()
        assert file_in_artifact is not None

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

    def create_artifact_relation(
        self,
        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_source_artifact(
        self,
        *,
        name: str = "hello",
        version: str = "1.0-1",
        workspace: Workspace | None = None,
        created_by: User | None = None,
        create_files: bool = False,
    ) -> Artifact:
        """Create an artifact for a source package."""
        with tempfile.TemporaryDirectory() as tempdir:
            workdir = Path(tempdir)
            source_package = self.create_source_package(
                workdir, name=name, version=version
            )
            source, _ = self.create_artifact(
                paths={
                    name: path.read_bytes()
                    for name, path in source_package.files.items()
                },
                category=ArtifactCategory.SOURCE_PACKAGE,
                workspace=workspace or self.get_default_workspace(),
                data=model_to_json_serializable_dict(source_package.data),
                created_by=created_by or self.get_default_user(),
                create_files=create_files,
            )
            return source

    def create_debian_environments_collection(
        self, name: str = "debian", *, workspace: Workspace | None = None
    ) -> Collection:
        """Create a debian:environments collection."""
        if workspace is None:
            workspace = self.get_default_workspace()
        collection, _ = Collection.objects.get_or_create(
            name=name,
            category=CollectionCategory.ENVIRONMENTS,
            workspace=workspace,
        )
        return collection

    def create_debian_environment(  # noqa: C901
        self,
        *,
        environment: Artifact | None = None,
        category: ArtifactCategory = ArtifactCategory.SYSTEM_TARBALL,
        codename: str = "bookworm",
        architecture: str = "amd64",
        variant: str | None = None,
        collection: Collection | None = None,
        user: User | None = None,
        variables: dict[str, Any] | None = None,
        workspace: Workspace | None = None,
    ) -> CollectionItem:
        """Create a debian build environment."""
        if workspace is None:
            workspace = self.get_default_workspace()
        if collection is None:
            collection = self.create_debian_environments_collection()
        if user is None:
            user = self.get_default_user()
        if variant is None and variables is not None:
            variant = variables.get("variant")
        if variant is None:
            variant = "apt"
        if environment is None:
            # Try looking up an existing environment
            lookup_string = (
                f"{collection.name}@{collection.category}"
                f"/match:codename={codename}"
                f":architecture={architecture}"
                # FIXME with this match fails when run twice in a row
                f":variant={variant or variables}"
            )
            if category == ArtifactCategory.SYSTEM_TARBALL:
                lookup_string += ":format=tarball"
            else:
                lookup_string += ":format=image"
            try:
                lookup_result = lookup_single(
                    lookup_string, workspace, user=user
                )
                assert lookup_result.collection_item is not None
                return lookup_result.collection_item
            except KeyError:
                pass
        if environment is None:
            # Lookup failed: create it
            data = {
                "codename": codename,
                "architecture": architecture,
                "variant": variant,
            }
            environment, _ = self.create_artifact(
                category=category,
                data={**data, "with_dev": True},
                workspace=workspace,
            )
        if variables is None:
            variables = {"backend": "unshare", "variant": variant}

        try:
            return CollectionItem.active_objects.get(
                parent_collection=collection,
                artifact=environment,
                data__contains=variables,
            )
        except CollectionItem.DoesNotExist:
            pass

        manager = collection.manager
        item = manager.add_artifact(environment, user=user, variables=variables)
        return item

    def create_sbuild_work_request(
        self,
        *,
        source: Artifact,
        architecture: str = "all",
        environment: Artifact,
        workflow: WorkRequest | None = None,
        workspace: Workspace | None = None,
    ) -> WorkRequest:
        """Create a sbuild work request."""
        task_data = SbuildData(
            input=SbuildInput(source_artifact=source.pk),
            host_architecture=architecture,
            environment=environment.pk,
            backend=BackendType.UNSHARE,
            build_components=[
                (
                    SbuildBuildComponent.ALL
                    if architecture == "all"
                    else SbuildBuildComponent.ANY
                )
            ],
        )
        if workflow:
            return workflow.create_child(
                task_name="sbuild",
                task_data=task_data.dict(),
                workflow_data=WorkRequestWorkflowData(
                    display_name=f"Build {architecture}",
                    step=f"build-{architecture}",
                ),
            )
        else:
            return self.create_work_request(
                task_name="sbuild",
                task_data=task_data.dict(),
                workspace=workspace or self.get_default_workspace(),
            )

    def simulate_package_build(
        self,
        source: Artifact,
        *,
        architecture: str = "all",
        workflow: WorkRequest | None = None,
        environment: Artifact | None = None,
        worker: Worker | None = None,
    ) -> WorkRequest:
        """Generate database objects as if a package build happened."""
        workspace = source.workspace
        name = source.data["name"]
        version = source.data["version"]

        if environment is None:
            environment_item = self.create_debian_environment()
            assert environment_item.artifact is not None
            environment = environment_item.artifact

        build_arch = environment.data["architecture"]

        # Create sbuild work request
        work_request = self.create_sbuild_work_request(
            source=source,
            architecture=architecture,
            environment=environment,
            workflow=workflow,
            workspace=workspace,
        )

        # Assign a worker
        if worker is None:
            worker = self.create_worker()
        work_request.assign_worker(worker)
        work_request.mark_running()

        with tempfile.TemporaryDirectory() as tempdir:
            workdir = Path(tempdir)

            # Add a build log
            buildlog_file = workdir / f"{name}_{version}_{build_arch}.build"
            buildlog_file.write_text(buildlog_file.name)
            buildlog = self.create_package_build_log(buildlog_file)
            buildlog_artifact, _ = self.create_artifact(
                paths={buildlog_file.name: buildlog_file.read_bytes()},
                category=buildlog.category,
                workspace=workspace,
                data=model_to_json_serializable_dict(buildlog.data),
                work_request=work_request,
                created_by=work_request.created_by,
                create_files=True,
            )
            self.create_artifact_relation(
                buildlog_artifact, source, ArtifactRelation.Relations.RELATES_TO
            )

            # Add the .deb
            binarypackage = self.create_binary_package(
                workdir, name=name, version=version, architecture=architecture
            )
            deb_file = next(iter(binarypackage.files.values()))
            binarypackage_artifact, _ = self.create_artifact(
                paths={deb_file.name: deb_file.read_bytes()},
                category=binarypackage.category,
                workspace=workspace,
                data=model_to_json_serializable_dict(binarypackage.data),
                work_request=work_request,
                created_by=work_request.created_by,
                create_files=True,
            )
            self.create_artifact_relation(
                binarypackage_artifact,
                source,
                ArtifactRelation.Relations.BUILT_USING,
            )
            self.create_artifact_relation(
                buildlog_artifact,
                binarypackage_artifact,
                ArtifactRelation.Relations.RELATES_TO,
            )

            # TODO: create the .changes file / Upload artifact
            # TODO:     adding the .dsc, the sources, the .deb
            # TODO: self.create_artifact_relation(
            # TODO:     upload,
            # TODO:     binarypackage_artifact,
            # TODO:     ArtifactRelation.Relations.EXTENDS,
            # TODO: )
            # TODO: self.create_artifact_relation(
            # TODO:     upload,
            # TODO:     binarypackage_artifact,
            # TODO:     ArtifactRelation.Relations.RELATES_TO,
            # TODO: )

            # Complete the task
            work_request.mark_completed(WorkRequest.Results.SUCCESS)

        return work_request
