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

import org.apache.kafka.clients.admin.ListOffsetsOptions;
import org.apache.kafka.clients.admin.ListOffsetsResult.ListOffsetsResultInfo;
import org.apache.kafka.clients.admin.OffsetSpec;
import org.apache.kafka.clients.admin.internals.AdminApiHandler.ApiResult;
import org.apache.kafka.common.IsolationLevel;
import org.apache.kafka.common.Node;
import org.apache.kafka.common.TopicPartition;
import org.apache.kafka.common.errors.UnsupportedVersionException;
import org.apache.kafka.common.message.ListOffsetsRequestData.ListOffsetsPartition;
import org.apache.kafka.common.message.ListOffsetsRequestData.ListOffsetsTopic;
import org.apache.kafka.common.message.ListOffsetsResponseData;
import org.apache.kafka.common.message.ListOffsetsResponseData.ListOffsetsPartitionResponse;
import org.apache.kafka.common.message.ListOffsetsResponseData.ListOffsetsTopicResponse;
import org.apache.kafka.common.protocol.Errors;
import org.apache.kafka.common.requests.ListOffsetsRequest;
import org.apache.kafka.common.requests.ListOffsetsResponse;
import org.apache.kafka.common.utils.LogContext;

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.Set;
import java.util.stream.Collectors;

import static java.util.Collections.emptyList;
import static java.util.Collections.emptyMap;
import static java.util.Collections.emptySet;
import static java.util.Collections.singleton;
import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertTrue;

public final class ListOffsetsHandlerTest {

    private final LogContext logContext = new LogContext();
    private final int defaultApiTimeoutMs = 100;

    private final TopicPartition t0p0 = new TopicPartition("t0", 0);
    private final TopicPartition t0p1 = new TopicPartition("t0", 1);
    private final TopicPartition t1p0 = new TopicPartition("t1", 0);
    private final TopicPartition t1p1 = new TopicPartition("t1", 1);
    private final TopicPartition t2p0 = new TopicPartition("t2", 0);
    private final TopicPartition t2p1 = new TopicPartition("t2", 1);

    private final Node node = new Node(1, "host", 1234);

    private final Map<TopicPartition, Long> offsetTimestampsByPartition = new HashMap<>() {
        {
            put(t0p0, ListOffsetsRequest.LATEST_TIMESTAMP);
            put(t0p1, ListOffsetsRequest.EARLIEST_TIMESTAMP);
            put(t1p0, 123L);
            put(t1p1, ListOffsetsRequest.MAX_TIMESTAMP);
            put(t2p0, ListOffsetsRequest.EARLIEST_LOCAL_TIMESTAMP);
            put(t2p1, ListOffsetsRequest.LATEST_TIERED_TIMESTAMP);
        }
    };

    @Test
    public void testBuildRequestSimple() {
        ListOffsetsHandler handler =
            new ListOffsetsHandler(offsetTimestampsByPartition, new ListOffsetsOptions(), logContext, defaultApiTimeoutMs);
        ListOffsetsRequest request = handler.buildBatchedRequest(node.id(), Set.of(t0p0, t0p1)).build();
        List<ListOffsetsTopic> topics = request.topics();
        assertEquals(1, topics.size());
        ListOffsetsTopic topic = topics.get(0);
        assertEquals(2, topic.partitions().size());
        for (ListOffsetsPartition partition : topic.partitions()) {
            TopicPartition topicPartition = new TopicPartition(topic.name(), partition.partitionIndex());
            assertExpectedTimestamp(topicPartition, partition.timestamp());
        }
        assertEquals(IsolationLevel.READ_UNCOMMITTED, request.isolationLevel());
    }

