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

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

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

from django.db.models import Q
from django.db.models.fields.json import KT

from debusine.artifacts.models import BareDataCategory, CollectionCategory
from debusine.client.models import LookupChildType, LookupResultType
from debusine.db.context import context
from debusine.db.models import (
    Artifact,
    Collection,
    CollectionItem,
    WorkRequest,
    Workspace,
)
from debusine.db.models.lookups import (
    LookupResult,
    LookupResultArtifact,
    LookupResultBare,
    LookupResultCollection,
)
from debusine.db.permissioncontext import PermissionContext
from debusine.tasks.models import (
    LookupDict,
    LookupMultiple,
    LookupSingle,
    build_lookup_string,
    parse_lookup_string,
)


def _lookup_result_expected(
    result: LookupResult, expect_type: LookupChildType
) -> bool:
    """Check whether a lookup result is of the expected type."""
    if expect_type != LookupChildType.ANY:
        expected_result_types = {
            LookupChildType.BARE: [LookupResultType.BARE],
            LookupChildType.ARTIFACT: [LookupResultType.ARTIFACT],
            LookupChildType.ARTIFACT_OR_PROMISE: [
                LookupResultType.ARTIFACT,
                LookupResultType.BARE,
            ],
            LookupChildType.COLLECTION: [LookupResultType.COLLECTION],
        }[expect_type]
        if result.result_type not in expected_result_types:
            return False
    return True


def _normalize_lookup_segments(
    lookup: LookupSingle,
    default_category: CollectionCategory | None = None,
    expect_type: LookupChildType = LookupChildType.ANY,
) -> tuple[list[str], dict[str, Any] | None]:
    """Normalize a single lookup into a list of segments."""
    if isinstance(lookup, int):
        if expect_type in (
            LookupChildType.ARTIFACT,
            LookupChildType.ARTIFACT_OR_PROMISE,
        ):
            lookup = f"{lookup}@artifacts"
        elif expect_type == LookupChildType.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")

    try:
        segments, options = parse_lookup_string(lookup)
    except ValueError:
        raise LookupError(
            f"Invalid syntax in lookup options part of {lookup!r}"
        )

    first_segment = segments[0]
    if "@" not in first_segment:
        if default_category is not None:
            first_segment = f"{first_segment}@{default_category}"
        else:
            raise LookupError(
                f"{first_segment!r} does not specify a category and the "
                f"context does not supply a default"
            )

    return ([first_segment, *segments[1:]], options)


def _lookup_single_first(
    lookup: str,
    workspace: Workspace,
    *,
    original_lookup: LookupSingle,
    pc: PermissionContext,
    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):
            # TODO: filter using the current scope?
            # TODO: support user is None to mean the user in the current
            #       context?
            workspaces = Workspace.objects.can_display(pc)
            workspace_filter = Q(workspace__in=workspaces)
            visible_artifacts = Artifact.objects.filter(workspace_filter)
            return LookupResultArtifact(
                result_type=LookupResultType.ARTIFACT,
                lookup=original_lookup,
                artifact=visible_artifacts.get(id=int(m.group(1))),
            )
        elif m := re.match(r"^([0-9]+)@collections$", lookup):
            # TODO: filter using the current scope?
            # TODO: support user is None to mean the user in the current
            #       context?
            workspaces = Workspace.objects.can_display(pc)
            workspace_filter = Q(workspace__in=workspaces)
            visible_collections = Collection.objects.filter(workspace_filter)
            return LookupResultCollection(
                result_type=LookupResultType.COLLECTION,
                lookup=original_lookup,
                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)
            else:
                # Ensured by lookup_single.
                assert "@" in lookup
                name, category = lookup.rsplit("@", 1)
            return LookupResultCollection(
                result_type=LookupResultType.COLLECTION,
                lookup=original_lookup,
                collection=workspace.get_collection(
                    pc=pc, category=category, name=name
                ),
            )
    except (Artifact.DoesNotExist, Collection.DoesNotExist):
        raise KeyError(f"{lookup!r} does not exist or is hidden")


def _process_lookup_options(
    options: dict[str, Any], workspace: Workspace, pc: PermissionContext
) -> None:
    for name, value in options.items():
        match name:
            case "scope":
                if "workspace" not in options:
                    raise LookupError(
                        "The 'scope' lookup option requires the 'workspace' one"
                    )
            case "workspace":
                scope_name = options.get("scope", workspace.scope.name)
                try:
                    requested_workspace = Workspace.objects.can_display(pc).get(
                        name=value, scope__name=scope_name
                    )
                    options["workspace"] = requested_workspace
                except Workspace.DoesNotExist:
                    raise LookupError(
                        f"Invalid or hidden workspace in lookup options: "
                        f"scope={scope_name!r} workspace={value!r}"
                    )
            case _:
                raise LookupError(f"Unknown lookup option: {name!r}")


