/*
 * 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.assignor;

import org.junit.jupiter.api.Test;

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 static org.apache.kafka.common.utils.Utils.mkEntry;
import static org.apache.kafka.common.utils.Utils.mkMap;
import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertNotNull;
import static org.junit.jupiter.api.Assertions.assertThrows;
import static org.junit.jupiter.api.Assertions.assertTrue;


public class MockAssignorTest {

    private final MockAssignor assignor = new MockAssignor();

    @Test
    public void testZeroMembers() {

        TaskAssignorException ex = assertThrows(TaskAssignorException.class, () -> assignor.assign(
            new GroupSpecImpl(
                Map.of(),
                new HashMap<>()
            ),
            new TopologyDescriberImpl(5, List.of("test-subtopology"))
        ));

        assertEquals("No member available to assign task 0 of subtopology test-subtopology", ex.getMessage());
    }

    @Test
    public void testDoubleAssignment() {

        final AssignmentMemberSpec memberSpec1 = new AssignmentMemberSpec(
            Optional.empty(),
            Optional.empty(),
            Map.of("test-subtopology", Set.of(0)),
            Map.of(),
            Map.of(),
            "test-process",
            Map.of(),
            Map.of(),
            Map.of()
        );

        final AssignmentMemberSpec memberSpec2 = new AssignmentMemberSpec(
            Optional.empty(),
            Optional.empty(),
            Map.of("test-subtopology", Set.of(0)),
            Map.of(),
            Map.of(),
            "test-process",
            Map.of(),
            Map.of(),
            Map.of()
        );

        TaskAssignorException ex = assertThrows(TaskAssignorException.class, () -> assignor.assign(
            new GroupSpecImpl(
                Map.of("member1", memberSpec1, "member2", memberSpec2),
                new HashMap<>()
            ),
            new TopologyDescriberImpl(5, List.of("test-subtopology"))
        ));

        assertEquals("Task 0 of subtopology test-subtopology is assigned to multiple members", ex.getMessage());
    }

    @Test
    public void testBasicScenario() {

        final GroupAssignment result = assignor.assign(
            new GroupSpecImpl(
                Map.of(),
                new HashMap<>()
            ),
            new TopologyDescriberImpl(5, List.of())
        );

        assertEquals(0, result.members().size());
    }


    @Test
    public void testSingleMember() {

        final AssignmentMemberSpec memberSpec = new AssignmentMemberSpec(
            Optional.empty(),
            Optional.empty(),
            Map.of(),
            Map.of(),
            Map.of(),
            "test-process",
            Map.of(),
            Map.of(),
            Map.of()
        );

        final GroupAssignment result = assignor.assign(
            new GroupSpecImpl(
                Map.of("test_member", memberSpec),
                new HashMap<>()
            ),
            new TopologyDescriberImpl(4, List.of("test-subtopology"))
        );

        assertEquals(1, result.members().size());
        final MemberAssignment testMember = result.members().get("test_member");
        assertNotNull(testMember);
        assertEquals(mkMap(
            mkEntry("test-subtopology", Set.of(0, 1, 2, 3))
        ), testMember.activeTasks());
    }


    @Test
    public void testTwoMembersTwoSubtopologies() {

        final AssignmentMemberSpec memberSpec1 = new AssignmentMemberSpec(
            Optional.empty(),
            Optional.empty(),
            Map.of(),
            Map.of(),
            Map.of(),
            "test-process",
            Map.of(),
            Map.of(),
            Map.of()
        );

        final AssignmentMemberSpec memberSpec2 = new AssignmentMemberSpec(
            Optional.empty(),
            Optional.empty(),
            Map.of(),
            Map.of(),
            Map.of(),
            "test-process",
            Map.of(),
            Map.of(),
            Map.of()
        );

        final GroupAssignment result = assignor.assign(
            new GroupSpecImpl(
                mkMap(mkEntry("test_member1", memberSpec1), mkEntry("test_member2", memberSpec2)),
                new HashMap<>()
            ),
            new TopologyDescriberImpl(4, List.of("test-subtopology1", "test-subtopology2"))
        );

        final Map<String, Set<Integer>> expected1 = mkMap(
            mkEntry("test-subtopology1", Set.of(1, 3)),
            mkEntry("test-subtopology2", Set.of(1, 3))
        );
        final Map<String, Set<Integer>> expected2 = mkMap(
            mkEntry("test-subtopology1", Set.of(0, 2)),
            mkEntry("test-subtopology2", Set.of(0, 2))
        );

        assertEquals(2, result.members().size());
        final MemberAssignment testMember1 = result.members().get("test_member1");
        final MemberAssignment testMember2 = result.members().get("test_member2");
        assertNotNull(testMember1);
        assertNotNull(testMember2);
        assertTrue(expected1.equals(testMember1.activeTasks()) || expected2.equals(testMember1.activeTasks()));
        assertTrue(expected1.equals(testMember2.activeTasks()) || expected2.equals(testMember2.activeTasks()));
    }

    @Test
    public void testTwoMembersTwoSubtopologiesStickiness() {

        final AssignmentMemberSpec memberSpec1 = new AssignmentMemberSpec(
            Optional.empty(),
            Optional.empty(),
            mkMap(
                mkEntry("test-subtopology1", new HashSet<>(List.of(0, 2, 3))),
                mkEntry("test-subtopology2", new HashSet<>(List.of(0)))
            ),
            Map.of(),
            Map.of(),
            "test-process",
            Map.of(),
            Map.of(),
            Map.of()
        );

        final AssignmentMemberSpec memberSpec2 = new AssignmentMemberSpec(
            Optional.empty(),
            Optional.empty(),
            mkMap(
                mkEntry("test-subtopology1", new HashSet<>(List.of(1))),
                mkEntry("test-subtopology2", new HashSet<>(List.of(3)))
            ),
            Map.of(),
            Map.of(),
            "test-process",
            Map.of(),
            Map.of(),
            Map.of()
        );
        final GroupAssignment result = assignor.assign(
            new GroupSpecImpl(
                mkMap(mkEntry("test_member1", memberSpec1), mkEntry("test_member2", memberSpec2)),
                new HashMap<>()
            ),
            new TopologyDescriberImpl(4, List.of("test-subtopology1", "test-subtopology2"))
        );

        assertEquals(2, result.members().size());
        final MemberAssignment testMember1 = result.members().get("test_member1");
        final MemberAssignment testMember2 = result.members().get("test_member2");
        assertNotNull(testMember1);
        assertNotNull(testMember2);
        assertEquals(mkMap(
            mkEntry("test-subtopology1", Set.of(0, 2, 3)),
            mkEntry("test-subtopology2", Set.of(0))
        ), testMember1.activeTasks());
        assertEquals(mkMap(
            mkEntry("test-subtopology1", Set.of(1)),
            mkEntry("test-subtopology2", Set.of(1, 2, 3))
        ), testMember2.activeTasks());
    }

    private record TopologyDescriberImpl(int numPartitions, List<String> subtopologies) implements TopologyDescriber {

        @Override
        public List<String> subtopologies() {
            return subtopologies;
        }

        @Override
        public int maxNumInputPartitions(String subtopologyId) {
            return numPartitions;
        }

        @Override
        public boolean isStateful(String subtopologyId) {
            return false;
        }

    }

}
