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

"""Lookups of items in collections."""

import re
from collections.abc import Sequence
from dataclasses import dataclass
from typing import Literal, assert_never, overload

from django.db.models import Q, TextField
from django.db.models.fields.json import KeyTextTransform
from django.db.models.functions import Cast

from debusine.artifacts.models import CollectionCategory
from debusine.db.models import (
    Artifact,
    Collection,
    CollectionItem,
    User,
    WorkRequest,
    Workspace,
)
from debusine.tasks.models import (
    LookupChildType,
    LookupDict,
    LookupMultiple,
    LookupSingle,
    parse_lookup_string_segments,
)


@dataclass
class LookupResult:
    """
    The result of a collection item lookup.

    None of :py:class:`Artifact`, :py:class:`Collection`, or
    :py:class:`CollectionItem` are entirely suitable return types: a lookup
    that returns a bare item needs to return the collection item, but in the
    case of a single-segment lookup there's no collection item available,
    only an artifact or collection.

    The least bad option seems to be to define a custom return type
    returning all the different things callers might need.
    """

    result_type: CollectionItem.Types
    collection_item: CollectionItem | None = None
    artifact: Artifact | None = None
    collection: Collection | None = None


class LookupResultArtifact(LookupResult):
    """
    A collection item lookup containing an artifact.

    Used to assist type annotations.
    """

    result_type: Literal[CollectionItem.Types.ARTIFACT]
    artifact: Artifact


class LookupResultCollection(LookupResult):
    """
    A collection item lookup containing a collection.

    Used to assist type annotations.
    """

    result_type: Literal[CollectionItem.Types.COLLECTION]
    collection: Collection


def _lookup_single_first(
    lookup: str,
    workspace: Workspace,
    *,
    user: User | None,
    default_category: CollectionCategory | None = None,
    workflow_root: WorkRequest | None = None,
) -> LookupResult:
    """
    Resolve the first segment of a string lookup.

    This is handled differently from the other segments, because the lookup
    doesn't happen in the context of a collection.
    """
    try:
        if m := re.match(r"^([0-9]+)@artifacts$", lookup):
            workspace_filter = Q(workspace=workspace) | Q(
                workspace__public=True
            )
            visible_artifacts = Artifact.objects.filter(workspace_filter)
            return LookupResult(
                result_type=CollectionItem.Types.ARTIFACT,
                artifact=visible_artifacts.get(id=int(m.group(1))),
            )
        elif m := re.match(r"^([0-9]+)@collections$", lookup):
            workspace_filter = Q(workspace=workspace) | Q(
                workspace__public=True
            )
            visible_collections = Collection.objects.filter(workspace_filter)
            return LookupResult(
                result_type=CollectionItem.Types.COLLECTION,
                collection=visible_collections.get(id=int(m.group(1))),
            )
        else:
            if lookup == "internal@collections":
                if workflow_root is None:
                    raise LookupError(
                        "internal@collections is only valid in the context of "
                        "a workflow"
                    )
                name = f"workflow-{workflow_root.id}"
                category = str(CollectionCategory.WORKFLOW_INTERNAL)
            elif "@" in lookup:
                name, category = lookup.rsplit("@", 1)
            elif default_category is not None:
                name = lookup
                category = default_category
            else:
                raise LookupError(
                    f"{lookup!r} does not specify a category and the context "
                    f"does not supply a default"
                )
            return LookupResult(
                result_type=CollectionItem.Types.COLLECTION,
                collection=workspace.get_collection(
                    user=user, category=category, name=name
                ),
            )
    except (Artifact.DoesNotExist, Collection.DoesNotExist):
        raise KeyError(f"{lookup!r} does not exist or is hidden")


@overload
def lookup_single(
    lookup: LookupSingle,
    workspace: Workspace,
    *,
    user: User | None,
    default_category: CollectionCategory | None = None,
    workflow_root: WorkRequest | None = None,
    expect_type: Literal[CollectionItem.Types.ARTIFACT],
) -> LookupResultArtifact: ...


@overload
def lookup_single(
    lookup: LookupSingle,
    workspace: Workspace,
    *,
    user: User | None,
    default_category: CollectionCategory | None = None,
    workflow_root: WorkRequest | None = None,
    expect_type: Literal[CollectionItem.Types.COLLECTION],
) -> LookupResultCollection: ...


@overload
def lookup_single(
    lookup: LookupSingle,
    workspace: Workspace,
    *,
    user: User | None,
    default_category: CollectionCategory | None = None,
    workflow_root: WorkRequest | None = None,
    expect_type: CollectionItem.Types | None = None,
) -> LookupResult: ...


