# Copyright 2023 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 Web forms."""
import functools
import shutil
import tempfile
from datetime import timedelta
from pathlib import Path
from typing import Any, Generic, TYPE_CHECKING, TypeVar, cast

from django.core.exceptions import ValidationError
from django.db import transaction
from django.db.models import Model
from django.forms import (
    BaseForm,
    CharField,
    ChoiceField,
    ClearableFileInput,
    DateTimeField,
    DateTimeInput,
    Field,
    FileField,
    Form,
    IntegerField,
    ModelChoiceField,
    ModelForm,
    Textarea,
)

try:
    import pydantic.v1 as pydantic
except ImportError:
    import pydantic  # type: ignore

import yaml

# from debusine.client import LocalArtifact
from debusine.artifacts import LocalArtifact
from debusine.db.models import (
    Artifact,
    Collection,
    CollectionItem,
    File,
    FileInArtifact,
    Token,
    WorkRequest,
    Workspace,
)
from debusine.tasks import BaseTask
from debusine.tasks.models import TaskTypes
from debusine.web.utils import ui_prototype

if TYPE_CHECKING:
    ModelFormBase = ModelForm
else:
    # Django's ModelForm doesn't support generic types at run-time yet.
    class _ModelFormBase:
        def __class_getitem__(*args):
            return ModelForm

    ModelFormBase = _ModelFormBase


_KT = TypeVar("_KT")
_VT = TypeVar("_VT")


class DictWithCallback(dict[_KT, _VT], Generic[_KT, _VT]):
    r"""
    Dictionary that when setting a value it calls a method to process it.

    Usage: d = DictWithCallback(callback, \*dict_args, \*\*dict_kwargs)

    when doing:
    d["key"] = value

    Before setting the value it calls "callback" which can change, in place,
    the value. Then sets the value to the dictionary.

    When accessing the value (d["key"]) it return the value as it was modified
    by the callback.
    """

    def __init__(self, callback, *args, **kwargs):
        """Create the object."""
        self._callback = callback
        super().__init__(*args, **kwargs)

    def __setitem__(self, key, value):
        """Call self._callback(value) and set the value."""
        self._callback(value)
        super().__setitem__(key, value)


class BootstrapMixin:
    """
    Mixin that adjusts the CSS classes of form fields with Bootstrap's UI.

    This mixin is intended to be used in combination with Django's form classes.
    """

    def __init__(self, *args, **kwargs):
        """Initialize the mixin."""
        super().__init__(*args, **kwargs)

        self._adjust_bootstrap_classes_all_fields()

        self.fields = DictWithCallback[str, Field](
            self._adjust_bootstrap_for_field, self.fields
        )

    @staticmethod
    def _adjust_bootstrap_for_field(field: Field):
        """Adjust the CSS class for a field."""
        existing_class = field.widget.attrs.get("class", "")
        bootstrap_class = None

        if field.required:
            suffix = " *"
            if not (field.label_suffix or "").endswith(suffix):
                field.label_suffix = (field.label_suffix or "") + suffix

        if isinstance(field, ChoiceField):
            bootstrap_class = "form-select"
        elif isinstance(field, (CharField, FileField, DateTimeField)):
            bootstrap_class = "form-control"

        if bootstrap_class and bootstrap_class not in existing_class.split():
            field.widget.attrs["class"] = (
                f"{existing_class} {bootstrap_class}".strip()
            )

    def _adjust_bootstrap_classes_all_fields(self):
        """Adjust the CSS classes of form fields to be Bootstrap-compatible."""
        for field in self.fields.values():
            self._adjust_bootstrap_for_field(field)


class TokenForm(BootstrapMixin, ModelFormBase[Token]):
    """Form for creating or editing a token."""

    def __init__(self, *args, **kwargs):
        """Initialize TokenForm."""
        self.user = kwargs.pop("user")
        super().__init__(*args, **kwargs)

        if not self.instance.pk:
            # New instance (not loaded from the DB). Set defaults
            self.fields["enabled"].initial = True

    def save(self, commit=True):
        """Save TokenForm."""
        instance = super().save(commit=False)

        instance.user = self.user

        if commit:
            instance.save()

        return instance

    class Meta:
        model = Token
        fields = ["comment", "enabled"]

        labels = {
            "comment": "Comment",
        }


