/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements. See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License. You may obtain a copy of the License at
 *
 *    http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package org.apache.kafka.coordinator.group.streams;

import org.apache.kafka.common.errors.FencedMemberEpochException;

import java.util.HashMap;
import java.util.HashSet;
import java.util.Map;
import java.util.Objects;
import java.util.Optional;
import java.util.Set;
import java.util.function.BiFunction;
import java.util.function.BiPredicate;

/**
 * The CurrentAssignmentBuilder class encapsulates the reconciliation engine of the streams group protocol. Given the current state of a
 * member and a desired or target assignment state, the state machine takes the necessary steps to converge them.
 */
public class CurrentAssignmentBuilder {

    /**
     * The streams group member which is reconciled.
     */
    private final StreamsGroupMember member;

    /**
     * The target assignment epoch.
     */
    private int targetAssignmentEpoch;

    /**
     * The target assignment.
     */
    private TasksTuple targetAssignment;

    /**
     * A function which returns the current process ID of an active task or null if the active task
     * is not assigned. The current process ID is the process ID of the current owner.
     */
    private BiFunction<String, Integer, String> currentActiveTaskProcessId;

    /**
     * A function which returns the current process IDs of a standby task or null if the standby
     * task is not assigned. The current process IDs are the process IDs of all current owners.
     */
    private BiFunction<String, Integer, Set<String>> currentStandbyTaskProcessIds;

    /**
     * A function which returns the current process IDs of a warmup task or null if the warmup task
     * is not assigned. The current process IDs are the process IDs of all current owners.
     */
    private BiFunction<String, Integer, Set<String>> currentWarmupTaskProcessIds;

    /**
     * The tasks owned by the member. This may be provided by the member in the StreamsGroupHeartbeat request.
     */
    private Optional<TasksTuple> ownedTasks = Optional.empty();

    /**
     * Constructs the CurrentAssignmentBuilder based on the current state of the provided streams group member.
     *
     * @param member The streams group member that must be reconciled.
     */
    public CurrentAssignmentBuilder(StreamsGroupMember member) {
        this.member = Objects.requireNonNull(member);
    }

    /**
     * Sets the target assignment epoch and the target assignment that the streams group member must be reconciled to.
     *
     * @param targetAssignmentEpoch The target assignment epoch.
     * @param targetAssignment      The target assignment.
     * @return This object.
     */
    public CurrentAssignmentBuilder withTargetAssignment(int targetAssignmentEpoch,
                                                         TasksTuple targetAssignment) {
        this.targetAssignmentEpoch = targetAssignmentEpoch;
        this.targetAssignment = Objects.requireNonNull(targetAssignment);
        return this;
    }

    /**
     * Sets a BiFunction which allows to retrieve the current process ID of an active task. This is
     * used by the state machine to determine if an active task is free or still used by another
     * member, and if there is still a task on a specific process that is not yet revoked.
     *
     * @param currentActiveTaskProcessId A BiFunction which gets the process ID of a subtopology ID /
     *                                   partition ID pair.
     * @return This object.
     */
    public CurrentAssignmentBuilder withCurrentActiveTaskProcessId(BiFunction<String, Integer, String> currentActiveTaskProcessId) {
        this.currentActiveTaskProcessId = Objects.requireNonNull(currentActiveTaskProcessId);
        return this;
    }

    /**
     * Sets a BiFunction which allows to retrieve the current process IDs of a standby task. This is
     * used by the state machine to determine if there is still a task on a specific process that is
     * not yet revoked.
     *
     * @param currentStandbyTaskProcessIds A BiFunction which gets the process IDs of a subtopology
     *                                     ID / partition ID pair.
     * @return This object.
     */
    public CurrentAssignmentBuilder withCurrentStandbyTaskProcessIds(
        BiFunction<String, Integer, Set<String>> currentStandbyTaskProcessIds
    ) {
        this.currentStandbyTaskProcessIds = Objects.requireNonNull(currentStandbyTaskProcessIds);
        return this;
    }

    /**
     * Sets a BiFunction which allows to retrieve the current process IDs of a warmup task. This is
     * used by the state machine to determine if there is still a task on a specific process that is
     * not yet revoked.
     *
     * @param currentWarmupTaskProcessIds A BiFunction which gets the process IDs of a subtopology ID
     *                                    / partition ID pair.
     * @return This object.
     */
    public CurrentAssignmentBuilder withCurrentWarmupTaskProcessIds(BiFunction<String, Integer, Set<String>> currentWarmupTaskProcessIds) {
        this.currentWarmupTaskProcessIds = Objects.requireNonNull(currentWarmupTaskProcessIds);
        return this;
    }

