/*
 * 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.clients.consumer.internals;

import org.junit.jupiter.api.Test;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.Set;
import java.util.UUID;
import java.util.stream.Collectors;

import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertFalse;
import static org.junit.jupiter.api.Assertions.assertNotEquals;
import static org.junit.jupiter.api.Assertions.assertNotSame;
import static org.junit.jupiter.api.Assertions.assertThrows;
import static org.junit.jupiter.api.Assertions.assertTrue;

public class StreamsRebalanceDataTest {

    @Test
    public void testTaskIdEqualsAndHashCode() {
        final StreamsRebalanceData.TaskId task = new StreamsRebalanceData.TaskId("subtopologyId1", 1);
        final StreamsRebalanceData.TaskId taskEqual = new StreamsRebalanceData.TaskId(task.subtopologyId(), task.partitionId());
        final StreamsRebalanceData.TaskId taskUnequalSubtopology = new StreamsRebalanceData.TaskId(task.subtopologyId() + "1", task.partitionId());
        final StreamsRebalanceData.TaskId taskUnequalPartition = new StreamsRebalanceData.TaskId(task.subtopologyId(), task.partitionId() + 1);

        assertEquals(task, taskEqual);
        assertEquals(task.hashCode(), taskEqual.hashCode());
        assertNotEquals(task, taskUnequalSubtopology);
        assertNotEquals(task.hashCode(), taskUnequalSubtopology.hashCode());
        assertNotEquals(task, taskUnequalPartition);
        assertNotEquals(task.hashCode(), taskUnequalSubtopology.hashCode());
    }

    @Test
    public void taskIdShouldNotAcceptNulls() {
        final Exception exception = assertThrows(NullPointerException.class, () -> new StreamsRebalanceData.TaskId(null, 1));
        assertEquals("Subtopology ID cannot be null", exception.getMessage());
    }

    @Test
    public void testTaskIdCompareTo() {
        final StreamsRebalanceData.TaskId task = new StreamsRebalanceData.TaskId("subtopologyId1", 1);

        assertEquals(0, task.compareTo(new StreamsRebalanceData.TaskId(task.subtopologyId(), task.partitionId())));
        assertTrue(task.compareTo(new StreamsRebalanceData.TaskId(task.subtopologyId() + "1", task.partitionId())) < 0);
        assertTrue(task.compareTo(new StreamsRebalanceData.TaskId(task.subtopologyId(), task.partitionId() + 1)) < 0);
        assertTrue(new StreamsRebalanceData.TaskId(task.subtopologyId() + "1", task.partitionId()).compareTo(task) > 0);
        assertTrue(new StreamsRebalanceData.TaskId(task.subtopologyId(), task.partitionId() + 1).compareTo(task) > 0);
    }

    @Test
    public void emptyAssignmentShouldNotBeModifiable() {
        final StreamsRebalanceData.Assignment emptyAssignment = StreamsRebalanceData.Assignment.EMPTY;

        assertThrows(
            UnsupportedOperationException.class,
            () -> emptyAssignment.activeTasks().add(new StreamsRebalanceData.TaskId("subtopologyId1", 1))
        );
        assertThrows(
            UnsupportedOperationException.class,
            () -> emptyAssignment.standbyTasks().add(new StreamsRebalanceData.TaskId("subtopologyId1", 1))
        );
        assertThrows(
            UnsupportedOperationException.class,
            () -> emptyAssignment.warmupTasks().add(new StreamsRebalanceData.TaskId("subtopologyId1", 1))
        );
    }

    @Test
    public void assignmentShouldNotBeModifiable() {
        final StreamsRebalanceData.Assignment assignment = new StreamsRebalanceData.Assignment(
            Set.of(new StreamsRebalanceData.TaskId("subtopologyId1", 1)),
            Set.of(new StreamsRebalanceData.TaskId("subtopologyId1", 2)),
            Set.of(new StreamsRebalanceData.TaskId("subtopologyId1", 3)),
            true
        );

        assertThrows(
            UnsupportedOperationException.class,
            () -> assignment.activeTasks().add(new StreamsRebalanceData.TaskId("subtopologyId2", 1))
        );
        assertThrows(
            UnsupportedOperationException.class,
            () -> assignment.standbyTasks().add(new StreamsRebalanceData.TaskId("subtopologyId2", 2))
        );
        assertThrows(
            UnsupportedOperationException.class,
            () -> assignment.warmupTasks().add(new StreamsRebalanceData.TaskId("subtopologyId2", 3))
        );
    }

    @Test
    public void assignmentShouldNotAcceptNulls() {
        final Exception exception1 = assertThrows(NullPointerException.class, () -> new StreamsRebalanceData.Assignment(null, Set.of(), Set.of(), true));
        assertEquals("Active tasks cannot be null", exception1.getMessage());
        final Exception exception2 = assertThrows(NullPointerException.class, () -> new StreamsRebalanceData.Assignment(Set.of(), null, Set.of(), true));
        assertEquals("Standby tasks cannot be null", exception2.getMessage());
        final Exception exception3 = assertThrows(NullPointerException.class, () -> new StreamsRebalanceData.Assignment(Set.of(), Set.of(), null, true));
        assertEquals("Warmup tasks cannot be null", exception3.getMessage());
    }

    @Test
    public void testAssignmentEqualsAndHashCode() {
        final StreamsRebalanceData.TaskId additionalTask = new StreamsRebalanceData.TaskId("subtopologyId2", 1);
        final StreamsRebalanceData.Assignment assignment = new StreamsRebalanceData.Assignment(
            Set.of(new StreamsRebalanceData.TaskId("subtopologyId1", 1)),
            Set.of(new StreamsRebalanceData.TaskId("subtopologyId1", 2)),
            Set.of(new StreamsRebalanceData.TaskId("subtopologyId1", 3)),
            true
        );
        final StreamsRebalanceData.Assignment assignmentEqual = new StreamsRebalanceData.Assignment(
            assignment.activeTasks(),
            assignment.standbyTasks(),
            assignment.warmupTasks(),
            assignment.isGroupReady()
        );
        Set<StreamsRebalanceData.TaskId> unequalActiveTasks = new HashSet<>(assignment.activeTasks());
        unequalActiveTasks.add(additionalTask);
        final StreamsRebalanceData.Assignment assignmentUnequalActiveTasks = new StreamsRebalanceData.Assignment(
            unequalActiveTasks,
            assignment.standbyTasks(),
            assignment.warmupTasks(),
            assignment.isGroupReady()
        );
        Set<StreamsRebalanceData.TaskId> unequalStandbyTasks = new HashSet<>(assignment.standbyTasks());
        unequalStandbyTasks.add(additionalTask);
        final StreamsRebalanceData.Assignment assignmentUnequalStandbyTasks = new StreamsRebalanceData.Assignment(
            assignment.activeTasks(),
            unequalStandbyTasks,
            assignment.warmupTasks(),
            assignment.isGroupReady()
        );
        Set<StreamsRebalanceData.TaskId> unequalWarmupTasks = new HashSet<>(assignment.warmupTasks());
        unequalWarmupTasks.add(additionalTask);
        final StreamsRebalanceData.Assignment assignmentUnequalWarmupTasks = new StreamsRebalanceData.Assignment(
            assignment.activeTasks(),
            assignment.standbyTasks(),
            unequalWarmupTasks,
            assignment.isGroupReady()
        );
        final StreamsRebalanceData.Assignment assignmentUnequalIsGroupReady = new StreamsRebalanceData.Assignment(
            assignment.activeTasks(),
            assignment.standbyTasks(),
            assignment.warmupTasks(),
            !assignment.isGroupReady()
        );

        assertEquals(assignment, assignmentEqual);
        assertNotEquals(assignment, assignmentUnequalActiveTasks);
        assertNotEquals(assignment, assignmentUnequalStandbyTasks);
        assertNotEquals(assignment, assignmentUnequalWarmupTasks);
        assertNotEquals(assignment, assignmentUnequalIsGroupReady);
        assertEquals(assignment.hashCode(), assignmentEqual.hashCode());
        assertNotEquals(assignment.hashCode(), assignmentUnequalActiveTasks.hashCode());
        assertNotEquals(assignment.hashCode(), assignmentUnequalStandbyTasks.hashCode());
        assertNotEquals(assignment.hashCode(), assignmentUnequalWarmupTasks.hashCode());
        assertNotEquals(assignment.hashCode(), assignmentUnequalIsGroupReady.hashCode());
    }

    @Test
    public void shouldCopyAssignment() {
        final StreamsRebalanceData.Assignment assignment = new StreamsRebalanceData.Assignment(
            Set.of(new StreamsRebalanceData.TaskId("subtopologyId1", 1)),
            Set.of(new StreamsRebalanceData.TaskId("subtopologyId1", 2)),
            Set.of(new StreamsRebalanceData.TaskId("subtopologyId1", 3)),
            true
        );

        final StreamsRebalanceData.Assignment copy = assignment.copy();

        assertEquals(assignment, copy);
        assertNotSame(assignment, copy);
    }

    @Test
    public void shouldCopyEmptyAssignment() {
        final StreamsRebalanceData.Assignment emptyAssignment = StreamsRebalanceData.Assignment.EMPTY;

        final StreamsRebalanceData.Assignment copy = emptyAssignment.copy();

        assertEquals(emptyAssignment, copy);
        assertNotSame(emptyAssignment, copy);
    }

    @Test
    public void subtopologyShouldNotAcceptNulls() {
        final Exception exception1 = assertThrows(
            NullPointerException.class,
            () -> new StreamsRebalanceData.Subtopology(null, Set.of(), Map.of(), Map.of(), List.of())
        );
        assertEquals("Subtopology ID cannot be null", exception1.getMessage());
        final Exception exception2 = assertThrows(
            NullPointerException.class,
            () -> new StreamsRebalanceData.Subtopology(Set.of(), null, Map.of(), Map.of(), List.of())
        );
        assertEquals("Repartition sink topics cannot be null", exception2.getMessage());
        final Exception exception3 = assertThrows(
            NullPointerException.class,
            () -> new StreamsRebalanceData.Subtopology(Set.of(), Set.of(), null, Map.of(), List.of())
        );
        assertEquals("Repartition source topics cannot be null", exception3.getMessage());
        final Exception exception4 = assertThrows(
            NullPointerException.class,
            () -> new StreamsRebalanceData.Subtopology(Set.of(), Set.of(), Map.of(), null, List.of())
        );
        assertEquals("State changelog topics cannot be null", exception4.getMessage());
        final Exception exception5 = assertThrows(
            NullPointerException.class,
            () -> new StreamsRebalanceData.Subtopology(Set.of(), Set.of(), Map.of(), Map.of(), null)
        );
        assertEquals("Co-partition groups cannot be null", exception5.getMessage());
    }

    @Test
    public void subtopologyShouldNotBeModifiable() {
        final StreamsRebalanceData.Subtopology subtopology = new StreamsRebalanceData.Subtopology(
            Set.of("sourceTopic1"),
            Set.of("repartitionSinkTopic1"),
            Map.of("repartitionSourceTopic1", new StreamsRebalanceData.TopicInfo(Optional.of(1), Optional.of((short) 1), Map.of()))
                .entrySet().stream()
                .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue)),
            Map.of("stateChangelogTopic1", new StreamsRebalanceData.TopicInfo(Optional.of(0), Optional.of((short) 1), Map.of()))
                .entrySet().stream()
                .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue)),
            new ArrayList<>(List.of(Set.of("sourceTopic1")))
        );

        assertThrows(
            UnsupportedOperationException.class,
            () -> subtopology.sourceTopics().add("sourceTopic2")
        );
        assertThrows(
            UnsupportedOperationException.class,
            () -> subtopology.repartitionSinkTopics().add("repartitionSinkTopic2")
        );
        assertThrows(
            UnsupportedOperationException.class,
            () -> subtopology.repartitionSourceTopics().put("repartitionSourceTopic2", new StreamsRebalanceData.TopicInfo(Optional.of(1), Optional.of((short) 1), Map.of()))
        );
        assertThrows(
            UnsupportedOperationException.class,
            () -> subtopology.stateChangelogTopics().put("stateChangelogTopic2", new StreamsRebalanceData.TopicInfo(Optional.of(0), Optional.of((short) 1), Map.of()))
        );
        assertThrows(
            UnsupportedOperationException.class,
            () -> subtopology.copartitionGroups().add(Set.of("sourceTopic2"))
        );
    }

    @Test
    public void topicInfoShouldNotAcceptNulls() {
        final Exception exception1 = assertThrows(
            NullPointerException.class,
            () -> new StreamsRebalanceData.TopicInfo(null, Optional.of((short) 1), Map.of())
        );
        assertEquals("Number of partitions cannot be null", exception1.getMessage());
        final Exception exception2 = assertThrows(
            NullPointerException.class,
            () -> new StreamsRebalanceData.TopicInfo(Optional.of(1), null, Map.of())
        );
        assertEquals("Replication factor cannot be null", exception2.getMessage());
        final Exception exception3 = assertThrows(
            NullPointerException.class,
            () -> new StreamsRebalanceData.TopicInfo(Optional.of(1), Optional.of((short) 1), null)
        );
        assertEquals("Additional topic configs cannot be null", exception3.getMessage());
    }

    @Test
    public void streamsRebalanceDataShouldNotHaveModifiableSubtopologiesAndClientTags() {
        final UUID processId = UUID.randomUUID();
        final Optional<StreamsRebalanceData.HostInfo> endpoint = Optional.of(new StreamsRebalanceData.HostInfo("localhost", 9090));
        final Map<String, StreamsRebalanceData.Subtopology> subtopologies = new HashMap<>();
        final Map<String, String> clientTags = Map.of("clientTag1", "clientTagValue1");
        final StreamsRebalanceData streamsRebalanceData = new StreamsRebalanceData(
            processId,
            endpoint,
            subtopologies,
            clientTags
        );

        assertThrows(
            UnsupportedOperationException.class,
            () -> streamsRebalanceData.subtopologies().put("subtopologyId2", new StreamsRebalanceData.Subtopology(
                Set.of(),
                Set.of(),
                Map.of(),
                Map.of(),
                List.of()
            ))
        );
        assertThrows(
            UnsupportedOperationException.class,
            () -> streamsRebalanceData.clientTags().put("clientTag1", "clientTagValue2")
        );
    }

    @Test
    public void streamsRebalanceDataShouldNotAcceptNullProcessId() {
        final Optional<StreamsRebalanceData.HostInfo> endpoint = Optional.of(new StreamsRebalanceData.HostInfo("localhost", 9090));
        final Map<String, StreamsRebalanceData.Subtopology> subtopologies = new HashMap<>();
        final Map<String, String> clientTags = Map.of("clientTag1", "clientTagValue1");

        final Exception exception = assertThrows(
            NullPointerException.class,
            () -> new StreamsRebalanceData(
                null,
                endpoint,
                subtopologies,
                clientTags
            )
        );
        assertEquals("Process ID cannot be null", exception.getMessage());
    }

    @Test
    public void streamsRebalanceDataShouldNotAcceptNullHostInfo() {
        final UUID processId = UUID.randomUUID();
        final Map<String, StreamsRebalanceData.Subtopology> subtopologies = new HashMap<>();
        final Map<String, String> clientTags = Map.of("clientTag1", "clientTagValue1");

        final Exception exception = assertThrows(
            NullPointerException.class,
            () -> new StreamsRebalanceData(
                processId,
                null,
                subtopologies,
                clientTags
            )
        );
        assertEquals("Endpoint cannot be null", exception.getMessage());
    }

    @Test
    public void streamsRebalanceDataShouldNotAcceptNullSubtopologies() {
        final UUID processId = UUID.randomUUID();
        final Optional<StreamsRebalanceData.HostInfo> endpoint = Optional.of(new StreamsRebalanceData.HostInfo("localhost", 9090));
        final Map<String, String> clientTags = Map.of("clientTag1", "clientTagValue1");

        final Exception exception = assertThrows(
            NullPointerException.class,
            () -> new StreamsRebalanceData(
                processId,
                endpoint,
                null,
                clientTags
            )
        );
        assertEquals("Subtopologies cannot be null", exception.getMessage());
    }

    @Test
    public void streamsRebalanceDataShouldNotAcceptNullClientTags() {
        final UUID processId = UUID.randomUUID();
        final Optional<StreamsRebalanceData.HostInfo> endpoint = Optional.of(new StreamsRebalanceData.HostInfo("localhost", 9090));
        final Map<String, StreamsRebalanceData.Subtopology> subtopologies = new HashMap<>();

        final Exception exception = assertThrows(
            NullPointerException.class,
            () -> new StreamsRebalanceData(
                processId,
                endpoint,
                subtopologies,
                null
            )
        );
        assertEquals("Client tags cannot be null", exception.getMessage());
    }

    @Test
    public void streamsRebalanceDataShouldBeConstructedWithEmptyAssignment() {
        final UUID processId = UUID.randomUUID();
        final Optional<StreamsRebalanceData.HostInfo> endpoint = Optional.of(new StreamsRebalanceData.HostInfo("localhost", 9090));
        final Map<String, StreamsRebalanceData.Subtopology> subtopologies = new HashMap<>();
        final Map<String, String> clientTags = Map.of("clientTag1", "clientTagValue1");
        final StreamsRebalanceData streamsRebalanceData = new StreamsRebalanceData(
            processId,
            endpoint,
            subtopologies,
            clientTags
        );

        assertEquals(StreamsRebalanceData.Assignment.EMPTY, streamsRebalanceData.reconciledAssignment());
    }

    @Test
    public void streamsRebalanceDataShouldBeConstructedWithEmptyPartitionsByHost() {
        final UUID processId = UUID.randomUUID();
        final Optional<StreamsRebalanceData.HostInfo> endpoint = Optional.of(new StreamsRebalanceData.HostInfo("localhost", 9090));
        final Map<String, StreamsRebalanceData.Subtopology> subtopologies = new HashMap<>();
        final Map<String, String> clientTags = Map.of("clientTag1", "clientTagValue1");
        final StreamsRebalanceData streamsRebalanceData = new StreamsRebalanceData(
            processId,
            endpoint,
            subtopologies,
            clientTags
        );

        assertTrue(streamsRebalanceData.partitionsByHost().isEmpty());
    }

    @Test
    public void streamsRebalanceDataShouldBeConstructedWithShutDownRequestedSetFalse() {
        final UUID processId = UUID.randomUUID();
        final Optional<StreamsRebalanceData.HostInfo> endpoint = Optional.of(new StreamsRebalanceData.HostInfo("localhost", 9090));
        final Map<String, StreamsRebalanceData.Subtopology> subtopologies = new HashMap<>();
        final Map<String, String> clientTags = Map.of("clientTag1", "clientTagValue1");
        final StreamsRebalanceData streamsRebalanceData = new StreamsRebalanceData(
            processId,
            endpoint,
            subtopologies,
            clientTags
        );

        assertFalse(streamsRebalanceData.shutdownRequested());
    }

    @Test
    public void streamsRebalanceDataShouldBeConstructedWithEmptyStatuses() {
        final UUID processId = UUID.randomUUID();
        final Optional<StreamsRebalanceData.HostInfo> endpoint = Optional.of(new StreamsRebalanceData.HostInfo("localhost", 9090));
        final Map<String, StreamsRebalanceData.Subtopology> subtopologies = Map.of();
        final Map<String, String> clientTags = Map.of("clientTag1", "clientTagValue1");
        final StreamsRebalanceData streamsRebalanceData = new StreamsRebalanceData(
            processId,
            endpoint,
            subtopologies,
            clientTags
        );

        assertTrue(streamsRebalanceData.statuses().isEmpty());
    }

    @Test
    public void streamsRebalanceDataShouldBeConstructedWithHeartbeatIntervalMsSetToMinusOne() {
        final UUID processId = UUID.randomUUID();
        final Optional<StreamsRebalanceData.HostInfo> endpoint = Optional.of(new
                StreamsRebalanceData.HostInfo("localhost", 9090));
        final Map<String, StreamsRebalanceData.Subtopology> subtopologies = Map.of();
        final Map<String, String> clientTags = Map.of("clientTag1",
                "clientTagValue1");
        final StreamsRebalanceData streamsRebalanceData = new StreamsRebalanceData(
                processId,
                endpoint,
                subtopologies,
                clientTags
        );

        assertEquals(-1, streamsRebalanceData.heartbeatIntervalMs());
    }

    @Test
    public void streamsRebalanceDataShouldBeAbleToUpdateHeartbeatIntervalMs() {
        final UUID processId = UUID.randomUUID();
        final Optional<StreamsRebalanceData.HostInfo> endpoint = Optional.of(new
                StreamsRebalanceData.HostInfo("localhost", 9090));
        final Map<String, StreamsRebalanceData.Subtopology> subtopologies = Map.of();
        final Map<String, String> clientTags = Map.of("clientTag1",
                "clientTagValue1");
        final StreamsRebalanceData streamsRebalanceData = new StreamsRebalanceData(
                processId,
                endpoint,
                subtopologies,
                clientTags
        );

        streamsRebalanceData.setHeartbeatIntervalMs(1000);
        assertEquals(1000, streamsRebalanceData.heartbeatIntervalMs());
    }

}
