# 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.
"""Input fields for DB-side tasks."""

from typing import Unpack, override

from debusine.artifacts.models import CollectionCategory
from debusine.client.models import LookupChildType
from debusine.db.tasks.models import DBArtifactInfo, DBInputCollectionSingle
from debusine.db.tasks.resolver import DBFieldResolver
from debusine.tasks import (
    TaskConfigError,
    ensure_artifact_categories,
    ensure_collection_category,
)
from debusine.tasks.inputs import FieldResolver, serializers
from debusine.tasks.inputs.inputs import (
    BaseArtifactInput,
    DataFieldInput,
    DataFieldInputKwargs,
)
from debusine.tasks.models import BaseTaskData, LookupSingle


class ArtifactInput(BaseArtifactInput[LookupSingle, DBArtifactInfo]):
    """One single input artifact specified in a task data field."""

    serializer = serializers.PydanticSerializer(DBArtifactInfo)

    @override
    def resolve(self, resolver: FieldResolver) -> DBArtifactInfo:
        assert isinstance(resolver, DBFieldResolver)
        lookup = resolver.get_field(self)
        assert lookup is not None
        try:
            artifact = resolver.work_request.lookup_single(
                lookup=lookup,
                expect_type=LookupChildType.ARTIFACT,
            ).artifact
        except KeyError:
            msg = f"Artifact {lookup!r} in {resolver.work_request.workspace}"
            msg += f" not found for field {self.field!r}"
            raise TaskConfigError(msg)

        if self.categories is not None:
            ensure_artifact_categories(
                configuration_key=self.field,
                category=artifact.category,
                expected=self.categories,
            )

        return DBArtifactInfo.from_artifact(
            artifact, lookup=lookup, label=self.field
        )

    @override
    def get_artifact_ids(self, value: DBArtifactInfo) -> list[int]:
        return [value.artifact_id]


class CollectionInput(DataFieldInput[LookupSingle, DBInputCollectionSingle]):
    """One single input collection specified in a task data field."""

    serializer = serializers.PydanticSerializer(DBInputCollectionSingle)

    def __init__(
        self,
        *,
        category: CollectionCategory | None = None,
        **kwargs: Unpack[DataFieldInputKwargs],
    ) -> None:
        """
        Look up artifacts enforcing their type.

        :param field: dot-separated list of lookup names in task data
        :param categories: list of acceptable artifact categories
        """
        super().__init__(**kwargs)
        self.category = category

    @override
    def resolve(self, resolver: FieldResolver) -> DBInputCollectionSingle:
        assert isinstance(resolver, DBFieldResolver)
        lookup = resolver.get_field(self)
        assert lookup is not None
        # Look up the collection
        try:
            result = resolver.work_request.lookup_single(
                lookup=lookup,
                default_category=self.category,
                expect_type=LookupChildType.COLLECTION,
            )
        except KeyError:
            msg = f"Collection {lookup!r} in {resolver.work_request.workspace}"
            if self.category is not None:
                msg += f" with category '{self.category}'"
            msg += f" not found for field {self.field!r}"
            raise TaskConfigError(msg)

        if (
            self.category is not None
            and result.collection.category != self.category
        ):
            ensure_collection_category(
                configuration_key=self.field,
                category=result.collection.category,
                expected=self.category,
            )

        return DBInputCollectionSingle.from_collection(
            result.collection, lookup=lookup, label=self.field
        )


class CollectionInSameWorkspace(CollectionInput):
    """A collection in the same workspace as the work request."""

    @override
    def resolve(self, resolver: FieldResolver) -> DBInputCollectionSingle:
        assert isinstance(resolver, DBFieldResolver)
        info = super().resolve(resolver)
        if info.collection.workspace != resolver.work_request.workspace:
            raise TaskConfigError(
                f"{info.lookup} is in workspace {info.collection.workspace},"
                f" not {resolver.work_request.workspace}"
            )
        return info


class TaskConfigurationCollection(
    DataFieldInput[LookupSingle, DBInputCollectionSingle | None]
):
    """Lookup the task configuration collection to use."""

    serializer = serializers.OptionalPydanticSerializer(DBInputCollectionSingle)

    @override
    def resolve(
        self, resolver: FieldResolver
    ) -> DBInputCollectionSingle | None:
        assert isinstance(resolver, DBFieldResolver)
        if (lookup := resolver.get_field(self)) is None:
            return None

        # Look up the debusine:task-configuration collection
        try:
            lookup_result = resolver.work_request.lookup_single(
                lookup=lookup,
                default_category=CollectionCategory.TASK_CONFIGURATION,
                expect_type=LookupChildType.COLLECTION,
            )
        except KeyError:
            if lookup != BaseTaskData.DEFAULT_TASK_CONFIGURATION_LOOKUP:
                raise
            else:
                return None

        return DBInputCollectionSingle.from_collection(
            lookup_result.collection, lookup=lookup, label=self.field
        )