    /**
     * Sets the tasks currently owned by the member. This comes directly from the last StreamsGroupHeartbeat request. This is used to
     * determine if the member has revoked the necessary tasks. Passing null into this function means that the member did not provide
     * its owned tasks in this heartbeat.
     *
     * @param ownedAssignment A collection of active, standby and warm-up tasks
     * @return This object.
     */
    public CurrentAssignmentBuilder withOwnedAssignment(TasksTuple ownedAssignment) {
        this.ownedTasks = Optional.ofNullable(ownedAssignment);
        return this;
    }

    /**
     * Builds the next state for the member or keep the current one if it is not possible to move forward with the current state.
     *
     * @return A new StreamsGroupMember or the current one.
     */
    public StreamsGroupMember build() {
        switch (member.state()) {
            case STABLE:
                // When the member is in the STABLE state, we verify if a newer
                // epoch (or target assignment) is available. If it is, we can
                // reconcile the member towards it. Otherwise, we return.
                if (member.memberEpoch() != targetAssignmentEpoch) {
                    return computeNextAssignment(
                        member.memberEpoch(),
                        member.assignedTasks()
                    );
                } else {
                    return member;
                }

            case UNREVOKED_TASKS:
                // When the member is in the UNREVOKED_TASKS state, we wait
                // until the member has revoked the necessary tasks. They are
                // considered revoked when they are not anymore reported in the
                // owned tasks set in the StreamsGroupHeartbeat API.

                // If the member provides its owned tasks, we verify if it still
                // owns any of the revoked tasks. If it did not provide it's
                // owned tasks, or we still own some of the revoked tasks, we
                // cannot progress.
                if (
                    ownedTasks.isEmpty() || ownedTasks.get().containsAny(member.tasksPendingRevocation())
                ) {
                    return member;
                }

                // When the member has revoked all the pending tasks, it can
                // transition to the next epoch (current + 1) and we can reconcile
                // its state towards the latest target assignment.
                return computeNextAssignment(
                    member.memberEpoch() + 1,
                    member.assignedTasks()
                );

            case UNRELEASED_TASKS:
                // When the member is in the UNRELEASED_TASKS, we reconcile the
                // member towards the latest target assignment. This will assign any
                // of the unreleased tasks when they become available.
                return computeNextAssignment(
                    member.memberEpoch(),
                    member.assignedTasks()
                );

            case UNKNOWN:
                // We could only end up in this state if a new state is added in the
                // future and the group coordinator is downgraded. In this case, the
                // best option is to fence the member to force it to rejoin the group
                // without any tasks and to reconcile it again from scratch.
                if ((ownedTasks.isEmpty() || !ownedTasks.get().isEmpty())) {
                    throw new FencedMemberEpochException(
                        "The streams group member is in a unknown state. "
                            + "The member must abandon all its tasks and rejoin.");
                }

                return computeNextAssignment(
                    targetAssignmentEpoch,
                    member.assignedTasks()
                );
        }

        return member;
    }

    /**
     * Takes the current currentAssignment and the targetAssignment, and generates three
     * collections:
     *
     * - the resultAssignedTasks: the tasks that are assigned in both the current and target
     * assignments.
     * - the resultTasksPendingRevocation: the tasks that are assigned in the current
     * assignment but not in the target assignment.
     * - the resultTasksPendingAssignment: the tasks that are assigned in the target assignment but
     * not in the current assignment, and can be assigned currently (i.e., they are not owned by
     * another member, as defined by the `isUnreleasedTask` predicate).
     */
    private boolean computeAssignmentDifference(Map<String, Set<Integer>> currentAssignment,
                                                Map<String, Set<Integer>> targetAssignment,
                                                Map<String, Set<Integer>> resultAssignedTasks,
                                                Map<String, Set<Integer>> resultTasksPendingRevocation,
                                                Map<String, Set<Integer>> resultTasksPendingAssignment,
                                                BiPredicate<String, Integer> isUnreleasedTask) {
        boolean hasUnreleasedTasks = false;

        Set<String> allSubtopologyIds = new HashSet<>(targetAssignment.keySet());
        allSubtopologyIds.addAll(currentAssignment.keySet());

        for (String subtopologyId : allSubtopologyIds) {
            hasUnreleasedTasks |= computeAssignmentDifferenceForOneSubtopology(
                subtopologyId,
                currentAssignment.getOrDefault(subtopologyId, Set.of()),
                targetAssignment.getOrDefault(subtopologyId, Set.of()),
                resultAssignedTasks,
                resultTasksPendingRevocation,
                resultTasksPendingAssignment,
                isUnreleasedTask
            );
        }
        return hasUnreleasedTasks;
    }