def lookup_single(
    lookup: LookupSingle,
    workspace: Workspace,
    *,
    user: User | None,
    default_category: CollectionCategory | None = None,
    workflow_root: WorkRequest | None = None,
    expect_type: CollectionItem.Types | None = None,
) -> LookupResult:
    """
    Look up a collection item using a single lookup syntax.

    See :ref:`lookup-single`.

    :raises KeyError: if the lookup does not resolve to an item.
    :raises LookupError: if the lookup is invalid in some way.
    """
    if isinstance(lookup, int):
        if expect_type == CollectionItem.Types.ARTIFACT:
            lookup = f"{lookup}@artifacts"
        elif expect_type == CollectionItem.Types.COLLECTION:
            lookup = f"{lookup}@collections"
        else:
            raise LookupError(
                "Integer lookups only work in contexts that expect an "
                "artifact or a collection"
            )

    if not lookup:
        raise LookupError("Empty lookup")

    segments = parse_lookup_string_segments(lookup)

    # Look up the first segment as a collection by name and category.
    result = _lookup_single_first(
        segments[0],
        workspace,
        user=user,
        default_category=default_category,
        workflow_root=workflow_root,
    )
    container: Collection | None = None

    # Resolve each subsequent segment by calling `lookup`.
    for i, segment in enumerate(segments[1:], start=1):
        container_name = "/".join(segments[:i])
        if result.result_type != CollectionItem.Types.COLLECTION:
            raise LookupError(f"{container_name!r} is not of type 'collection'")
        container = result.collection
        assert container is not None
        if ":" not in segment:
            segment = f"name:{segment}"
        item = container.manager.lookup(segment)
        if item is None:
            raise KeyError(f"{container_name!r} has no item {segment!r}")
        result = LookupResult(
            result_type=CollectionItem.Types(item.child_type),
            collection_item=item,
            artifact=item.artifact,
            collection=item.collection,
        )

    if expect_type is not None and result.result_type != expect_type:
        raise LookupError(
            f"{lookup!r} is not of type {expect_type.name.lower()!r}"
        )

    return result


def _lookup_dict(
    lookup: LookupDict,
    workspace: Workspace,
    *,
    user: User | None,
    default_category: CollectionCategory | None = None,
    workflow_root: WorkRequest | None = None,
    expect_type: CollectionItem.Types | None = None,
) -> list[LookupResult]:
    """
    Look up collection items using the dictionary lookup syntax.

    See :ref:`lookup-dict`.
    """
    # Find the containing collection.
    try:
        collection = lookup_single(
            lookup.collection,
            workspace,
            user=user,
            default_category=default_category,
            workflow_root=workflow_root,
            expect_type=CollectionItem.Types.COLLECTION,
        ).collection
    except KeyError:
        return []
    objects = CollectionItem.active_objects.filter(parent_collection=collection)

    # Prepare query conditions.  We don't need to check the workspace here;
    # it's good enough if the item is in a collection we can see.
    match lookup.child_type:
        case LookupChildType.BARE:
            child_type = CollectionItem.Types.BARE
        case LookupChildType.ARTIFACT:
            child_type = CollectionItem.Types.ARTIFACT
        case LookupChildType.COLLECTION:
            child_type = CollectionItem.Types.COLLECTION
        case _ as unreachable:
            assert_never(unreachable)
    if expect_type is not None and expect_type != child_type:
        raise LookupError(
            f"Only lookups for type {expect_type.name.lower()!r} are allowed "
            f"here"
        )
    objects = objects.filter(child_type=child_type)
    if lookup.category is not None:
        objects = objects.filter(category=lookup.category)
    if lookup.name_matcher is not None:
        objects = objects.filter(
            **{f"name__{lookup.name_matcher.kind}": lookup.name_matcher.value}
        )
    for key, matcher in lookup.data_matchers:
        annotation = f"data_text_{key}"
        objects = objects.annotate(
            **{annotation: Cast(KeyTextTransform(key, "data"), TextField())}
        ).filter(**{f"{annotation}__{matcher.kind}": matcher.value})

    # Execute the query.
    return [
        LookupResult(
            result_type=CollectionItem.Types(item.child_type),
            collection_item=item,
            artifact=item.artifact,
            collection=item.collection,
        )
        for item in objects.all()
    ]


@overload
def lookup_multiple(
    lookup: LookupMultiple,
    workspace: Workspace,
    *,
    user: User | None,
    default_category: CollectionCategory | None = None,
    workflow_root: WorkRequest | None = None,
    expect_type: Literal[CollectionItem.Types.ARTIFACT],
) -> Sequence[LookupResultArtifact]: ...


@overload
def lookup_multiple(
    lookup: LookupMultiple,
    workspace: Workspace,
    *,
    user: User | None,
    default_category: CollectionCategory | None = None,
    workflow_root: WorkRequest | None = None,
    expect_type: Literal[CollectionItem.Types.COLLECTION],
) -> Sequence[LookupResultCollection]: ...


@overload
def lookup_multiple(
    lookup: LookupMultiple,
    workspace: Workspace,
    *,
    user: User | None,
    default_category: CollectionCategory | None = None,
    workflow_root: WorkRequest | None = None,
    expect_type: CollectionItem.Types | None = None,
) -> Sequence[LookupResult]: ...


def lookup_multiple(
    lookup: LookupMultiple,
    workspace: Workspace,
    *,
    user: User | None,
    default_category: CollectionCategory | None = None,
    workflow_root: WorkRequest | None = None,
    expect_type: CollectionItem.Types | None = None,
) -> Sequence[LookupResult]:
    """
    Look up collection items using a multiple lookup syntax.

    See :ref:`lookup-multiple`.

    :raises LookupError: if the lookup is invalid in some way.
    """
    results: list[LookupResult] = []
    for alternative in lookup.__root__:
        match alternative:
            case int() | str():
                try:
                    results.append(
                        lookup_single(
                            alternative,
                            workspace,
                            user=user,
                            default_category=default_category,
                            workflow_root=workflow_root,
                            expect_type=expect_type,
                        )
                    )
                except KeyError:
                    pass
            case LookupDict():
                results.extend(
                    _lookup_dict(
                        alternative,
                        workspace,
                        user=user,
                        default_category=default_category,
                        workflow_root=workflow_root,
                        expect_type=expect_type,
                    )
                )
            case _ as unreachable:
                assert_never(unreachable)
    return tuple(results)
