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

"""Models used by DB-side debusine tasks."""

from collections.abc import Sequence
from typing import Self

import pydantic

from debusine.client.models import LookupResultType
from debusine.db.models import Artifact, Collection, CollectionItem
from debusine.db.models.lookups import LookupResult
from debusine.tasks.models import (
    BaseLabeledTaskInput,
    BaseTaskInputArtifact,
    InputArtifactMultiple,
    InputArtifactSingle,
    InputCollectionSingle,
    LookupDict,
    LookupSingle,
)


class DBBaseTaskInputArtifact(BaseTaskInputArtifact):
    """BaseTaskInputArtifact caching a DB artifact."""

    #: Cached Artifact
    _artifact: Artifact | None = None

    @property
    def artifact(self) -> Artifact:
        """Return the cached artifact, or look it up by ID."""
        if self.artifact_id is None:
            raise AssertionError(
                "cannot access the Artifact of a task input without an ID"
            )
        if self._artifact is None:
            self._artifact = Artifact.objects.get(pk=self.artifact_id)
        return self._artifact

    @classmethod
    def from_artifact(cls, artifact: Artifact) -> Self:
        """Instantiate from an Artifact."""
        result = cls(
            artifact_id=artifact.id,
            category=artifact.category,
            data=artifact.create_data(),
        )
        result._artifact = artifact
        return result


class DBArtifactInfo(InputArtifactSingle):
    """Information about an artifact."""

    #: Cached Artifact
    _artifact: Artifact | None = None

    @property
    def artifact(self) -> Artifact:
        """Return the cached artifact, or look it up by ID."""
        if self.artifact_id is None:
            raise AssertionError(
                "cannot access the Artifact of a task input without an ID"
            )
        if self._artifact is None:
            self._artifact = Artifact.objects.get(pk=self.artifact_id)
        return self._artifact

    @classmethod
    def from_artifact(
        cls, artifact: Artifact, *, lookup: LookupSingle, label: str
    ) -> Self:
        """Instantiate from an Artifact."""
        result = cls(
            lookup=lookup,
            label=label,
            artifact_id=artifact.id,
            category=artifact.category,
            data=artifact.create_data(),
        )
        result._artifact = artifact
        return result


class DBInputArtifactMultiple(InputArtifactMultiple):
    """Represent declared multiple input artifacts for a task."""

    artifacts: Sequence[DBBaseTaskInputArtifact] = pydantic.Field(
        default_factory=list
    )


class DBInputCollectionSingle(InputCollectionSingle):
    """Information about a collection."""

    #: Cached Collection
    _collection: Collection | None = None

    @property
    def collection(self) -> Collection:
        """Return the cached collection, or look it up by ID."""
        if self._collection is None:
            self._collection = Collection.objects.get(pk=self.collection_id)
        return self._collection

    @classmethod
    def from_collection(
        cls, collection: Collection, *, lookup: LookupSingle, label: str
    ) -> Self:
        """Instantiate from a Collection."""
        result = cls(
            lookup=lookup,
            label=label,
            collection_id=collection.id,
            scope_name=collection.workspace.scope.name,
            workspace_name=collection.workspace.name,
            category=collection.category,
            name=collection.name,
            data=collection.data,
        )
        result._collection = collection
        return result


class DBInputLookupResult(BaseLabeledTaskInput):
    """
    Information about a generic lookup result.

    This may be a bare data item, an artifact, or a collection.
    """

    result_type: LookupResultType
    lookup: LookupSingle | LookupDict
    parent_collection_lookup: LookupSingle | None = None
    collection_item_id: int | None = None
    artifact_id: int | None = None
    collection_id: int | None = None

    _result: LookupResult | None = None

    @property
    def result(self) -> LookupResult:
        """Return the cached lookup result, or reconstruct it."""
        if self._result is None:
            self._result = LookupResult.create(
                result_type=self.result_type,
                lookup=self.lookup,
                parent_collection_lookup=self.parent_collection_lookup,
                collection_item=(
                    None
                    if self.collection_item_id is None
                    else CollectionItem.objects.get(pk=self.collection_item_id)
                ),
                artifact=(
                    None
                    if self.artifact_id is None
                    else Artifact.objects.get(pk=self.artifact_id)
                ),
                collection=(
                    None
                    if self.collection_id is None
                    else Collection.objects.get(pk=self.collection_id)
                ),
            )
        return self._result

    @classmethod
    def from_lookup_result(cls, result: LookupResult, *, label: str) -> Self:
        """Instantiate from a LookupResult."""
        instance = cls(
            label=label,
            result_type=result.result_type,
            lookup=result.lookup,
            parent_collection_lookup=result.parent_collection_lookup,
            collection_item_id=(
                None
                if result.collection_item is None
                else result.collection_item.id
            ),
            artifact_id=None if result.artifact is None else result.artifact.id,
            collection_id=(
                None if result.collection is None else result.collection.id
            ),
        )
        instance._result = result
        return instance