    @Test
    public void testBuildRequestMultipleTopicsWithReadCommitted() {
        ListOffsetsHandler handler =
            new ListOffsetsHandler(
                offsetTimestampsByPartition, new ListOffsetsOptions(IsolationLevel.READ_COMMITTED), logContext, defaultApiTimeoutMs);
        ListOffsetsRequest request =
            handler.buildBatchedRequest(node.id(), offsetTimestampsByPartition.keySet()).build();
        List<ListOffsetsTopic> topics = request.topics();
        assertEquals(3, topics.size());
        Map<TopicPartition, ListOffsetsPartition> partitions = new HashMap<>();
        for (ListOffsetsTopic topic : topics) {
            for (ListOffsetsPartition partition : topic.partitions()) {
                partitions.put(new TopicPartition(topic.name(), partition.partitionIndex()), partition);
            }
        }
        assertEquals(6, partitions.size());
        for (Map.Entry<TopicPartition, ListOffsetsPartition> entry : partitions.entrySet()) {
            assertExpectedTimestamp(entry.getKey(), entry.getValue().timestamp());
        }
        assertEquals(IsolationLevel.READ_COMMITTED, request.isolationLevel());
    }

    @Test
    public void testBuildRequestAllowedVersions() {
        ListOffsetsHandler defaultOptionsHandler =
            new ListOffsetsHandler(offsetTimestampsByPartition, new ListOffsetsOptions(), logContext, defaultApiTimeoutMs);
        ListOffsetsRequest.Builder builder =
            defaultOptionsHandler.buildBatchedRequest(node.id(), Set.of(t0p0, t0p1, t1p0));
        assertEquals(1, builder.oldestAllowedVersion());

        ListOffsetsHandler readCommittedHandler =
            new ListOffsetsHandler(
                offsetTimestampsByPartition, new ListOffsetsOptions(IsolationLevel.READ_COMMITTED), logContext, defaultApiTimeoutMs);
        builder = readCommittedHandler.buildBatchedRequest(node.id(), Set.of(t0p0, t0p1, t1p0));
        assertEquals(2, builder.oldestAllowedVersion());

        builder = readCommittedHandler.buildBatchedRequest(node.id(), Set.of(t0p0, t0p1, t1p0, t1p1));
        assertEquals(7, builder.oldestAllowedVersion());

        builder = readCommittedHandler.buildBatchedRequest(node.id(), Set.of(t0p0, t0p1, t1p0, t1p1, t2p0));
        assertEquals(8, builder.oldestAllowedVersion());

        builder = readCommittedHandler.buildBatchedRequest(node.id(), Set.of(t0p0, t0p1, t1p0, t1p1, t2p0, t2p1));
        assertEquals(9, builder.oldestAllowedVersion());
    }

    @Test
    public void testHandleSuccessfulResponse() {
        ApiResult<TopicPartition, ListOffsetsResultInfo> result =
            handleResponse(createResponse(emptyMap()));

        assertResult(result, offsetTimestampsByPartition.keySet(), emptyMap(), emptyList(), emptySet());
    }

    @Test
    public void testHandleRetriablePartitionTimeoutResponse() {
        TopicPartition errorPartition = t0p0;
        Map<TopicPartition, Short> errorsByPartition = new HashMap<>();
        errorsByPartition.put(errorPartition, Errors.REQUEST_TIMED_OUT.code());

        ApiResult<TopicPartition, ListOffsetsResultInfo> result =
            handleResponse(createResponse(errorsByPartition));

        // Timeouts should be retried within the fulfillment stage as they are a common type of
        // retriable error.
        Set<TopicPartition> retriable = singleton(errorPartition);
        Set<TopicPartition> completed = new HashSet<>(offsetTimestampsByPartition.keySet());
        completed.removeAll(retriable);
        assertResult(result, completed, emptyMap(), emptyList(), retriable);
    }

    @Test
    public void testHandleLookupRetriablePartitionInvalidMetadataResponse() {
        TopicPartition errorPartition = t0p0;
        Errors error = Errors.NOT_LEADER_OR_FOLLOWER;
        Map<TopicPartition, Short> errorsByPartition = new HashMap<>();
        errorsByPartition.put(errorPartition, error.code());

        ApiResult<TopicPartition, ListOffsetsResultInfo> result =
            handleResponse(createResponse(errorsByPartition));

        // Some invalid metadata errors should be retried from the lookup stage as the partition-to-leader
        // mappings should be recalculated.
        List<TopicPartition> unmapped = new ArrayList<>();
        unmapped.add(errorPartition);
        Set<TopicPartition> completed = new HashSet<>(offsetTimestampsByPartition.keySet());
        completed.removeAll(unmapped);
        assertResult(result, completed, emptyMap(), unmapped, emptySet());
    }

