/*
 * 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.streams.processor.internals.assignment;

import org.apache.kafka.streams.processor.TaskId;
import org.apache.kafka.streams.processor.assignment.ProcessId;

import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import org.slf4j.Logger;

import java.util.Map;
import java.util.Set;
import java.util.function.Function;
import java.util.stream.Collectors;

import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.TASK_0_0;
import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.TASK_0_1;
import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.TASK_0_2;
import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.getClientStatesMap;
import static org.apache.kafka.streams.processor.internals.assignment.StandbyTaskAssignmentUtils.computeTasksToRemainingStandbys;
import static org.apache.kafka.streams.processor.internals.assignment.StandbyTaskAssignmentUtils.pollClientAndMaybeAssignAndUpdateRemainingStandbyTasks;
import static org.hamcrest.CoreMatchers.equalTo;
import static org.hamcrest.MatcherAssert.assertThat;
import static org.junit.jupiter.api.Assertions.assertTrue;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.ArgumentMatchers.anyInt;
import static org.mockito.ArgumentMatchers.anyString;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.times;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.verifyNoInteractions;

public class StandbyTaskAssignmentUtilsTest {

    private static final Set<TaskId> ACTIVE_TASKS = Set.of(TASK_0_0, TASK_0_1, TASK_0_2);

    private Map<ProcessId, ClientState> clients;
    private ConstrainedPrioritySet clientsByTaskLoad;

    @BeforeEach
    public void setup() {

        clients = getClientStatesMap(ACTIVE_TASKS.stream().map(StandbyTaskAssignmentUtilsTest::mkState).toArray(ClientState[]::new));
        clientsByTaskLoad = new ConstrainedPrioritySet(
            (client, task) -> !clients.get(client).hasAssignedTask(task),
            client -> clients.get(client).assignedTaskLoad()
        );
        clientsByTaskLoad.offerAll(clients.keySet());
    }

    @Test
    public void shouldReturnNumberOfStandbyTasksThatWereNotAssigned() {
        final Logger logMock = mock(Logger.class);
        final int numStandbyReplicas = 3;
        final Map<TaskId, Integer> tasksToRemainingStandbys = computeTasksToRemainingStandbys(numStandbyReplicas, ACTIVE_TASKS);

        tasksToRemainingStandbys.keySet().forEach(taskId -> pollClientAndMaybeAssignAndUpdateRemainingStandbyTasks(numStandbyReplicas,
                                                                                                                   clients,
                                                                                                                   tasksToRemainingStandbys,
                                                                                                                   clientsByTaskLoad,
                                                                                                                   taskId,
                                                                                                                   logMock));

        assertTrue(ACTIVE_TASKS.stream().allMatch(activeTask -> tasksToRemainingStandbys.get(activeTask) == 1));
        assertTrue(areStandbyTasksPresentForAllActiveTasks(2));
        verify(logMock, times(ACTIVE_TASKS.size())).warn(anyString(), anyInt(), anyInt(), any());
    }

    @Test
    public void shouldReturnZeroWhenAllStandbyTasksWereSuccessfullyAssigned() {
        final Logger logMock = mock(Logger.class);
        final int numStandbyReplicas = 1;
        final Map<TaskId, Integer> tasksToRemainingStandbys = computeTasksToRemainingStandbys(numStandbyReplicas, ACTIVE_TASKS);

        tasksToRemainingStandbys.keySet().forEach(taskId -> pollClientAndMaybeAssignAndUpdateRemainingStandbyTasks(numStandbyReplicas,
                                                                                                                   clients,
                                                                                                                   tasksToRemainingStandbys,
                                                                                                                   clientsByTaskLoad,
                                                                                                                   taskId,
                                                                                                                   logMock));

        assertTrue(ACTIVE_TASKS.stream().allMatch(activeTask -> tasksToRemainingStandbys.get(activeTask) == 0));
        assertTrue(areStandbyTasksPresentForAllActiveTasks(1));
        verifyNoInteractions(logMock);
    }

    @Test
    public void shouldComputeTasksToRemainingStandbys() {
        assertThat(
            computeTasksToRemainingStandbys(0, ACTIVE_TASKS),
            equalTo(
                ACTIVE_TASKS.stream().collect(Collectors.toMap(Function.identity(), it -> 0))
            )
        );
        assertThat(
            computeTasksToRemainingStandbys(5, ACTIVE_TASKS),
            equalTo(
                ACTIVE_TASKS.stream().collect(Collectors.toMap(Function.identity(), it -> 5))
            )
        );
    }

    private boolean areStandbyTasksPresentForAllActiveTasks(final int expectedNumberOfStandbyTasks) {
        return ACTIVE_TASKS.stream().allMatch(taskId -> clients.values()
                                                               .stream()
                                                               .filter(client -> client.hasStandbyTask(taskId))
                                                               .count() == expectedNumberOfStandbyTasks);
    }

    private static ClientState mkState(final TaskId... activeTasks) {
        return mkState(1, activeTasks);
    }

    private static ClientState mkState(final int capacity, final TaskId... activeTasks) {
        final ClientState clientState = new ClientState(capacity);
        for (final TaskId activeTask : activeTasks) {
            clientState.assignActive(activeTask);
        }
        return clientState;
    }
}