Corp/backend/tasks/group_ranking.py

"""Lexicographically sortable rank keys for boards inside a group."""

from __future__ import annotations

from collections.abc import Sequence

from django.db.models import QuerySet

from .models import TaskBoard, TaskBoardGroup, TaskBoardGroupMembership

RANK_ALPHABET = "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz"
RANK_BASE = len(RANK_ALPHABET)
RANK_WIDTH = 32
RANK_MAX_VALUE = RANK_BASE**RANK_WIDTH - 1
_RANK_DIGITS = {character: index for index, character in enumerate(RANK_ALPHABET)}


class RankSpaceExhausted(ValueError):
    """There is no representable key between two adjacent rank keys."""


def encode_rank(value: int) -> str:
    if not 0 <= value <= RANK_MAX_VALUE:
        raise ValueError("Rank value is outside the supported range.")
    characters = [RANK_ALPHABET[0]] * RANK_WIDTH
    for index in range(RANK_WIDTH - 1, -1, -1):
        value, remainder = divmod(value, RANK_BASE)
        characters[index] = RANK_ALPHABET[remainder]
    return "".join(characters)


def decode_rank(rank_key: str) -> int:
    if len(rank_key) != RANK_WIDTH:
        raise ValueError("Rank key has an invalid length.")
    value = 0
    for character in rank_key:
        try:
            digit = _RANK_DIGITS[character]
        except KeyError as error:
            raise ValueError("Rank key contains an invalid character.") from error
        value = value * RANK_BASE + digit
    return value


def sort_group_memberships_by_rank(
    memberships: Sequence[TaskBoardGroupMembership] | QuerySet,
) -> list[TaskBoardGroupMembership]:
    """Sort ranks independently from the database's text collation."""
    return sorted(
        memberships,
        key=lambda membership: (decode_rank(membership.rank_key), membership.id),
    )


def rank_between(lower_key: str | None, upper_key: str | None) -> str:
    lower_value = decode_rank(lower_key) if lower_key is not None else 0
    upper_value = decode_rank(upper_key) if upper_key is not None else RANK_MAX_VALUE
    if lower_value >= upper_value:
        raise ValueError("The lower rank must precede the upper rank.")
    if upper_value - lower_value <= 1:
        raise RankSpaceExhausted("No rank key remains between the neighbours.")
    return encode_rank((lower_value + upper_value) // 2)


def evenly_spaced_rank_keys(count: int) -> list[str]:
    if count <= 0:
        return []
    step = RANK_MAX_VALUE // (count + 1)
    if step <= 0:
        raise RankSpaceExhausted("The group contains too many boards to rank.")
    return [encode_rank(step * index) for index in range(1, count + 1)]


def allocate_rank_keys_after(last_key: str | None, count: int) -> list[str]:
    if count <= 0:
        return []
    lower_value = decode_rank(last_key) if last_key is not None else 0
    step = (RANK_MAX_VALUE - lower_value) // (count + 1)
    if step <= 0:
        raise RankSpaceExhausted("No rank keys remain after the last board.")
    return [encode_rank(lower_value + step * index) for index in range(1, count + 1)]


def rebalance_group_memberships(
    group: TaskBoardGroup,
    memberships: Sequence[TaskBoardGroupMembership] | QuerySet | None = None,
) -> list[TaskBoardGroupMembership]:
    ordered_memberships = sort_group_memberships_by_rank(
        memberships
        if memberships is not None
        else TaskBoardGroupMembership.objects.filter(group=group)
    )
    final_rank_keys = evenly_spaced_rank_keys(len(ordered_memberships))
    if ordered_memberships:
        # Rank keys are unique inside a group. Updating directly from the old
        # order to the new one can make PostgreSQL observe a transient
        # duplicate (A receives B's key before B releases it). Move every row
        # into a disjoint temporary namespace first, then assign final keys.
        for membership in ordered_memberships:
            temporary_suffix = format(membership.pk, "031x")
            if len(temporary_suffix) > RANK_WIDTH - 1:
                raise ValueError("Membership id is too large for a temporary rank key.")
            membership.rank_key = f"~{temporary_suffix}"
        TaskBoardGroupMembership.objects.bulk_update(
            ordered_memberships,
            ["rank_key"],
        )

        for membership, rank_key in zip(
            ordered_memberships,
            final_rank_keys,
            strict=True,
        ):
            membership.rank_key = rank_key
        TaskBoardGroupMembership.objects.bulk_update(
            ordered_memberships,
            ["rank_key"],
        )
    return ordered_memberships


def allocate_group_tail_rank_keys(
    group: TaskBoardGroup,
    count: int,
) -> list[str]:
    if count <= 0:
        return []
    rank_keys = list(
        TaskBoardGroupMembership.objects.filter(group=group).values_list(
            "rank_key", flat=True
        )
    )
    last_key = max(rank_keys, key=decode_rank, default=None)
    try:
        return allocate_rank_keys_after(last_key, count)
    except RankSpaceExhausted:
        memberships = rebalance_group_memberships(group)
        last_key = memberships[-1].rank_key if memberships else None
        return allocate_rank_keys_after(last_key, count)


def visible_group_position_map(
    user,
    group_ids: Sequence[int] | None = None,
) -> dict[tuple[int, int], int]:
    """Return consecutive positions for boards visible to one user."""
    from django.db.models import Q

    from .access import task_board_access_q

    accessible_boards = TaskBoard.objects.filter(task_board_access_q(user))
    memberships = (
        TaskBoardGroupMembership.objects.filter(
            board__in=accessible_boards,
        )
        .filter(
            Q(group__user=user) | Q(group__is_shared=True),
        )
        .exclude(group__hidden_by=user)
    )
    if group_ids is not None:
        memberships = memberships.filter(group_id__in=group_ids)

    membership_rows = list(
        memberships.values_list("group_id", "board_id", "rank_key", "id")
    )
    membership_rows.sort(
        key=lambda row: (row[0], decode_rank(row[2]), row[3]),
    )

    result: dict[tuple[int, int], int] = {}
    current_group_id = None
    position = 0
    for group_id, board_id, _rank_key, _membership_id in membership_rows:
        if group_id != current_group_id:
            current_group_id = group_id
            position = 0
        position += 1
        result[(group_id, board_id)] = position
    return result