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

"""The collection manager for debusine:task-configuration collections."""

import datetime as dt
from collections.abc import Generator, Iterable
from typing import Any, NamedTuple, override

from django.db import IntegrityError
from django.db.models import Case, JSONField, Q, Value, When
from django.db.models.functions import Coalesce
from django.utils import timezone

from debusine.artifacts.models import (
    BareDataCategory,
    CollectionCategory,
    DebusineTaskConfiguration,
)
from debusine.db.models import Collection, CollectionItem, User, WorkRequest
from debusine.db.orm_utils import CompareBoolean, JSONBTypeOf
from debusine.db.tasks import DBTask
from debusine.server.collections.base import (
    CollectionManagerInterface,
    ItemAdditionError,
)
from debusine.server.tags import (
    make_task_provided_tagset,
    make_task_required_tagset,
)
from debusine.tasks import TaskConfigError
from debusine.utils.tag_utils import DerivationRuleParser


class DebusineTaskConfigurationManager(CollectionManagerInterface):
    """Manage collection of category debusine:task-configuration."""

    COLLECTION_CATEGORY = CollectionCategory.TASK_CONFIGURATION
    VALID_BARE_DATA_CATEGORIES = frozenset(
        {BareDataCategory.TASK_CONFIGURATION}
    )

    @override
    def do_add_bare_data(
        self,
        category: BareDataCategory,
        *,
        user: User,
        workflow: WorkRequest | None = None,
        data: dict[str, Any] | None = None,
        name: str | None = None,  # noqa: U100
        replace: bool = False,
        created_at: dt.datetime | None = None,
        replaced_by: CollectionItem | None = None,
    ) -> CollectionItem:
        if data is None:
            raise ItemAdditionError(
                f"Adding to {CollectionCategory.TASK_CONFIGURATION} requires "
                f"data"
            )
        if name is not None:
            raise ItemAdditionError(
                f"Cannot use an explicit item name when adding to "
                f"{CollectionCategory.TASK_CONFIGURATION}"
            )

        config = DebusineTaskConfiguration(**data)
        name = config.name()

        if replace:
            self.remove_items_by_name(
                name=name,
                child_types=[CollectionItem.Types.BARE],
                user=user,
                workflow=workflow,
            )

        try:
            return CollectionItem.objects.create_from_bare_data(
                category,
                parent_collection=self.collection,
                name=name,
                data=data,
                created_at=created_at,
                created_by_user=user,
                created_by_workflow=workflow,
                replaced_by=replaced_by,
            )
        except IntegrityError as exc:
            raise ItemAdditionError(str(exc))

    @override
    def do_remove_item(
        self,
        item: CollectionItem,
        *,
        user: User | None = None,
        workflow: WorkRequest | None = None,
        remove_associated: bool = False,  # noqa: ARG002
    ) -> None:
        item.removed_by_user = user
        item.removed_by_workflow = workflow
        item.removed_at = timezone.now()
        item.save()