@overload
def lookup_single(
    lookup: LookupSingle,
    workspace: Workspace,
    *,
    pc: PermissionContext,
    default_category: CollectionCategory | None = None,
    workflow_root: WorkRequest | None = None,
    expect_type: Literal[LookupChildType.BARE],
    precache_result: LookupResultBare | None = None,
) -> LookupResultBare: ...


@overload
def lookup_single(
    lookup: LookupSingle,
    workspace: Workspace,
    *,
    pc: PermissionContext,
    default_category: CollectionCategory | None = None,
    workflow_root: WorkRequest | None = None,
    expect_type: Literal[LookupChildType.ARTIFACT],
    precache_result: LookupResultArtifact | None = None,
) -> LookupResultArtifact: ...


@overload
def lookup_single(
    lookup: LookupSingle,
    workspace: Workspace,
    *,
    pc: PermissionContext,
    default_category: CollectionCategory | None = None,
    workflow_root: WorkRequest | None = None,
    expect_type: Literal[LookupChildType.COLLECTION],
    precache_result: LookupResultCollection | None = None,
) -> LookupResultCollection: ...


@overload
def lookup_single(
    lookup: LookupSingle,
    workspace: Workspace,
    *,
    pc: PermissionContext,
    default_category: CollectionCategory | None = None,
    workflow_root: WorkRequest | None = None,
    expect_type: LookupChildType = LookupChildType.ANY,
    precache_result: LookupResult | None = None,
) -> LookupResult: ...


