/*
 * 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.Metadata;
import org.apache.kafka.clients.MockClient;
import org.apache.kafka.common.Node;
import org.apache.kafka.common.TopicPartition;
import org.apache.kafka.common.errors.TopicAuthorizationException;
import org.apache.kafka.common.internals.ClusterResourceListeners;
import org.apache.kafka.common.message.OffsetForLeaderEpochResponseData;
import org.apache.kafka.common.message.OffsetForLeaderEpochResponseData.EpochEndOffset;
import org.apache.kafka.common.message.OffsetForLeaderEpochResponseData.OffsetForLeaderTopicResult;
import org.apache.kafka.common.protocol.Errors;
import org.apache.kafka.common.requests.OffsetsForLeaderEpochResponse;
import org.apache.kafka.common.utils.LogContext;
import org.apache.kafka.common.utils.MockTime;
import org.apache.kafka.common.utils.Time;

import org.junit.jupiter.api.Test;

import java.util.Collections;
import java.util.HashMap;
import java.util.Map;
import java.util.Optional;

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

public class OffsetForLeaderEpochClientTest {

    private final TopicPartition tp0 = new TopicPartition("topic", 0);
    private ConsumerNetworkClient consumerClient;
    private MockClient client;

    @Test
    public void testEmptyResponse() {
        OffsetsForLeaderEpochClient offsetClient = newOffsetClient();
        RequestFuture<OffsetsForLeaderEpochUtils.OffsetForEpochResult> future =
                offsetClient.sendAsyncRequest(Node.noNode(), Collections.emptyMap());

        OffsetsForLeaderEpochResponse resp = new OffsetsForLeaderEpochResponse(
            new OffsetForLeaderEpochResponseData());
        client.prepareResponse(resp);
        consumerClient.pollNoWakeup();

        OffsetsForLeaderEpochUtils.OffsetForEpochResult result = future.value();
        assertTrue(result.partitionsToRetry().isEmpty());
        assertTrue(result.endOffsets().isEmpty());
    }

    @Test
    public void testUnexpectedEmptyResponse() {
        Map<TopicPartition, SubscriptionState.FetchPosition> positionMap = new HashMap<>();
        positionMap.put(tp0, new SubscriptionState.FetchPosition(0, Optional.of(1),
                new Metadata.LeaderAndEpoch(Optional.empty(), Optional.of(1))));

        OffsetsForLeaderEpochClient offsetClient = newOffsetClient();
        RequestFuture<OffsetsForLeaderEpochUtils.OffsetForEpochResult> future =
                offsetClient.sendAsyncRequest(Node.noNode(), positionMap);

        OffsetsForLeaderEpochResponse resp = new OffsetsForLeaderEpochResponse(
            new OffsetForLeaderEpochResponseData());
        client.prepareResponse(resp);
        consumerClient.pollNoWakeup();

        OffsetsForLeaderEpochUtils.OffsetForEpochResult result = future.value();
        assertFalse(result.partitionsToRetry().isEmpty());
        assertTrue(result.endOffsets().isEmpty());
    }

    @Test
    public void testOkResponse() {
        Map<TopicPartition, SubscriptionState.FetchPosition> positionMap = new HashMap<>();
        positionMap.put(tp0, new SubscriptionState.FetchPosition(0, Optional.of(1),
                new Metadata.LeaderAndEpoch(Optional.empty(), Optional.of(1))));

        OffsetsForLeaderEpochClient offsetClient = newOffsetClient();
        RequestFuture<OffsetsForLeaderEpochUtils.OffsetForEpochResult> future =
                offsetClient.sendAsyncRequest(Node.noNode(), positionMap);

        client.prepareResponse(prepareOffsetForLeaderEpochResponse(
            tp0, Errors.NONE, 1, 10L));
        consumerClient.pollNoWakeup();

        OffsetsForLeaderEpochUtils.OffsetForEpochResult result = future.value();
        assertTrue(result.partitionsToRetry().isEmpty());
        assertTrue(result.endOffsets().containsKey(tp0));
        assertEquals(result.endOffsets().get(tp0).errorCode(), Errors.NONE.code());
        assertEquals(1, result.endOffsets().get(tp0).leaderEpoch());
        assertEquals(10L, result.endOffsets().get(tp0).endOffset());
    }

    @Test
    public void testUnauthorizedTopic() {
        Map<TopicPartition, SubscriptionState.FetchPosition> positionMap = new HashMap<>();
        positionMap.put(tp0, new SubscriptionState.FetchPosition(0, Optional.of(1),
                new Metadata.LeaderAndEpoch(Optional.empty(), Optional.of(1))));

        OffsetsForLeaderEpochClient offsetClient = newOffsetClient();
        RequestFuture<OffsetsForLeaderEpochUtils.OffsetForEpochResult> future =
                offsetClient.sendAsyncRequest(Node.noNode(), positionMap);

        client.prepareResponse(prepareOffsetForLeaderEpochResponse(
            tp0, Errors.TOPIC_AUTHORIZATION_FAILED, -1, -1));
        consumerClient.pollNoWakeup();

        assertTrue(future.failed());
        assertEquals(TopicAuthorizationException.class, future.exception().getClass());
        assertTrue(((TopicAuthorizationException) future.exception()).unauthorizedTopics().contains(tp0.topic()));
    }

    @Test
    public void testRetriableError() {
        Map<TopicPartition, SubscriptionState.FetchPosition> positionMap = new HashMap<>();
        positionMap.put(tp0, new SubscriptionState.FetchPosition(0, Optional.of(1),
                new Metadata.LeaderAndEpoch(Optional.empty(), Optional.of(1))));

        OffsetsForLeaderEpochClient offsetClient = newOffsetClient();
        RequestFuture<OffsetsForLeaderEpochUtils.OffsetForEpochResult> future =
                offsetClient.sendAsyncRequest(Node.noNode(), positionMap);

        client.prepareResponse(prepareOffsetForLeaderEpochResponse(
            tp0, Errors.LEADER_NOT_AVAILABLE, -1, -1));
        consumerClient.pollNoWakeup();

        assertFalse(future.failed());
        OffsetsForLeaderEpochUtils.OffsetForEpochResult result = future.value();
        assertTrue(result.partitionsToRetry().contains(tp0));
        assertFalse(result.endOffsets().containsKey(tp0));
    }

    private OffsetsForLeaderEpochClient newOffsetClient() {
        buildDependencies(AutoOffsetResetStrategy.EARLIEST);
        return new OffsetsForLeaderEpochClient(consumerClient, new LogContext());
    }

    private void buildDependencies(AutoOffsetResetStrategy offsetResetStrategy) {
        LogContext logContext = new LogContext();
        Time time = new MockTime(1);
        SubscriptionState subscriptions = new SubscriptionState(logContext, offsetResetStrategy);
        Metadata metadata = new ConsumerMetadata(0, 0, Long.MAX_VALUE, false, false,
                subscriptions, logContext, new ClusterResourceListeners());
        client = new MockClient(time, metadata);
        consumerClient = new ConsumerNetworkClient(logContext, client, metadata, time,
                100, 1000, Integer.MAX_VALUE);
    }

    private static OffsetsForLeaderEpochResponse prepareOffsetForLeaderEpochResponse(
            TopicPartition tp, Errors error, int leaderEpoch, long endOffset) {
        OffsetForLeaderEpochResponseData data = new OffsetForLeaderEpochResponseData();
        OffsetForLeaderTopicResult topic = new OffsetForLeaderTopicResult()
            .setTopic(tp.topic());
        data.topics().add(topic);
        topic.partitions().add(new EpochEndOffset()
            .setPartition(tp.partition())
            .setErrorCode(error.code())
            .setLeaderEpoch(leaderEpoch)
            .setEndOffset(endOffset));
        return new OffsetsForLeaderEpochResponse(data);
    }
}
