# 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-admin command to manage workers."""

import io
import os
import shutil
import subprocess
import tempfile
from collections.abc import Callable
from typing import Any, NoReturn, cast

import yaml
from django.core.management import CommandError, CommandParser
from django.db import transaction

from debusine.db.models import Token, WorkRequest, Worker
from debusine.django.management.debusine_base_command import DebusineBaseCommand
from debusine.server.management.management_utils import Workers
from debusine.tasks.models import WorkerType


class WorkerStaticMetadataEditor:
    """Edit Worker static metadata."""

    def __init__(
        self,
        worker: Worker,
        yaml_file: str | None,
        stdout: io.TextIOBase,
        stderr: io.TextIOBase,
    ) -> None:
        """
        Initialize fields of WorkerStaticMetadataEditor.

        :param worker: worker that is going to be edited
        :param yaml_file: edit() method will read the metadata from it.
          If it is None it will load the metadata from the database.
        :param stdout: stdout file.
        :param stderr: stderr file.
        """
        self._worker = worker
        self._yaml_file = yaml_file
        self._stdout = stdout
        self._stderr = stderr

    def edit(self) -> bool:
        """Edit the worker's static_metadata."""
        if self._yaml_file:
            metadata = self._read_yaml_file(self._yaml_file)
        else:
            metadata = self._edit_worker_metadata(
                self._worker.static_metadata, self._stdout, self._stderr
            )

        if metadata is not None:
            self._worker.static_metadata = metadata
            self._worker.save()
            self._stdout.write(
                f"debusine: metadata set for {self._worker.name}"
            )
            return True
        else:
            return False

    def _read_yaml_file(self, yaml_path: str) -> dict[Any, Any] | None:
        """
        Read yaml_path and returns it as dictionary (or None if invalid).

        :param yaml_path: parses it. Be aware:
          * If it's an empty object returns {}.

          * If it does not contain a dictionary structure: writes a message to
            self.stderr and returns None.

          * If it cannot be opened writes a message to self.stderr and
            it raises SystemExit(3).
        """
        try:
            with open(yaml_path) as yaml_file:
                try:
                    contents = yaml.safe_load(yaml_file)
                except (yaml.YAMLError, yaml.scanner.ScannerError) as exc:
                    self._stderr.write(f"Invalid YAML: {exc}")
                    return None
        except OSError as exc:
            self._stderr.write(
                f"Error: cannot open worker configuration file "
                f"'{yaml_path}': {exc}\n"
            )
            raise SystemExit(3)

        if contents is None:
            # Empty YAML is normalized to an empty dictionary
            contents = {}
        elif not isinstance(contents, dict):
            self._stderr.write("Worker metadata must be a dictionary or empty")
            return None

        assert isinstance(contents, dict)
        return contents

    @staticmethod
    def _input() -> str:
        return input()  # pragma: no cover

    def _edit_worker_metadata(
        self,
        metadata: dict[str, Any] | None,
        stdout: io.TextIOBase,
        stderr: io.TextIOBase,
    ) -> dict[str, Any] | None:
        edit_metadata = tempfile.NamedTemporaryFile(
            prefix="debusine-worker-edit-metadata-",
            suffix=".tmp.yaml",
            delete=False,
        )
        edit_metadata.close()

        original_metadata = tempfile.NamedTemporaryFile(
            prefix="debusine-worker-edit-metadata-",
            suffix=".yaml",
            delete=False,
        )
        original_metadata.close()

        with open(original_metadata.name, "w") as original_yaml_file:
            yaml.safe_dump(data=metadata, stream=original_yaml_file)

        shutil.copy2(original_metadata.name, edit_metadata.name)

        while True:
            self._open_editor(edit_metadata.name)

            metadata = self._read_yaml_file(edit_metadata.name)

            if metadata is not None:
                os.unlink(edit_metadata.name)
                os.unlink(original_metadata.name)
                return metadata
            else:
                stderr.write("Error reading metadata.\n")
                stdout.write("Do you want to retry the same edit? (y/n)\n")
                answer = self._input()

                if answer != "y":
                    os.unlink(original_metadata.name)
                    stdout.write(
                        f"debusine: edits left in {edit_metadata.name}"
                    )
                    return None

    @staticmethod
    def _open_editor(file_path: str) -> None:
        # sensible-editor is in sensible-utils package
        subprocess.run(['sensible-editor', file_path])  # pragma: no cover