class WorkspaceChoiceField(ModelChoiceField):
    """ChoiceField for the workspaces: set the label and order by name."""

    def __init__(self, user, *args, **kwargs):
        """Set the queryset."""
        kwargs["queryset"] = Workspace.objects.order_by("name")

        if user is None:
            # Non-authenticated users can list only public workspaces
            kwargs["queryset"] = kwargs["queryset"].filter(public=True)

        super().__init__(*args, **kwargs)

    # TODO: Once https://github.com/typeddjango/django-stubs/pull/1889 is in
    # a released version, we can improve this by changing the base class to
    # ModelChoiceField[Workspace].
    def label_from_instance(self, obj: Model) -> str:
        """Return name of the workspace."""
        assert isinstance(obj, Workspace)
        return obj.name


class YamlMixin:
    """
    Mixin that that handles fields and validate/convert from YAML to a dict.

    Usage:
    In the class inheriting from YamlMixin:

    yaml_fields = ["task_data", "some_other_field"]
    """

    def __init__(self, *args, **kwargs):
        """For each self.yaml_fields: create the CharField."""
        super().__init__(*args, **kwargs)

        for name in self.yaml_fields:
            self.fields[name] = CharField(widget=Textarea, required=False)

            # This might need to change if it's possible to edit
            # a form loading data from the database
            self.fields[name].value = ""

            self.__setattr__(
                f"clean_{name}", functools.partial(self._clean_data_yaml, name)
            )

    def _clean_data_yaml(self, field_name: str) -> dict[Any, Any]:
        """Return dictionary representing the YAML input."""
        data_yaml = cast(BaseForm, self).cleaned_data[field_name]

        try:
            task_data = yaml.safe_load(data_yaml)
        except yaml.YAMLError as exc:
            raise ValidationError(f"Invalid YAML: {exc}")

        return {} if task_data is None else task_data


class WorkRequestForm(YamlMixin, BootstrapMixin, ModelFormBase[WorkRequest]):
    """Form for creating a Work Request."""

    yaml_fields = ["task_data"]

    def __init__(self, *args, **kwargs):
        """Initialize WorkRequestForm."""
        self.user = kwargs.pop("user")

        # self.field_order = ["workspace", "task_name", "data"]

        super().__init__(*args, **kwargs)

        self.fields["workspace"] = WorkspaceChoiceField(user=self.user)
        if self.user.is_superuser:
            tasks = BaseTask.task_names(TaskTypes.WORKER) + BaseTask.task_names(
                TaskTypes.SERVER
            )
        else:
            tasks = BaseTask.worker_task_names()
        self.fields["task_name"] = ChoiceField(
            choices=lambda: [[name, name] for name in sorted(tasks)]
        )

    def _get_task_type(self, name: str) -> TaskTypes:
        """Infer a task type from a task name."""
        if self.user.is_superuser and not BaseTask.is_worker_task(name):
            return TaskTypes.SERVER
        return TaskTypes.WORKER

    def full_clean(self):
        """Make sure task_type is filled."""
        if "task_name" in self.data:
            self.instance.task_type = self._get_task_type(
                self.data["task_name"]
            )
        super().full_clean()

    def save(self, commit=True):
        """Save the work request."""
        instance = super().save(commit=False)
        instance.created_by = self.user
        instance.task_type = self._get_task_type(instance.task_name)

        if commit:
            instance.save()

        return instance

    class Meta:
        model = WorkRequest
        fields = ["workspace", "task_name", "task_data"]


class MultipleFileInput(ClearableFileInput):
    """ClearableFileInput allowing to select multiple files."""

    allow_multiple_selected = True


class MultipleFileField(FileField):
    """
    FileField using the widget MultipleFileInput.

    Implementation as suggested by Django documentation:
    https://docs.djangoproject.com/en/4.2/topics/http/file-uploads/#uploading-multiple-files
    """

    def __init__(self, *args, **kwargs):
        """Initialize object: use MultipleFileInput() as a widget."""
        kwargs.setdefault("widget", MultipleFileInput())
        super().__init__(*args, **kwargs, allow_empty_file=True)

    def clean(self, data: list[Any] | tuple[Any, ...], initial=None):
        """Call super().clean() for each file."""  # noqa: D402
        single_file_clean = super().clean
        return [single_file_clean(file, initial) for file in data]


