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

import org.apache.kafka.common.ClusterResource;
import org.apache.kafka.common.Endpoint;
import org.apache.kafka.common.security.auth.SecurityProtocol;

import org.junit.jupiter.api.Test;

import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.Set;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.CompletionException;

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.assertThrows;
import static org.junit.jupiter.api.Assertions.assertTrue;

public final class EndpointReadyFuturesTest {
    private static final Endpoint EXTERNAL =
            new Endpoint("EXTERNAL", SecurityProtocol.SSL, "127.0.0.1", 9092);

    private static final Endpoint INTERNAL =
            new Endpoint("INTERNAL", SecurityProtocol.PLAINTEXT, "127.0.0.1", 9093);

    private static final KafkaAuthorizerServerInfo INFO = new KafkaAuthorizerServerInfo(
        new ClusterResource("S6-01LPiQOCBhhFIunQUcQ"),
        1,
        List.of(EXTERNAL, INTERNAL),
        INTERNAL,
        List.of("INTERNAL"));

    static void assertComplete(
            EndpointReadyFutures readyFutures,
            Endpoint... endpoints
    ) {
        for (Endpoint endpoint : endpoints) {
            String name = endpoint.listener();
            CompletableFuture<Void> future = readyFutures.futures().get(endpoint);
            assertNotNull(future, "Unable to find future for " + name);
            assertTrue(future.isDone(), "Future for " + name + " is not done.");
            assertFalse(future.isCompletedExceptionally(),
                    "Future for " + name + " is completed exceptionally.");
        }
    }

    static void assertIncomplete(
            EndpointReadyFutures readyFutures,
            Endpoint... endpoints
    ) {
        for (Endpoint endpoint : endpoints) {
            CompletableFuture<Void> future = readyFutures.futures().get(endpoint);
            assertNotNull(future, "Unable to find future for " + endpoint);
            assertFalse(future.isDone(), "Future for " + endpoint + " is done.");
        }
    }

    static void assertException(
            EndpointReadyFutures readyFutures,
            Throwable throwable,
            Endpoint... endpoints
    ) {
        for (Endpoint endpoint : endpoints) {
            CompletableFuture<Void> future = readyFutures.futures().get(endpoint);
            assertNotNull(future, "Unable to find future for " + endpoint);
            assertTrue(future.isCompletedExceptionally(),
                    "Future for " + endpoint + " is not completed exceptionally.");
            Throwable cause = assertThrows(CompletionException.class,
                () -> future.getNow(null)).getCause();
            assertNotNull(cause, "Unable to find CompletionException cause for " + endpoint);
            assertEquals(throwable.getClass(), cause.getClass());
            assertEquals(throwable.getMessage(), cause.getMessage());
        }
    }

    @Test
    public void testImmediateCompletion() {
        EndpointReadyFutures readyFutures = new EndpointReadyFutures.Builder().
                build(Optional.empty(), INFO);
        assertEquals(Set.of(EXTERNAL, INTERNAL),
                readyFutures.futures().keySet());
        assertComplete(readyFutures, EXTERNAL, INTERNAL);
    }

    @Test
    public void testAddReadinessFuture() {
        CompletableFuture<Void> foo = new CompletableFuture<>();
        EndpointReadyFutures readyFutures = new EndpointReadyFutures.Builder().
                addReadinessFuture("foo", foo).
                build(Optional.empty(), INFO);
        assertEquals(Set.of(EXTERNAL, INTERNAL),
                readyFutures.futures().keySet());
        assertIncomplete(readyFutures, EXTERNAL, INTERNAL);
        foo.complete(null);
        assertComplete(readyFutures, EXTERNAL, INTERNAL);
    }

    @Test
    public void testAddMultipleReadinessFutures() {
        CompletableFuture<Void> foo = new CompletableFuture<>();
        CompletableFuture<Void> bar = new CompletableFuture<>();
        EndpointReadyFutures readyFutures = new EndpointReadyFutures.Builder().
                addReadinessFuture("foo", foo).
                addReadinessFuture("bar", bar).
                build(Optional.empty(), INFO);
        assertEquals(Set.of(EXTERNAL, INTERNAL),
                readyFutures.futures().keySet());
        assertIncomplete(readyFutures, EXTERNAL, INTERNAL);
        foo.complete(null);
        assertIncomplete(readyFutures, EXTERNAL, INTERNAL);
        bar.complete(null);
        assertComplete(readyFutures, EXTERNAL, INTERNAL);
    }

    @Test
    public void testAddReadinessFutures() {
        Map<Endpoint, CompletableFuture<Void>> bazFutures = new HashMap<>();
        bazFutures.put(EXTERNAL, new CompletableFuture<>());
        bazFutures.put(INTERNAL, new CompletableFuture<>());
        EndpointReadyFutures readyFutures = new EndpointReadyFutures.Builder().
                addReadinessFutures("baz", bazFutures).
                build(Optional.empty(), INFO);
        assertEquals(Set.of(EXTERNAL, INTERNAL),
                readyFutures.futures().keySet());
        assertIncomplete(readyFutures, EXTERNAL, INTERNAL);
        bazFutures.get(EXTERNAL).complete(null);
        assertComplete(readyFutures, EXTERNAL);
        assertIncomplete(readyFutures, INTERNAL);
        bazFutures.get(INTERNAL).complete(null);
        assertComplete(readyFutures, EXTERNAL, INTERNAL);
    }

    @Test
    public void testFailedReadinessFuture() {
        CompletableFuture<Void> foo = new CompletableFuture<>();
        CompletableFuture<Void> bar = new CompletableFuture<>();
        EndpointReadyFutures readyFutures = new EndpointReadyFutures.Builder().
                addReadinessFuture("foo", foo).
                addReadinessFuture("bar", bar).
                build(Optional.empty(), INFO);
        assertEquals(Set.of(EXTERNAL, INTERNAL),
                readyFutures.futures().keySet());
        assertIncomplete(readyFutures, EXTERNAL, INTERNAL);
        foo.complete(null);
        assertIncomplete(readyFutures, EXTERNAL, INTERNAL);
        bar.completeExceptionally(new RuntimeException("Failed."));
        assertException(readyFutures, new RuntimeException("Failed."),
                EXTERNAL, INTERNAL);
    }
}
