/*
 * 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.server.share.session;

import org.apache.kafka.common.Uuid;
import org.apache.kafka.common.utils.ImplicitLinkedHashCollection;
import org.apache.kafka.server.share.CachedSharePartition;
import org.apache.kafka.server.share.ShareGroupListener;
import org.apache.kafka.test.TestUtils;

import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import org.mockito.Mockito;

import java.util.Iterator;
import java.util.List;

import static org.apache.kafka.server.share.fetch.ShareFetchTestUtils.clearYammerMetrics;
import static org.apache.kafka.server.share.fetch.ShareFetchTestUtils.yammerMetricValue;
import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertFalse;
import static org.junit.jupiter.api.Assertions.assertNotNull;
import static org.junit.jupiter.api.Assertions.assertNull;

public class ShareSessionCacheTest {

    @BeforeEach
    public void setUp() {
        clearYammerMetrics();
    }

    @Test
    public void testShareSessionCache() throws InterruptedException {
        ShareSessionCache cache = new ShareSessionCache(3);
        assertEquals(0, cache.size());
        ShareSessionKey key1 = cache.maybeCreateSession("grp", Uuid.randomUuid(), mockedSharePartitionMap(10), "conn-1");
        ShareSessionKey key2 = cache.maybeCreateSession("grp", Uuid.randomUuid(), mockedSharePartitionMap(20), "conn-2");
        ShareSessionKey key3 = cache.maybeCreateSession("grp", Uuid.randomUuid(), mockedSharePartitionMap(30), "conn-3");
        assertNull(cache.maybeCreateSession("grp", Uuid.randomUuid(), mockedSharePartitionMap(40), "conn-4"));
        assertNull(cache.maybeCreateSession("grp", Uuid.randomUuid(), mockedSharePartitionMap(5), "conn-5"));
        assertShareCacheContains(cache, List.of(key1, key2, key3));

        assertMetricsValues(3, 60, 0, cache);
    }

    @Test
    public void testResizeCachedSessions() throws InterruptedException {
        ShareSessionCache cache = new ShareSessionCache(2);
        assertEquals(0, cache.size());
        assertEquals(0, cache.totalPartitions());
        ShareSessionKey key1 = cache.maybeCreateSession("grp", Uuid.randomUuid(), mockedSharePartitionMap(2), "conn-1");
        assertNotNull(key1);
        assertShareCacheContains(cache, List.of(key1));
        ShareSession session1 = cache.get(key1);
        assertEquals(2, session1.size());
        assertEquals(2, cache.totalPartitions());
        assertEquals(1, cache.size());

        assertMetricsValues(1, 2, 0, cache);

        ShareSessionKey key2 = cache.maybeCreateSession("grp", Uuid.randomUuid(), mockedSharePartitionMap(4), "conn-2");
        assertNotNull(key2);
        assertShareCacheContains(cache, List.of(key1, key2));
        ShareSession session2 = cache.get(key2);
        assertEquals(6, cache.totalPartitions());
        assertEquals(2, cache.size());
        cache.updateNumPartitions(session1);
        cache.updateNumPartitions(session2);

        assertMetricsValues(2, 6, 0, cache);

        ShareSessionKey key3 = cache.maybeCreateSession("grp", Uuid.randomUuid(), mockedSharePartitionMap(5), "conn-3");
        assertNull(key3);
        assertShareCacheContains(cache, List.of(key1, key2));
        assertEquals(6, cache.totalPartitions());
        assertEquals(2, cache.size());
        cache.remove(key1);
        assertShareCacheContains(cache, List.of(key2));
        assertEquals(1, cache.size());
        assertEquals(4, cache.totalPartitions());

        assertMetricsValues(1, 4, 0, cache);

        Iterator<CachedSharePartition> iterator = session2.partitionMap().iterator();
        iterator.next();
        iterator.remove();
        // Session size should get updated as it's backed by the partition map.
        assertEquals(3, session2.size());
        // Cached size should not get updated as it shall update on touch.
        assertEquals(4, session2.cachedSize());
        assertEquals(4, cache.totalPartitions());
        // Touch the session to update the changes in cache and session's cached size.
        cache.updateNumPartitions(session2);
        assertEquals(3, session2.cachedSize());
        assertEquals(3, cache.totalPartitions());

        assertMetricsValues(1, 3, 0, cache);
    }

    @Test
    public void testRemoveConnection() throws InterruptedException {
        ShareSessionCache cache = new ShareSessionCache(3);
        assertEquals(0, cache.size());
        ShareSessionKey key1 = cache.maybeCreateSession("grp", Uuid.randomUuid(), mockedSharePartitionMap(1), "conn-1");
        ShareSessionKey key2 = cache.maybeCreateSession("grp", Uuid.randomUuid(), mockedSharePartitionMap(2), "conn-2");
        ShareSessionKey key3 = cache.maybeCreateSession("grp", Uuid.randomUuid(), mockedSharePartitionMap(3), "conn-3");

        assertMetricsValues(3, 6, 0, cache);

        // Since cache size is now equal to max entries allowed(3), no new session can be created.
        assertNull(cache.maybeCreateSession("grp", Uuid.randomUuid(), mockedSharePartitionMap(40), "conn-4"));
        assertNull(cache.maybeCreateSession("grp", Uuid.randomUuid(), mockedSharePartitionMap(5), "conn-5"));
        assertShareCacheContains(cache, List.of(key1, key2, key3));

        assertMetricsValues(3, 6, 0, cache);

        // Simulating the disconnection of client with connection id conn-1
        cache.connectionDisconnectListener().onDisconnect("conn-1");
        assertShareCacheContains(cache, List.of(key2, key3));

        assertMetricsValues(2, 5, 1, cache);

        // Since one client got disconnected, we can add another one now
        ShareSessionKey key4 = cache.maybeCreateSession("grp", Uuid.randomUuid(), mockedSharePartitionMap(4), "conn-6");
        assertShareCacheContains(cache, List.of(key2, key3, key4));

        assertMetricsValues(3, 9, 1, cache);
    }

    @Test
    public void testRemoveAllSessions() {
        ShareSessionCache cache = new ShareSessionCache(3);
        assertEquals(0, cache.size());
        assertEquals(0, cache.totalPartitions());
        cache.maybeCreateSession("grp", Uuid.randomUuid(), mockedSharePartitionMap(10), "conn-1");
        cache.maybeCreateSession("grp", Uuid.randomUuid(), mockedSharePartitionMap(20), "conn-2");
        cache.maybeCreateSession("grp", Uuid.randomUuid(), mockedSharePartitionMap(30), "conn-3");
        assertEquals(3, cache.size());
        assertEquals(60, cache.totalPartitions());
        cache.removeAllSessions();
        assertEquals(0, cache.size());
        assertEquals(0, cache.totalPartitions());
    }

    @Test
    public void testShareGroupListenerEvents() {
        ShareGroupListener mockListener = Mockito.mock(ShareGroupListener.class);
        ShareSessionCache cache = new ShareSessionCache(3);
        cache.registerShareGroupListener(mockListener);

        String groupId = "grp";
        Uuid memberId1 = Uuid.randomUuid();
        Uuid memberId2 = Uuid.randomUuid();
        ShareSessionKey key1 = cache.maybeCreateSession(groupId, memberId1, mockedSharePartitionMap(1), "conn-1");
        ShareSessionKey key2 = cache.maybeCreateSession(groupId, memberId2, mockedSharePartitionMap(1), "conn-2");

        // Verify member count is tracked
        assertEquals(2, cache.size());
        assertNotNull(cache.get(key1));
        assertNotNull(cache.get(key2));
        assertEquals(2, cache.numMembers(groupId));

        // Remove session and verify listener are not called as connection disconnect listener didn't
        // remove the session.
        cache.remove(key1);
        Mockito.verify(mockListener, Mockito.times(0)).onMemberLeave(groupId, memberId1);
        Mockito.verify(mockListener, Mockito.times(0)).onGroupEmpty(groupId);
        // Verify member count is updated
        assertEquals(1, cache.numMembers(groupId));

        // Re-create session for memberId1.
        cache.maybeCreateSession(groupId, memberId1, mockedSharePartitionMap(1), "conn-1");
        assertEquals(2, cache.numMembers(groupId));

        // Simulate connection disconnect for memberId1.
        cache.connectionDisconnectListener().onDisconnect("conn-1");
        // Verify only member leave event is triggered for memberId1. Empty group event should not be triggered.
        Mockito.verify(mockListener, Mockito.times(1)).onMemberLeave(groupId, memberId1);
        Mockito.verify(mockListener, Mockito.times(0)).onMemberLeave(groupId, memberId2);
        Mockito.verify(mockListener, Mockito.times(0)).onGroupEmpty(groupId);
        assertEquals(1, cache.numMembers(groupId));

        // Simulate connection disconnect for memberId2.
        cache.connectionDisconnectListener().onDisconnect("conn-2");
        // Verify both member leave event and empty group event should be triggered.
        Mockito.verify(mockListener, Mockito.times(1)).onMemberLeave(groupId, memberId1);
        Mockito.verify(mockListener, Mockito.times(1)).onMemberLeave(groupId, memberId2);
        Mockito.verify(mockListener, Mockito.times(1)).onGroupEmpty(groupId);
        assertNull(cache.numMembers(groupId));
    }

    @Test
    public void testShareGroupListenerEventsMultipleGroups() {
        ShareGroupListener mockListener = Mockito.mock(ShareGroupListener.class);
        ShareSessionCache cache = new ShareSessionCache(3);
        cache.registerShareGroupListener(mockListener);

        String groupId1 = "grp1";
        String groupId2 = "grp2";
        Uuid memberId1 = Uuid.randomUuid();
        Uuid memberId2 = Uuid.randomUuid();
        ShareSessionKey key1 = cache.maybeCreateSession(groupId1, memberId1, mockedSharePartitionMap(1), "conn-1");
        ShareSessionKey key2 = cache.maybeCreateSession(groupId2, memberId2, mockedSharePartitionMap(1), "conn-2");

        // Verify member count is tracked
        assertEquals(2, cache.size());
        assertNotNull(cache.get(key1));
        assertNotNull(cache.get(key2));
        assertEquals(1, cache.numMembers(groupId1));
        assertEquals(1, cache.numMembers(groupId2));

        // Remove session for group1 and verify listeners are only called for group1.
        cache.connectionDisconnectListener().onDisconnect("conn-1");
        Mockito.verify(mockListener, Mockito.times(1)).onMemberLeave(groupId1, memberId1);
        Mockito.verify(mockListener, Mockito.times(1)).onGroupEmpty(groupId1);
        // Listener should not be called for group2.
        Mockito.verify(mockListener, Mockito.times(0)).onMemberLeave(groupId2, memberId2);
        Mockito.verify(mockListener, Mockito.times(0)).onGroupEmpty(groupId2);
        // Verify member count is updated.
        assertNull(cache.numMembers(groupId1));
        assertEquals(1, cache.numMembers(groupId2));
    }

    @Test
    public void testNoShareGroupListenerRegistered() {
        ShareSessionCache cache = new ShareSessionCache(3);

        String groupId = "grp";
        Uuid memberId = Uuid.randomUuid();
        ShareSessionKey key = cache.maybeCreateSession(groupId, memberId, mockedSharePartitionMap(1), "conn-1");

        // Verify member count is still tracked even without listener
        assertEquals(1, cache.numMembers(groupId));
        assertNotNull(cache.get(key));

        // Remove session should not throw any exceptions.
        cache.connectionDisconnectListener().onDisconnect("conn-1");
        assertNull(cache.numMembers(groupId));
    }

    private ImplicitLinkedHashCollection<CachedSharePartition> mockedSharePartitionMap(int size) {
        ImplicitLinkedHashCollection<CachedSharePartition> cacheMap = new
                ImplicitLinkedHashCollection<>(size);
        for (int i = 0; i < size; i++)
            cacheMap.add(new CachedSharePartition("test", Uuid.randomUuid(), i, false));
        return cacheMap;
    }

    private void assertShareCacheContains(ShareSessionCache cache,
                                         List<ShareSessionKey> sessionKeys) {
        int i = 0;
        assertEquals(sessionKeys.size(), cache.size());
        for (ShareSessionKey sessionKey : sessionKeys) {
            assertFalse(cache.get(sessionKey).isEmpty(),
                    "Missing session " + ++i + " out of " + sessionKeys.size() + " ( " + sessionKey + " )");
        }
    }

    private void assertMetricsValues(
        int shareSessionsCount,
        int sharePartitionsCount,
        int evictionsCount,
        ShareSessionCache cache
    ) throws InterruptedException {
        TestUtils.waitForCondition(() -> yammerMetricValue(ShareSessionCache.SHARE_SESSIONS_COUNT).intValue() == shareSessionsCount,
            "Share session count should be " + shareSessionsCount);
        TestUtils.waitForCondition(() -> yammerMetricValue(ShareSessionCache.SHARE_PARTITIONS_COUNT).intValue() == sharePartitionsCount,
            "Share partition count should be " + sharePartitionsCount);
        assertEquals(evictionsCount, cache.evictionsMeter().count());
    }
}