    @Test
    public void testHandleUnexpectedPartitionErrorResponse() {
        TopicPartition errorPartition = t0p0;
        Errors error = Errors.UNKNOWN_SERVER_ERROR;
        Map<TopicPartition, Short> errorsByPartition = new HashMap<>();
        errorsByPartition.put(errorPartition, error.code());

        ApiResult<TopicPartition, ListOffsetsResultInfo> result =
            handleResponse(createResponse(errorsByPartition));

        Map<TopicPartition, Throwable> failed = new HashMap<>();
        failed.put(errorPartition, error.exception());
        Set<TopicPartition> completed = new HashSet<>(offsetTimestampsByPartition.keySet());
        completed.removeAll(failed.keySet());
        assertResult(result, completed, failed, emptyList(), emptySet());
    }

    @Test
    public void testHandleResponseSanityCheck() {
        TopicPartition errorPartition = t0p0;
        Map<TopicPartition, Long> specsByPartition = new HashMap<>(offsetTimestampsByPartition);
        specsByPartition.remove(errorPartition);

        ApiResult<TopicPartition, ListOffsetsResultInfo> result =
            handleResponse(createResponse(emptyMap(), specsByPartition));

        assertEquals(offsetTimestampsByPartition.size() - 1, result.completedKeys.size());
        assertEquals(1, result.failedKeys.size());
        assertEquals(errorPartition, result.failedKeys.keySet().iterator().next());
        String sanityCheckMessage = result.failedKeys.get(errorPartition).getMessage();
        assertTrue(sanityCheckMessage.contains("did not contain a result for topic partition"));
        assertTrue(result.unmappedKeys.isEmpty());
    }

    @Test
    public void testHandleResponseUnsupportedVersion() {
        int brokerId = 1;
        UnsupportedVersionException uve = new UnsupportedVersionException("");
        Map<TopicPartition, OffsetSpec> maxTimestampPartitions = new HashMap<>();
        maxTimestampPartitions.put(t1p1, OffsetSpec.maxTimestamp());

        ListOffsetsHandler handler =
            new ListOffsetsHandler(offsetTimestampsByPartition, new ListOffsetsOptions(), logContext, defaultApiTimeoutMs);

        final Map<TopicPartition, Long> nonMaxTimestampPartitions = new HashMap<>(offsetTimestampsByPartition);
        maxTimestampPartitions.forEach((k, v) -> nonMaxTimestampPartitions.remove(k));
        // Unsupported version exceptions currently cannot be handled if there's no partition with a
        // MAX_TIMESTAMP spec...
        Set<TopicPartition> keysToTest = nonMaxTimestampPartitions.keySet();
        Set<TopicPartition> expectedFailures = keysToTest;
        assertEquals(
            mapToError(expectedFailures, uve),
            handler.handleUnsupportedVersionException(brokerId, uve, keysToTest));

        // ...or if there are only partitions with MAX_TIMESTAMP specs.
        keysToTest = maxTimestampPartitions.keySet();
        expectedFailures = keysToTest;
        assertEquals(
            mapToError(expectedFailures, uve),
            handler.handleUnsupportedVersionException(brokerId, uve, keysToTest));

        // What can be handled is a request with a mix of partitions with MAX_TIMESTAMP specs
        // and partitions with non-MAX_TIMESTAMP specs.
        keysToTest = offsetTimestampsByPartition.keySet();
        expectedFailures = maxTimestampPartitions.keySet();
        assertEquals(
            mapToError(expectedFailures, uve),
            handler.handleUnsupportedVersionException(brokerId, uve, keysToTest));
    }

    @Test
    public void testBuildRequestWithDefaultApiTimeoutMs() {
        ListOffsetsOptions options = new ListOffsetsOptions();
        ListOffsetsHandler handler =
                new ListOffsetsHandler(offsetTimestampsByPartition, options, logContext, defaultApiTimeoutMs);
        ListOffsetsRequest request = handler.buildBatchedRequest(node.id(), Set.of(t0p0, t0p1)).build();
        assertEquals(defaultApiTimeoutMs, request.timeoutMs());
    }