class Command(DebusineBaseCommand):
    """Command to manage workers."""

    help = "Manage workers."

    def add_arguments(self, parser: CommandParser) -> None:
        """Add CLI arguments for the worker command."""
        subparsers = parser.add_subparsers(dest="action", required=True)

        create = subparsers.add_parser(
            "create", help="Create a worker and print its activation token."
        )
        create.add_argument(
            "--worker-type",
            choices=["external", "signing"],
            type=WorkerType,
            default=WorkerType.EXTERNAL,
        )
        create.add_argument(
            "fqdn", help="Fully-qualified domain name of the worker"
        )

        enable = subparsers.add_parser("enable", help="Enable a worker.")
        enable.add_argument(
            "--worker-type",
            choices=["external", "signing"],
            type=WorkerType,
            default=WorkerType.EXTERNAL,
        )
        enable.add_argument("worker", help="Name of the worker to enable")

        disable = subparsers.add_parser(
            "disable",
            help=(
                "Disable a worker. This disables the worker's token (so the "
                "worker cannot communicate with the server) and de-assigns "
                "work requests for the worker, which will be assigned to "
                "another worker. No attempt is made to stop any worker's "
                "operations, but the server will not accept any results."
            ),
        )
        disable.add_argument(
            "--worker-type",
            choices=["external", "signing"],
            type=WorkerType,
            default=WorkerType.EXTERNAL,
        )
        disable.add_argument("worker", help="Name of the worker to disable")

        edit_metadata = subparsers.add_parser(
            "edit_metadata", help="Edit a worker's metadata."
        )
        edit_metadata.add_argument(
            "--set",
            help=(
                "Filename with the metadata in YAML"
                " (note: its contents will be displayed to users in the web UI)"
            ),
            metavar="PATH",
        )
        edit_metadata.add_argument(
            "worker", help="Name of the worker whose metadata should be edited"
        )

        list_ = subparsers.add_parser("list", help="List workers.")
        list_.add_argument(
            '--yaml', action="store_true", help="Machine readable YAML output"
        )

    def get_worker(self, name_or_token: str, worker_type: WorkerType) -> Worker:
        """Get a worker by name or token, checking its type."""
        worker = Worker.objects.get_worker_or_none(name_or_token)
        if worker is None:
            worker = Worker.objects.get_worker_by_token_key_or_none(
                name_or_token
            )

        if worker is None:
            raise CommandError("Worker not found", returncode=3)
        if worker.worker_type != worker_type:
            raise CommandError(
                f'Worker "{worker.name}" is of type "{worker.worker_type}", '
                f'not "{worker_type}"',
                returncode=4,
            )

        return worker

    def handle_create(self, *args: Any, **options: Any) -> NoReturn:
        """Create a worker and print its activation token."""
        activation_token = Token.objects.create_worker_activation()
        Worker.objects.create_with_fqdn(
            options["fqdn"],
            activation_token=activation_token,
            worker_type=options["worker_type"],
        )
        print(activation_token.key)

        raise SystemExit(0)

    def handle_enable(self, *args: Any, **options: Any) -> NoReturn:
        """Enable a worker."""
        worker = self.get_worker(options["worker"], options["worker_type"])
        # By this point the worker cannot be a Celery worker, and a database
        # constraint ensures that all non-Celery workers have a token.
        assert worker.token is not None

        worker.token.enable()

        raise SystemExit(0)

    def handle_disable(self, *args: Any, **options: Any) -> NoReturn:
        """Disable a worker."""
        worker = self.get_worker(options["worker"], options["worker_type"])
        # By this point the worker cannot be a Celery worker, and a database
        # constraint ensures that all non-Celery workers have a token.
        assert worker.token is not None

        with transaction.atomic():
            worker.token.disable()
            for work_request in worker.assigned_work_requests.filter(
                status__in={
                    WorkRequest.Statuses.PENDING,
                    WorkRequest.Statuses.RUNNING,
                }
            ):
                work_request.de_assign_worker()

        raise SystemExit(0)

    def handle_edit_metadata(self, *args: Any, **options: Any) -> NoReturn:
        """Edit a worker's metadata."""
        worker_name = options["worker"]

        try:
            worker = Worker.objects.get(name=worker_name)
        except Worker.DoesNotExist:
            raise CommandError(
                f'Error: worker "{worker_name}" is not registered\n'
                f'Use the command `debusine-admin worker list` to list the '
                f'existing workers',
                returncode=3,
            )

        worker_editor = WorkerStaticMetadataEditor(
            worker, options["set"], self.stdout, self.stderr
        )
        if worker_editor.edit():
            raise SystemExit(0)
        else:
            raise SystemExit(3)

    def handle_list(self, *args: Any, **options: Any) -> NoReturn:
        """List the workers."""
        workers = Worker.objects.all()
        Workers(options["yaml"]).print(workers, self.stdout)
        raise SystemExit(0)

    def handle(self, *args: Any, **options: Any) -> NoReturn:
        """Dispatch the requested action."""
        func = cast(
            Callable[..., NoReturn],
            getattr(self, f"handle_{options['action']}", None),
        )
        func(*args, **options)