def lookup_single(
    lookup: LookupSingle,
    workspace: Workspace,
    *,
    pc: PermissionContext,
    default_category: CollectionCategory | None = None,
    workflow_root: WorkRequest | None = None,
    expect_type: LookupChildType = LookupChildType.ANY,
    precache_result: LookupResult | None = None,
) -> LookupResult:
    """
    Look up a collection item using a single lookup syntax.

    See :ref:`lookup-single`.

    :param precache_result: Rather than doing a normal lookup, insert this
      result into the lookup cache and return it.  This is useful to improve
      workflow performance in some situations, such as workflows that do
      bulk database lookups and then reconstruct lookups to pass down to
      their children.
    :raises KeyError: if the lookup does not resolve to an item.
    :raises LookupError: if the lookup is invalid in some way, or does not
      resolve to the expected collection item type.
    """
    # TODO: make PermissionContext hashable and use that in the cache key
    cache = context.lookup_cache.single
    cache_key = (
        lookup,
        workspace,
        pc.user,
        default_category,
        workflow_root,
        expect_type,
    )
    if precache_result is not None:
        cache[cache_key] = precache_result
        return precache_result
    if (cached := cache.get(cache_key)) is not None:
        return cached

    segments, options = _normalize_lookup_segments(
        lookup, default_category=default_category, expect_type=expect_type
    )

    if options:
        _process_lookup_options(options, workspace=workspace, pc=pc)
        workspace = options.get("workspace", workspace)

    # Look up the first segment as a collection by name and category.
    result = _lookup_single_first(
        segments[0],
        workspace,
        original_lookup=lookup,
        pc=pc,
        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 = build_lookup_string(*segments[:i])
        if result.result_type != LookupResultType.COLLECTION:
            raise LookupError(
                f"{container_name!r} is of type"
                f" {result.result_type.name.lower()!r}"
                " instead of expected '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.create(
            result_type=LookupResultType(item.child_type),
            lookup=lookup,
            parent_collection_lookup=container_name,
            collection_item=item,
            artifact=item.artifact,
            collection=item.collection,
        )

    if not _lookup_result_expected(result, expect_type):
        raise LookupError(
            f"{lookup!r} is of type {result.result_type.name.lower()!r}"
            f" instead of expected {expect_type.name.lower()!r}"
        )

    cache[cache_key] = result
    return result


def _lookup_dict(  # noqa: C901
    lookup: LookupDict,
    workspace: Workspace,
    *,
    pc: PermissionContext,
    default_category: CollectionCategory | None = None,
    workflow_root: WorkRequest | None = None,
    expect_type: LookupChildType = LookupChildType.ANY,
) -> Sequence[LookupResult]:
    """
    Look up collection items using the dictionary lookup syntax.

    Results are returned in a predictable order (currently ordered by
    collection item ID, although this function only guarantees that the
    order will be stable between successive calls).

    See :ref:`lookup-dict`.
    """
    # TODO: use PermissionContext in the cache key instead of user
    cache = context.lookup_cache.dictionary
    cache_key = (
        lookup,
        workspace,
        pc.user,
        default_category,
        workflow_root,
        expect_type,
    )
    if (cached := cache.get(cache_key)) is not None:
        return cached

    parent_collection_segments, options = _normalize_lookup_segments(
        lookup.collection,
        default_category=default_category,
        expect_type=LookupChildType.COLLECTION,
    )

    parent_collection_lookup = build_lookup_string(
        *parent_collection_segments, options=options
    )

    # Find the containing collection.
    collection = lookup_single(
        lookup.collection,
        workspace,
        pc=pc,
        default_category=default_category,
        workflow_root=workflow_root,
        expect_type=LookupChildType.COLLECTION,
    ).collection
    objects = collection.child_items.active()

    # 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:
            objects = objects.filter(child_type=CollectionItem.Types.BARE)
        case LookupChildType.ARTIFACT:
            objects = objects.filter(child_type=CollectionItem.Types.ARTIFACT)
        case LookupChildType.ARTIFACT_OR_PROMISE:
            objects = objects.filter(
                Q(child_type=CollectionItem.Types.ARTIFACT)
                | Q(
                    child_type=CollectionItem.Types.BARE,
                    category=BareDataCategory.PROMISE,
                )
            )
        case LookupChildType.COLLECTION:
            objects = objects.filter(child_type=CollectionItem.Types.COLLECTION)
        case LookupChildType.ANY:
            pass
        case _ as unreachable:
            assert_never(unreachable)
    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: KT(f"data__{key}")}).filter(
            **{f"{annotation}__{matcher.kind}": matcher.value}
        )
    for key, value in lookup.lookup_filters:
        objects = objects.filter(
            collection.manager.lookup_filter(
                key,
                value,
                workspace=workspace,
                pc=pc,
                workflow_root=workflow_root,
            )
        )
    objects = objects.select_related("artifact", "collection")

    # Execute the query.
    results = [
        LookupResult.create(
            result_type=LookupResultType(item.child_type),
            lookup=lookup,
            parent_collection_lookup=parent_collection_lookup,
            collection_item=item,
            artifact=item.artifact,
            collection=item.collection,
        )
        for item in objects.order_by("id")
    ]
    if lookup.error_on_empty and not results:
        raise LookupError(f"No results found from {lookup.export()}")
    for result in results:
        if not _lookup_result_expected(result, expect_type):
            raise LookupError(
                f"Result from {lookup.export()} is of type"
                f" {result.result_type.name.lower()!r} instead of expected"
                f" {expect_type.name.lower()!r}"
            )

    frozen_results = tuple(results)
    cache[cache_key] = frozen_results
    return frozen_results


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


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


@overload
def lookup_multiple(
    lookup: LookupMultiple,
    workspace: Workspace,
    *,
    pc: PermissionContext,
    default_category: CollectionCategory | None = None,
    workflow_root: WorkRequest | None = None,
    expect_type: LookupChildType = LookupChildType.ANY,
) -> Sequence[LookupResult]: ...


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

    See :ref:`lookup-multiple`.

    :raises KeyError: if any of the lookups does not resolve to an item.
    :raises LookupError: if any of the lookups is invalid in some way, or
      does not resolve to the expected collection item type.
    """
    results: list[LookupResult] = []
    for alternative in lookup:
        match alternative:
            case int() | str():
                results.append(
                    lookup_single(
                        alternative,
                        workspace,
                        pc=pc,
                        default_category=default_category,
                        workflow_root=workflow_root,
                        expect_type=expect_type,
                    )
                )
            case LookupDict():
                results.extend(
                    _lookup_dict(
                        alternative,
                        workspace,
                        pc=pc,
                        default_category=default_category,
                        workflow_root=workflow_root,
                        expect_type=expect_type,
                    )
                )
            case _ as unreachable:
                assert_never(unreachable)
    return tuple(results)


def artifact_ids(results: Sequence[LookupResult]) -> list[int]:
    """Return the artifact ids."""
    return [
        result.artifact.id for result in results if result.artifact is not None
    ]