class ConfigurationLookup:
    """Machinery to look up task configuration items."""

    def __init__(self, collection: Collection) -> None:
        """Set the collection to use for lookups."""
        self.collection = collection
        self.templates = self.collection.child_items.active().filter(
            child_type=CollectionItem.Types.BARE
        )
        self.candidates = (
            self.collection.child_items.active()
            .filter(child_type=CollectionItem.Types.BARE)
            .annotate(
                # Match via postgres only when the field is an array
                tags_provided=Case(
                    When(
                        CompareBoolean(
                            JSONBTypeOf("data__tags_provided"),
                            "=",
                            Value("array"),
                        ),
                        then="data__tags_provided",
                    ),
                    default=Value("[]"),
                    output_field=JSONField(),
                ),
                # Match via postgres only when the field is an array
                tags_required=Case(
                    When(
                        CompareBoolean(
                            JSONBTypeOf("data__tags_required"),
                            "=",
                            Value("array"),
                        ),
                        then="data__tags_required",
                    ),
                    default=Value("[]"),
                    output_field=JSONField(),
                ),
                subject=Coalesce(
                    "data__subject", Value("null"), output_field=JSONField()
                ),
                context=Coalesce(
                    "data__context", Value("null"), output_field=JSONField()
                ),
            )
        )
        self.task_tags_provided: set[str]
        self.task_tags_required: set[str]

    def set_task(self, task: DBTask[Any, Any]) -> None:
        """Set lookup constraints from a task instance."""
        self.set_task_type_name_constraints(task.TASK_TYPE, task.name)
        subject, context = task.get_task_configuration_subject_context()
        self.set_subject_context_constraints(subject, context)

        # Compute the pre-configuration scheduler tags
        config_provided = make_task_provided_tagset()
        config_required = make_task_required_tagset()
        provided, required = task.compute_scheduler_tags()
        config_provided.add_candidates(provided)
        config_provided.finalize()
        config_required.add_candidates(required)
        config_required.finalize()

        # Filter candidate entries to those matching the given tags
        self.set_tags_provided(config_provided.tags)
        self.set_tags_required(config_required.tags)

    def set_task_type_name_constraints(
        self, task_type: str, task_name: str
    ) -> None:
        """Set constraints on task type and name."""
        self.candidates = self.candidates.filter(
            data__task_type=task_type,
            data__task_name=task_name,
        )

    def set_subject_context_constraints(
        self, subject: str | None, context: str | None
    ) -> None:
        """Set constraints on subject and context fields."""
        if subject is None:
            self.candidates = self.candidates.filter(subject=None)
        else:
            self.candidates = self.candidates.filter(
                Q(subject=None) | Q(subject=subject)
            )
        if context is None:
            self.candidates = self.candidates.filter(context=None)
        else:
            self.candidates = self.candidates.filter(
                Q(context=None) | Q(context=context)
            )

    def set_tags_provided(self, tags: set[str]) -> None:
        """Set the set of task-provided tags used for filtering."""
        self.candidates = self.candidates.filter(
            tags_provided__contained_by=sorted(tags)
        )
        self.task_tags_provided = tags

    def set_tags_required(self, tags: set[str]) -> None:
        """Set the set of task-required tags used for filtering."""
        self.candidates = self.candidates.filter(
            tags_required__contained_by=sorted(tags)
        )
        self.task_tags_required = tags

    @staticmethod
    def sort_key(
        value: tuple[DebusineTaskConfiguration, int],
    ) -> tuple[str, int, int, int]:
        """Sorting key for task configuration enties."""
        # Sorting by presence of subject and context
        sc: int
        match value[0].subject, value[0].context:
            case None, None:
                sc = 0
            case None, str():
                sc = 1
            case str(), None:
                sc = 2
            case _:
                sc = 3
        return (
            # Sort by path and position
            value[0].path or "",
            value[0].position or 0,
            # Then by presence of subject and context
            sc,
            # Lastly by item ID, as a tie breaker
            value[1],
        )

    def iter_matching_entries(self) -> Generator[DebusineTaskConfiguration]:
        """Filter configuration entries and iterate them in order."""
        parser = DerivationRuleParser()
        entries: list[tuple[DebusineTaskConfiguration, int]] = []

        # Instantiate entries
        for item in self.candidates:
            try:
                parsed = DebusineTaskConfiguration.model_validate(item.data)
            except ValueError as exc:
                raise TaskConfigError(
                    f"debusine:task-configuration item {item.id}"
                    " does not validate",
                    exc,
                )
            if isinstance(parsed.tags_provided, str):
                if not parser.as_function(
                    parsed.tags_provided, "task configuration"
                )(self.task_tags_provided):
                    continue
            if isinstance(parsed.tags_required, str):
                if not parser.as_function(
                    parsed.tags_required, "task configuration"
                )(self.task_tags_required):
                    continue
            entries.append((parsed, item.id))

        # Sort items
        entries.sort(key=self.sort_key)

        yield from (e[0] for e in entries)

    def lookup_template(self, name: str) -> DebusineTaskConfiguration | None:
        """Lookup a DebusineTaskConfiguration template by name."""
        try:
            item = self.templates.get(name=f"template:{name}")
        except CollectionItem.DoesNotExist:
            return None

        try:
            return DebusineTaskConfiguration(**item.data)
        except ValueError as exc:
            raise TaskConfigError(
                "debusine:task-configuration item does not validate", exc
            )

    def lookup_templates(
        self,
        entry: DebusineTaskConfiguration,
        stack: tuple[str, ...] = (),
    ) -> Generator[DebusineTaskConfiguration]:
        """
        Generate Task Configuration objects for a list of template names.

        :param config_collection: collection to use for lookups
        :param entry: entry with the list of template names to lookup
        :param stack: template names we are currently resolving, used to detect
                      circular loops
        """
        for name in entry.use_templates:
            if name in stack:
                raise TaskConfigError(
                    f"{entry.name()}: template lookup cycle detected: "
                    + "→".join((*stack, name))
                )
            if (template_entry := self.lookup_template(name)) is None:
                raise TaskConfigError(
                    f"{entry.name()} references missing template {name!r}"
                )
            if template_entry.use_templates:
                yield from self.lookup_templates(template_entry, (*stack, name))
            yield template_entry

    def iter_items(self) -> Generator[DebusineTaskConfiguration]:
        """
        Lookup DebusineTaskConfiguration objects for an item name.

        The resulting sequence also contains referenced template elements.
        """
        for entry in self.iter_matching_entries():
            yield entry

            # Look up referenced templates
            yield from self.lookup_templates(entry)