class DateTimeInputPicker(DateTimeInput):
    """DateTimeInput with browser's native picker."""

    input_type = "datetime-local"


class ArtifactForm(BootstrapMixin, YamlMixin, ModelFormBase[Artifact]):
    """Form for creating artifacts."""

    # Deliberately incompatible with BaseForm.files.
    files = MultipleFileField()  # type: ignore[assignment]
    category = ChoiceField()
    expiration_delay_in_days = IntegerField(
        min_value=0, initial=None, required=False
    )

    yaml_fields = ["data"]

    def __init__(self, *args: Any, **kwargs: Any) -> None:
        """Initialize object."""
        self.user = kwargs.pop("user")

        super().__init__(*args, **kwargs)

        self.fields["workspace"] = WorkspaceChoiceField(user=self.user)

        self.fields["category"].choices = [
            (artifact_category, artifact_category)
            for artifact_category in sorted(LocalArtifact.artifact_categories())
        ]

        # Populated on clean(), used on save()
        self._local_files: dict[str, Path] = {}
        self._temporary_directory = None

    @transaction.atomic
    def save(self, commit=True):
        """Create the artifact."""
        instance = super().save(commit=False)
        instance.created_by = self.user
        if self.cleaned_data["expiration_delay_in_days"] is not None:
            instance.expiration_delay = timedelta(
                days=self.cleaned_data["expiration_delay_in_days"]
            )

        if commit:
            instance.save()

            for file in self.cleaned_data["files"]:
                local_file_path = self._local_files[file.name]
                # Add file to the store
                file_obj = File.from_local_path(local_file_path)

                file_backend = (
                    instance.workspace.default_file_store.get_backend_object()
                )
                file_backend.add_file(local_file_path, fileobj=file_obj)

                # Add file to the artifact
                FileInArtifact.objects.create(
                    artifact=instance, path=file.name, file=file_obj
                )
                instance.files.add(file_obj)

        return instance

    def clean(self):
        """
        Create a LocalArtifact model and validate it.

        :raise ValidationError: if the LocalArtifact model is not valid.
        """
        cleaned_data = super().clean()

        artifact_category = cleaned_data["category"]

        SubLocalArtifact = LocalArtifact.class_from_category(artifact_category)

        self._temporary_directory = tempfile.TemporaryDirectory(
            prefix="debusine-form-artifact"
        )
        self._local_files = {}
        for file in cleaned_data["files"]:
            file_path = Path(self._temporary_directory.name) / file.name

            with file_path.open("wb") as local_file:
                shutil.copyfileobj(file.file, local_file)
            file.file.close()

            self._local_files[file.name] = file_path

        # If adding any new fields in this LocalArtifact fields,
        # make sure that the form has a field with the same name.
        # If not, adjust the code handling the ValidationError.
        sub_local_artifact_kwargs = {
            "category": artifact_category,
            "data": cleaned_data.get("data"),
            "files": self._local_files,
        }

        try:
            SubLocalArtifact(**sub_local_artifact_kwargs)
        except pydantic.ValidationError as exc:
            for error in exc.errors():
                field_name = error["loc"][0]
                # This assumes that the fields that can raise ValidationErrors
                # in the LocalArtifact have a field in the form with the same
                # name. If some day there are fields with different names
                # need to add a mapping or add errors via
                # self.add_error(None, ...) which adds the errors on the
                # top of the form.
                self.add_error(field_name, error["msg"])

        return cleaned_data

    class Meta:
        model = Artifact
        fields = ["category", "workspace", "files", "data"]


@ui_prototype
class CollectionSearchForm(Form):
    """Form for collection search fields."""

    category = ChoiceField(required=False)
    name = CharField(required=False)

    def __init__(self, *args, instance: Collection, **kwargs):
        """Initialize category choices from the database."""
        super().__init__(*args, **kwargs)
        self.instance = instance
        choices = [("", "All")]
        choices.extend(
            (name, name)
            for name in CollectionItem.objects.filter(
                parent_collection=self.instance
            )
            .values_list("category", flat=True)
            .distinct()
            .order_by("category")
        )
        cast(ChoiceField, self.fields["category"]).choices = choices
