/*
 * 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.apache.kafka.clients.MockClient;
import org.apache.kafka.common.Node;
import org.apache.kafka.common.PartitionInfo;
import org.apache.kafka.common.TopicPartition;
import org.apache.kafka.common.Uuid;
import org.apache.kafka.common.errors.InvalidTopicException;
import org.apache.kafka.common.errors.TimeoutException;
import org.apache.kafka.common.errors.TopicAuthorizationException;
import org.apache.kafka.common.internals.ClusterResourceListeners;
import org.apache.kafka.common.metrics.MetricConfig;
import org.apache.kafka.common.metrics.Metrics;
import org.apache.kafka.common.protocol.Errors;
import org.apache.kafka.common.requests.MetadataResponse;
import org.apache.kafka.common.requests.RequestTestUtils;
import org.apache.kafka.common.utils.LogContext;
import org.apache.kafka.common.utils.MockTime;

import org.junit.jupiter.api.AfterEach;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;

import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.Set;

import static java.util.Collections.singleton;
import static java.util.Collections.singletonMap;
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;
import static org.junit.jupiter.api.Assertions.assertThrows;
import static org.junit.jupiter.api.Assertions.fail;

public class TopicMetadataFetcherTest {

    private final String topicName = "test";
    private final Uuid topicId = Uuid.randomUuid();
    private final Map<String, Uuid> topicIds = new HashMap<>() {
        {
            put(topicName, topicId);
        }
    };
    private final TopicPartition tp0 = new TopicPartition(topicName, 0);
    private final int validLeaderEpoch = 0;
    private final MetadataResponse initialUpdateResponse =
        RequestTestUtils.metadataUpdateWithIds(1, singletonMap(topicName, 4), topicIds);

    private MockTime time = new MockTime(1);
    private SubscriptionState subscriptions;
    private ConsumerMetadata metadata;
    private MockClient client;
    private Metrics metrics;
    private ConsumerNetworkClient consumerClient;
    private TopicMetadataFetcher topicMetadataFetcher;

    @BeforeEach
    public void setup() {
    }

    private void assignFromUser(Set<TopicPartition> partitions) {
        subscriptions.assignFromUser(partitions);
        client.updateMetadata(initialUpdateResponse);

        // A dummy metadata update to ensure valid leader epoch.
        metadata.updateWithCurrentRequestVersion(RequestTestUtils.metadataUpdateWithIds("dummy", 1,
            Collections.emptyMap(), singletonMap(topicName, 4),
            tp -> validLeaderEpoch, topicIds), false, 0L);
    }

    @AfterEach
    public void teardown() throws Exception {
        if (metrics != null)
            this.metrics.close();
    }

    @Test
    public void testGetAllTopics() {
        // sending response before request, as getTopicMetadata is a blocking call
        buildFetcher();
        assignFromUser(singleton(tp0));
        client.prepareResponse(newMetadataResponse(Errors.NONE));

        Map<String, List<PartitionInfo>> allTopics = topicMetadataFetcher.getAllTopicMetadata(time.timer(5000L));

        assertEquals(initialUpdateResponse.topicMetadata().size(), allTopics.size());
    }

    @Test
    public void testGetAllTopicsDisconnect() {
        // first try gets a disconnect, next succeeds
        buildFetcher();
        assignFromUser(singleton(tp0));
        client.prepareResponse(null, true);
        client.prepareResponse(newMetadataResponse(Errors.NONE));
        Map<String, List<PartitionInfo>> allTopics = topicMetadataFetcher.getAllTopicMetadata(time.timer(5000L));
        assertEquals(initialUpdateResponse.topicMetadata().size(), allTopics.size());
    }

    @Test
    public void testGetAllTopicsTimeout() {
        // since no response is prepared, the request should time out
        buildFetcher();
        assignFromUser(singleton(tp0));
        assertThrows(TimeoutException.class, () -> topicMetadataFetcher.getAllTopicMetadata(time.timer(50L)));
    }

    @Test
    public void testGetAllTopicsUnauthorized() {
        buildFetcher();
        assignFromUser(singleton(tp0));
        client.prepareResponse(newMetadataResponse(Errors.TOPIC_AUTHORIZATION_FAILED));
        try {
            topicMetadataFetcher.getAllTopicMetadata(time.timer(10L));
            fail();
        } catch (TopicAuthorizationException e) {
            assertEquals(singleton(topicName), e.unauthorizedTopics());
        }
    }

    @Test
    public void testGetTopicMetadataInvalidTopic() {
        buildFetcher();
        assignFromUser(singleton(tp0));
        client.prepareResponse(newMetadataResponse(Errors.INVALID_TOPIC_EXCEPTION));
        assertThrows(InvalidTopicException.class, () -> topicMetadataFetcher.getTopicMetadata(topicName, true, time.timer(5000L)));
    }

    @Test
    public void testGetTopicMetadataUnknownTopic() {
        buildFetcher();
        assignFromUser(singleton(tp0));
        client.prepareResponse(newMetadataResponse(Errors.UNKNOWN_TOPIC_OR_PARTITION));

        List<PartitionInfo> topicMetadata = topicMetadataFetcher.getTopicMetadata(topicName, true, time.timer(5000L));
        assertNull(topicMetadata);
    }

    @Test
    public void testGetTopicMetadataLeaderNotAvailable() {
        buildFetcher();
        assignFromUser(singleton(tp0));
        client.prepareResponse(newMetadataResponse(Errors.LEADER_NOT_AVAILABLE));
        client.prepareResponse(newMetadataResponse(Errors.NONE));

        List<PartitionInfo> topicMetadata = topicMetadataFetcher.getTopicMetadata(topicName, true, time.timer(5000L));
        assertNotNull(topicMetadata);
    }

    @Test
    public void testGetTopicMetadataOfflinePartitions() {
        buildFetcher();
        assignFromUser(singleton(tp0));
        MetadataResponse originalResponse = newMetadataResponse(Errors.NONE); //baseline ok response

        //create a response based on the above one with all partitions being leaderless
        List<MetadataResponse.TopicMetadata> altTopics = new ArrayList<>();
        for (MetadataResponse.TopicMetadata item : originalResponse.topicMetadata()) {
            List<MetadataResponse.PartitionMetadata> partitions = item.partitionMetadata();
            List<MetadataResponse.PartitionMetadata> altPartitions = new ArrayList<>();
            for (MetadataResponse.PartitionMetadata p : partitions) {
                altPartitions.add(new MetadataResponse.PartitionMetadata(
                    p.error,
                    p.topicPartition,
                    Optional.empty(), //no leader
                    Optional.empty(),
                    p.replicaIds,
                    p.inSyncReplicaIds,
                    p.offlineReplicaIds
                ));
            }
            MetadataResponse.TopicMetadata alteredTopic = new MetadataResponse.TopicMetadata(
                item.error(),
                item.topic(),
                item.isInternal(),
                altPartitions
            );
            altTopics.add(alteredTopic);
        }
        Node controller = originalResponse.controller();
        MetadataResponse altered = RequestTestUtils.metadataResponse(
            originalResponse.brokers(),
            originalResponse.clusterId(),
            controller != null ? controller.id() : MetadataResponse.NO_CONTROLLER_ID,
            altTopics);

        client.prepareResponse(altered);

        List<PartitionInfo> topicMetadata = topicMetadataFetcher.getTopicMetadata(topicName, false, time.timer(5000L));

        assertNotNull(topicMetadata);
        assertFalse(topicMetadata.isEmpty());
        //noinspection ConstantConditions
        assertEquals(metadata.fetch().partitionCountForTopic(topicName).longValue(), topicMetadata.size());
    }

    private MetadataResponse newMetadataResponse(Errors error) {
        List<MetadataResponse.PartitionMetadata> partitionsMetadata = new ArrayList<>();
        if (error == Errors.NONE) {
            Optional<MetadataResponse.TopicMetadata> foundMetadata = initialUpdateResponse.topicMetadata()
                    .stream()
                    .filter(topicMetadata -> topicMetadata.topic().equals(topicName))
                    .findFirst();
            foundMetadata.ifPresent(topicMetadata -> partitionsMetadata.addAll(topicMetadata.partitionMetadata()));
        }

        MetadataResponse.TopicMetadata topicMetadata = new MetadataResponse.TopicMetadata(error, topicName, false,
                partitionsMetadata);
        List<Node> brokers = new ArrayList<>(initialUpdateResponse.brokers());
        return RequestTestUtils.metadataResponse(brokers, initialUpdateResponse.clusterId(),
                initialUpdateResponse.controller().id(), Collections.singletonList(topicMetadata));
    }

    private void buildFetcher() {
        MetricConfig metricConfig = new MetricConfig();
        long metadataExpireMs = Long.MAX_VALUE;
        long retryBackoffMs = 100;
        long retryBackoffMaxMs = 1000;
        LogContext logContext = new LogContext();
        SubscriptionState subscriptionState = new SubscriptionState(logContext, AutoOffsetResetStrategy.EARLIEST);
        buildDependencies(metricConfig, metadataExpireMs, subscriptionState, logContext);
        topicMetadataFetcher = new TopicMetadataFetcher(logContext, consumerClient, retryBackoffMs, retryBackoffMaxMs);
    }

    private void buildDependencies(MetricConfig metricConfig,
                                   long metadataExpireMs,
                                   SubscriptionState subscriptionState,
                                   LogContext logContext) {
        time = new MockTime(1);
        subscriptions = subscriptionState;
        metadata = new ConsumerMetadata(0, 0, metadataExpireMs, false, false,
                subscriptions, logContext, new ClusterResourceListeners());
        client = new MockClient(time, metadata);
        metrics = new Metrics(metricConfig, time);
        consumerClient = new ConsumerNetworkClient(logContext, client, metadata, time,
                100, 1000, Integer.MAX_VALUE);
    }
}