    private static boolean computeAssignmentDifferenceForOneSubtopology(final String subtopologyId,
                                                                        final Set<Integer> currentTasksForThisSubtopology,
                                                                        final Set<Integer> targetTasksForThisSubtopology,
                                                                        final Map<String, Set<Integer>> resultAssignedTasks,
                                                                        final Map<String, Set<Integer>> resultTasksPendingRevocation,
                                                                        final Map<String, Set<Integer>> resultTasksPendingAssignment,
                                                                        final BiPredicate<String, Integer> isUnreleasedTask) {
        // Result Assigned Tasks = Current Tasks ∩ Target Tasks
        // i.e. we remove all tasks from the current assignment that are not in the target
        //         assignment
        Set<Integer> resultAssignedTasksForThisSubtopology = new HashSet<>(currentTasksForThisSubtopology);
        resultAssignedTasksForThisSubtopology.retainAll(targetTasksForThisSubtopology);

        // Result Tasks Pending Revocation = Current Tasks - Result Assigned Tasks
        // i.e. we will ask the member to revoke all tasks in its current assignment that
        //      are not in the target assignment
        Set<Integer> resultTasksPendingRevocationForThisSubtopology = new HashSet<>(currentTasksForThisSubtopology);
        resultTasksPendingRevocationForThisSubtopology.removeAll(resultAssignedTasksForThisSubtopology);

        // Result Tasks Pending Assignment = Target Tasks - Result Assigned Tasks - Unreleased Tasks
        // i.e. we will ask the member to assign all tasks in its target assignment,
        //      except those that are already assigned, and those that are unreleased
        Set<Integer> resultTasksPendingAssignmentForThisSubtopology = new HashSet<>(targetTasksForThisSubtopology);
        resultTasksPendingAssignmentForThisSubtopology.removeAll(resultAssignedTasksForThisSubtopology);
        boolean hasUnreleasedTasks = resultTasksPendingAssignmentForThisSubtopology.removeIf(taskId ->
            isUnreleasedTask.test(subtopologyId, taskId)
        );

        if (!resultAssignedTasksForThisSubtopology.isEmpty()) {
            resultAssignedTasks.put(subtopologyId, resultAssignedTasksForThisSubtopology);
        }

        if (!resultTasksPendingRevocationForThisSubtopology.isEmpty()) {
            resultTasksPendingRevocation.put(subtopologyId, resultTasksPendingRevocationForThisSubtopology);
        }

        if (!resultTasksPendingAssignmentForThisSubtopology.isEmpty()) {
            resultTasksPendingAssignment.put(subtopologyId, resultTasksPendingAssignmentForThisSubtopology);
        }

        return hasUnreleasedTasks;
    }

    /**
     * Computes the next assignment.
     *
     * @param memberEpoch         The epoch of the member to use. This may be different from
     *                            the epoch in {@link CurrentAssignmentBuilder#member}.
     * @param memberAssignedTasks The assigned tasks of the member to use.
     * @return A new StreamsGroupMember.
     */
    private StreamsGroupMember computeNextAssignment(int memberEpoch,
                                                     TasksTuple memberAssignedTasks) {
        Map<String, Set<Integer>> newActiveAssignedTasks = new HashMap<>();
        Map<String, Set<Integer>> newActiveTasksPendingRevocation = new HashMap<>();
        Map<String, Set<Integer>> newActiveTasksPendingAssignment = new HashMap<>();
        Map<String, Set<Integer>> newStandbyAssignedTasks = new HashMap<>();
        Map<String, Set<Integer>> newStandbyTasksPendingRevocation = new HashMap<>();
        Map<String, Set<Integer>> newStandbyTasksPendingAssignment = new HashMap<>();
        Map<String, Set<Integer>> newWarmupAssignedTasks = new HashMap<>();
        Map<String, Set<Integer>> newWarmupTasksPendingRevocation = new HashMap<>();
        Map<String, Set<Integer>> newWarmupTasksPendingAssignment = new HashMap<>();

        boolean hasUnreleasedActiveTasks = computeAssignmentDifference(
            memberAssignedTasks.activeTasks(),
            targetAssignment.activeTasks(),
            newActiveAssignedTasks,
            newActiveTasksPendingRevocation,
            newActiveTasksPendingAssignment,
            (subtopologyId, partitionId) ->
                currentActiveTaskProcessId.apply(subtopologyId, partitionId) != null ||
                    currentStandbyTaskProcessIds.apply(subtopologyId, partitionId)
                        .contains(member.processId()) ||
                    currentWarmupTaskProcessIds.apply(subtopologyId, partitionId)
                        .contains(member.processId())
        );

        boolean hasUnreleasedStandbyTasks = computeAssignmentDifference(
            memberAssignedTasks.standbyTasks(),
            targetAssignment.standbyTasks(),
            newStandbyAssignedTasks,
            newStandbyTasksPendingRevocation,
            newStandbyTasksPendingAssignment,
            (subtopologyId, partitionId) ->
                Objects.equals(currentActiveTaskProcessId.apply(subtopologyId, partitionId),
                    member.processId()) ||
                    currentStandbyTaskProcessIds.apply(subtopologyId, partitionId)
                        .contains(member.processId()) ||
                    currentWarmupTaskProcessIds.apply(subtopologyId, partitionId)
                        .contains(member.processId())
        );

        boolean hasUnreleasedWarmupTasks = computeAssignmentDifference(
            memberAssignedTasks.warmupTasks(),
            targetAssignment.warmupTasks(),
            newWarmupAssignedTasks,
            newWarmupTasksPendingRevocation,
            newWarmupTasksPendingAssignment,
            (subtopologyId, partitionId) ->
                Objects.equals(currentActiveTaskProcessId.apply(subtopologyId, partitionId),
                    member.processId()) ||
                    currentStandbyTaskProcessIds.apply(subtopologyId, partitionId)
                        .contains(member.processId()) ||
                    currentWarmupTaskProcessIds.apply(subtopologyId, partitionId)
                        .contains(member.processId())
        );

        return buildNewMember(
            memberEpoch,
            new TasksTuple(
                newActiveTasksPendingRevocation,
                newStandbyTasksPendingRevocation,
                newWarmupTasksPendingRevocation
            ),
            new TasksTuple(
                newActiveAssignedTasks,
                newStandbyAssignedTasks,
                newWarmupAssignedTasks
            ),
            new TasksTuple(
                newActiveTasksPendingAssignment,
                newStandbyTasksPendingAssignment,
                newWarmupTasksPendingAssignment
            ),
            hasUnreleasedActiveTasks || hasUnreleasedStandbyTasks || hasUnreleasedWarmupTasks
        );
    }