    @Test
    public void testBuildRequestWithTimeoutMs() {
        Integer timeoutMs = 200;
        ListOffsetsOptions options = new ListOffsetsOptions().timeoutMs(timeoutMs);
        ListOffsetsHandler handler =
                new ListOffsetsHandler(offsetTimestampsByPartition, options, logContext, defaultApiTimeoutMs);
        ListOffsetsRequest request = handler.buildBatchedRequest(node.id(), Set.of(t0p0, t0p1)).build();
        assertEquals(timeoutMs, request.timeoutMs());
    }

    private static Map<TopicPartition, Throwable> mapToError(Set<TopicPartition> keys, Throwable t) {
        return keys.stream().collect(Collectors.toMap(k -> k, k -> t));
    }

    private void assertExpectedTimestamp(TopicPartition topicPartition, long actualTimestamp) {
        Long expectedTimestamp = offsetTimestampsByPartition.get(topicPartition);
        assertEquals(expectedTimestamp, actualTimestamp);
    }

    private ListOffsetsResponse createResponse(Map<TopicPartition, Short> errorsByPartition) {
        return createResponse(errorsByPartition, offsetTimestampsByPartition);
    }

    private static ListOffsetsResponse createResponse(
        Map<TopicPartition, Short> errorsByPartition,
        Map<TopicPartition, Long> specsByPartition
    ) {
        Map<String, ListOffsetsTopicResponse> responsesByTopic = new HashMap<>();
        for (Map.Entry<TopicPartition, Long> offsetSpecEntry : specsByPartition.entrySet()) {
            TopicPartition topicPartition = offsetSpecEntry.getKey();
            ListOffsetsTopicResponse topicResponse = responsesByTopic.computeIfAbsent(
                topicPartition.topic(), t -> new ListOffsetsTopicResponse());
            topicResponse.setName(topicPartition.topic());
            ListOffsetsPartitionResponse partitionResponse = new ListOffsetsPartitionResponse();
            partitionResponse.setPartitionIndex(topicPartition.partition());
            partitionResponse.setOffset(getOffset(topicPartition, offsetSpecEntry.getValue()));
            partitionResponse.setErrorCode(errorsByPartition.getOrDefault(topicPartition, (short) 0));
            topicResponse.partitions().add(partitionResponse);
        }
        ListOffsetsResponseData responseData = new ListOffsetsResponseData();
        responseData.setTopics(new ArrayList<>(responsesByTopic.values()));
        return new ListOffsetsResponse(responseData);
    }

    private ApiResult<TopicPartition, ListOffsetsResultInfo> handleResponse(ListOffsetsResponse response) {
        ListOffsetsHandler handler =
            new ListOffsetsHandler(offsetTimestampsByPartition, new ListOffsetsOptions(), logContext, defaultApiTimeoutMs);
        return handler.handleResponse(node, offsetTimestampsByPartition.keySet(), response);
    }

    private void assertResult(
        ApiResult<TopicPartition, ListOffsetsResultInfo> result,
        Set<TopicPartition> expectedCompleted,
        Map<TopicPartition, Throwable> expectedFailed,
        List<TopicPartition> expectedUnmapped,
        Set<TopicPartition> expectedRetriable
    ) {
        assertEquals(expectedCompleted, result.completedKeys.keySet());
        assertEquals(expectedFailed, result.failedKeys);
        assertEquals(expectedUnmapped, result.unmappedKeys);
        Set<TopicPartition> actualRetriable = new HashSet<>(offsetTimestampsByPartition.keySet());
        actualRetriable.removeAll(result.completedKeys.keySet());
        actualRetriable.removeAll(result.failedKeys.keySet());
        actualRetriable.removeAll(new HashSet<>(result.unmappedKeys));
        assertEquals(expectedRetriable, actualRetriable);
    }

    private static long getOffset(TopicPartition topicPartition, Long offsetQuery) {
        long base = 1 << 10;
        if (offsetQuery == ListOffsetsRequest.EARLIEST_TIMESTAMP) {
            return topicPartition.hashCode() & (base - 1);
        } else if (offsetQuery >= 0L) {
            return base;
        } else if (offsetQuery == ListOffsetsRequest.LATEST_TIMESTAMP) {
            return base + 1 + (topicPartition.hashCode() & (base - 1));
        }
        return 2 * base + 1;
    }
}