class Changes(NamedTuple):
    """Computed changes to apply to task configuration."""

    #: Values to set if not already set
    default_values: dict[str, Any]
    #: Values to set regardless
    override_values: dict[str, Any]
    #: Extra tags to provide, with the ``workspace`` provenance
    provide_tags: set[str]
    #: Extra tags to require, with the ``workspace`` provenance
    require_tags: set[str]


def build_configuration(
    items: Iterable[DebusineTaskConfiguration],
) -> Changes:
    """Turn config items into configuration to apply to task data."""
    default_values: dict[str, Any] = {}
    override_values: dict[str, Any] = {}
    provide_tags: set[str] = set()
    require_tags: set[str] = set()
    locked_values = set()

    for config_item in items:
        # Drop all the entries referenced in `delete_values` (except
        # locked values)
        for key in config_item.delete_values:
            if key in locked_values:
                continue
            default_values.pop(key, None)
            override_values.pop(key, None)

        # Merge the default/override values in the response
        # (except locked values)
        for key, value in config_item.default_values.items():
            if key in locked_values:
                continue
            default_values[key] = value
        for key, value in config_item.override_values.items():
            if key in locked_values:
                continue
            override_values[key] = value

        # Update the set of locked values
        locked_values.update(config_item.lock_values)

        # Collected extra provided and required tags
        provide_tags.update(config_item.provide_tags)
        require_tags.update(config_item.require_tags)

    return Changes(
        default_values=default_values,
        override_values=override_values,
        provide_tags=provide_tags,
        require_tags=require_tags,
    )


def apply_configuration(
    task_data: dict[str, Any],
    config_collection: Collection,
    task: DBTask[Any, Any],
) -> Changes:
    """
    Apply task configuration to task_data.

    :param task_data: the task data dict to configure, modified in place
    :param config_collection: the task configuration collection to use
    :param task: the task to use for task configuration lookups
    :returns: the set of changes that have been applied
    """
    # Build the base setup for looking up configuration entries
    lookup = ConfigurationLookup(config_collection)
    lookup.set_task(task)

    # Generate a list of matching items and their used template entries
    config_items = lookup.iter_items()

    # Turn items into a configuration to apply
    changes = build_configuration(config_items)

    # Apply default values (add missing values, but also replace explicit
    # None values)
    for k, v in changes.default_values.items():
        if task_data.get(k) is None:
            task_data[k] = v

    # Apply overrides
    task_data.update(changes.override_values)

    # Store the ID of the task configuration collection that was used
    task_data["task_configuration"] = config_collection.id

    return changes