    private StreamsGroupMember buildNewMember(final int memberEpoch,
                                              final TasksTuple newTasksPendingRevocation,
                                              final TasksTuple newAssignedTasks,
                                              final TasksTuple newTasksPendingAssignment,
                                              final boolean hasUnreleasedTasks) {

        final boolean hasTasksToBeRevoked =
            (!newTasksPendingRevocation.isEmpty())
                && (ownedTasks.isEmpty() || ownedTasks.get().containsAny(newTasksPendingRevocation));

        if (hasTasksToBeRevoked) {
            // If there are tasks to be revoked, the member remains in its current
            // epoch and requests the revocation of those tasks. It transitions to
            // the UNREVOKED_TASKS state to wait until the client acknowledges the
            // revocation of the tasks.
            return new StreamsGroupMember.Builder(member)
                .setState(MemberState.UNREVOKED_TASKS)
                .updateMemberEpoch(memberEpoch)
                .setAssignedTasks(newAssignedTasks)
                .setTasksPendingRevocation(newTasksPendingRevocation)
                .build();
        } else if (!newTasksPendingAssignment.isEmpty()) {
            // If there are tasks to be assigned, the member transitions to the
            // target epoch and requests the assignment of those tasks. Note that
            // the tasks are directly added to the assigned tasks set. The
            // member transitions to the STABLE state or to the UNRELEASED_TASKS
            // state depending on whether there are unreleased tasks or not.
            MemberState newState =
                hasUnreleasedTasks
                    ? MemberState.UNRELEASED_TASKS
                    : MemberState.STABLE;
            return new StreamsGroupMember.Builder(member)
                .setState(newState)
                .updateMemberEpoch(targetAssignmentEpoch)
                .setAssignedTasks(newAssignedTasks.merge(newTasksPendingAssignment))
                .setTasksPendingRevocation(TasksTuple.EMPTY)
                .build();
        } else if (hasUnreleasedTasks) {
            // If there are no tasks to be revoked nor to be assigned but some
            // tasks are not available yet, the member transitions to the target
            // epoch, to the UNRELEASED_TASKS state and waits.
            return new StreamsGroupMember.Builder(member)
                .setState(MemberState.UNRELEASED_TASKS)
                .updateMemberEpoch(targetAssignmentEpoch)
                .setAssignedTasks(newAssignedTasks)
                .setTasksPendingRevocation(TasksTuple.EMPTY)
                .build();
        } else {
            // Otherwise, the member transitions to the target epoch and to the
            // STABLE state.
            return new StreamsGroupMember.Builder(member)
                .setState(MemberState.STABLE)
                .updateMemberEpoch(targetAssignmentEpoch)
                .setAssignedTasks(newAssignedTasks)
                .setTasksPendingRevocation(TasksTuple.EMPTY)
                .build();
        }
    }
}
