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

import org.apache.kafka.common.Node;
import org.apache.kafka.common.utils.MockTime;

import org.junit.jupiter.api.Test;

import java.util.List;
import java.util.Optional;
import java.util.Random;
import java.util.stream.Collectors;
import java.util.stream.IntStream;

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

public class RequestManagerTest {
    private final MockTime time = new MockTime();
    private final int requestTimeoutMs = 30000;
    private final int retryBackoffMs = 100;
    private final Random random = new Random(1);

    @Test
    public void testResetAllConnections() {
        Node node1 = new Node(1, "mock-host-1", 4321);
        Node node2 = new Node(2, "mock-host-2", 4321);

        RequestManager cache = new RequestManager(
            makeBootstrapList(3),
            retryBackoffMs,
            requestTimeoutMs,
            random
        );

        // One host has an inflight request
        cache.onRequestSent(node1, 1, time.milliseconds());
        assertFalse(cache.isReady(node1, time.milliseconds()));

        // Another is backing off
        cache.onRequestSent(node2, 2, time.milliseconds());
        cache.onResponseResult(node2, 2, false, time.milliseconds());
        assertFalse(cache.isReady(node2, time.milliseconds()));

        cache.resetAll();

        // Now both should be ready
        assertTrue(cache.isReady(node1, time.milliseconds()));
        assertTrue(cache.isReady(node2, time.milliseconds()));
    }

    @Test
    public void testBackoffAfterFailure() {
        Node node = new Node(1, "mock-host-1", 4321);

        RequestManager cache = new RequestManager(
            makeBootstrapList(3),
            retryBackoffMs,
            requestTimeoutMs,
            random
        );

        assertTrue(cache.isReady(node, time.milliseconds()));

        long correlationId = 1;
        cache.onRequestSent(node, correlationId, time.milliseconds());
        assertFalse(cache.isReady(node, time.milliseconds()));

        cache.onResponseResult(node, correlationId, false, time.milliseconds());
        assertFalse(cache.isReady(node, time.milliseconds()));

        time.sleep(retryBackoffMs);
        assertTrue(cache.isReady(node, time.milliseconds()));
    }

    @Test
    public void testSuccessfulResponse() {
        Node node = new Node(1, "mock-host-1", 4321);

        RequestManager cache = new RequestManager(
            makeBootstrapList(3),
            retryBackoffMs,
            requestTimeoutMs,
            random
        );

        long correlationId = 1;
        cache.onRequestSent(node, correlationId, time.milliseconds());
        assertFalse(cache.isReady(node, time.milliseconds()));
        cache.onResponseResult(node, correlationId, true, time.milliseconds());
        assertTrue(cache.isReady(node, time.milliseconds()));
    }

    @Test
    public void testIgnoreUnexpectedResponse() {
        Node node = new Node(1, "mock-host-1", 4321);

        RequestManager cache = new RequestManager(
            makeBootstrapList(3),
            retryBackoffMs,
            requestTimeoutMs,
            random
        );

        long correlationId = 1;
        cache.onRequestSent(node, correlationId, time.milliseconds());
        assertFalse(cache.isReady(node, time.milliseconds()));
        cache.onResponseResult(node, correlationId + 1, true, time.milliseconds());
        assertFalse(cache.isReady(node, time.milliseconds()));
    }

    @Test
    public void testRequestTimeout() {
        Node node = new Node(1, "mock-host-1", 4321);

        RequestManager cache = new RequestManager(
            makeBootstrapList(3),
            retryBackoffMs,
            requestTimeoutMs,
            random
        );

        long correlationId = 1;
        cache.onRequestSent(node, correlationId, time.milliseconds());
        assertFalse(cache.isReady(node, time.milliseconds()));

        time.sleep(requestTimeoutMs - 1);
        assertFalse(cache.isReady(node, time.milliseconds()));

        time.sleep(1);
        assertTrue(cache.isReady(node, time.milliseconds()));
    }

    @Test
    public void testRequestToBootstrapList() {
        List<Node> bootstrapList = makeBootstrapList(2);
        RequestManager cache = new RequestManager(
            bootstrapList,
            retryBackoffMs,
            requestTimeoutMs,
            random
        );

        // Find a ready node with the starting state
        Node bootstrapNode1 = cache.findReadyBootstrapServer(time.milliseconds()).get();
        assertTrue(
            bootstrapList.contains(bootstrapNode1),
            String.format("%s is not in %s", bootstrapNode1, bootstrapList)
        );
        assertEquals(0, cache.backoffBeforeAvailableBootstrapServer(time.milliseconds()));

        // Send a request and check the cache state
        cache.onRequestSent(bootstrapNode1, 1, time.milliseconds());
        assertEquals(
            Optional.empty(),
            cache.findReadyBootstrapServer(time.milliseconds())
        );
        assertEquals(requestTimeoutMs, cache.backoffBeforeAvailableBootstrapServer(time.milliseconds()));

        // Fail the request
        time.sleep(100);
        cache.onResponseResult(bootstrapNode1, 1, false, time.milliseconds());
        Node bootstrapNode2 = cache.findReadyBootstrapServer(time.milliseconds()).get();
        assertNotEquals(bootstrapNode1, bootstrapNode2);
        assertEquals(0, cache.backoffBeforeAvailableBootstrapServer(time.milliseconds()));

        // Send a request to the second node and check the state
        cache.onRequestSent(bootstrapNode2, 2, time.milliseconds());
        assertEquals(
            Optional.empty(),
            cache.findReadyBootstrapServer(time.milliseconds())
        );
        assertEquals(requestTimeoutMs, cache.backoffBeforeAvailableBootstrapServer(time.milliseconds()));


        // Fail the second request before the request timeout
        time.sleep(retryBackoffMs - 1);
        cache.onResponseResult(bootstrapNode2, 2, false, time.milliseconds());
        assertEquals(
            Optional.empty(),
            cache.findReadyBootstrapServer(time.milliseconds())
        );
        assertEquals(1, cache.backoffBeforeAvailableBootstrapServer(time.milliseconds()));

        // Timeout the first backoff and show that that node is ready
        time.sleep(1);
        Node bootstrapNode3 = cache.findReadyBootstrapServer(time.milliseconds()).get();
        assertEquals(bootstrapNode1, bootstrapNode3);
        assertEquals(0, cache.backoffBeforeAvailableBootstrapServer(time.milliseconds()));
    }

    @Test
    public void testFindReadyWithInflightRequest() {
        Node otherNode = new Node(1, "other-node", 1234);
        List<Node> bootstrapList = makeBootstrapList(3);
        RequestManager cache = new RequestManager(
            bootstrapList,
            retryBackoffMs,
            requestTimeoutMs,
            random
        );

        // Send request to a node that is not in the bootstrap list
        cache.onRequestSent(otherNode, 1, time.milliseconds());
        assertEquals(Optional.empty(), cache.findReadyBootstrapServer(time.milliseconds()));
    }

    @Test
    public void testFindReadyWithRequestTimedOut() {
        Node otherNode = new Node(1, "other-node", 1234);
        List<Node> bootstrapList = makeBootstrapList(3);
        RequestManager cache = new RequestManager(
            bootstrapList,
            retryBackoffMs,
            requestTimeoutMs,
            random
        );

        // Send request to a node that is not in the bootstrap list
        cache.onRequestSent(otherNode, 1, time.milliseconds());
        assertTrue(cache.isResponseExpected(otherNode, 1));
        assertEquals(Optional.empty(), cache.findReadyBootstrapServer(time.milliseconds()));

        // Timeout the request
        time.sleep(requestTimeoutMs);
        Node bootstrapNode = cache.findReadyBootstrapServer(time.milliseconds()).get();
        assertTrue(bootstrapList.contains(bootstrapNode));
        assertFalse(cache.isResponseExpected(otherNode, 1));
    }

    @Test
    public void testHasRequestTimedOut() {
        List<Node> bootstrapList = makeBootstrapList(2);
        RequestManager cache = new RequestManager(
            bootstrapList,
            retryBackoffMs,
            requestTimeoutMs,
            random
        );

        // Find a ready node with the starting state
        Node bootstrapNode1 = cache.findReadyBootstrapServer(time.milliseconds()).get();
        assertTrue(
            bootstrapList.contains(bootstrapNode1),
            String.format("%s is not in %s", bootstrapNode1, bootstrapList)
        );
        // Before sending a request, no request should have timed out
        assertFalse(cache.hasRequestTimedOut(bootstrapNode1, time.milliseconds()));

        // Send a request
        cache.onRequestSent(bootstrapNode1, 1, time.milliseconds());
        assertEquals(
            Optional.empty(),
            cache.findReadyBootstrapServer(time.milliseconds())
        );
        assertFalse(cache.hasRequestTimedOut(bootstrapNode1, time.milliseconds()));

        time.sleep(requestTimeoutMs - 1);
        assertFalse(cache.hasRequestTimedOut(bootstrapNode1, time.milliseconds()));

        // Timeout the request
        time.sleep(1);
        assertTrue(cache.hasRequestTimedOut(bootstrapNode1, time.milliseconds()));
    }

    @Test
    public void testAnyInflightRequestWithAnyRequest() {
        Node otherNode = new Node(1, "other-node", 1234);
        List<Node> bootstrapList = makeBootstrapList(3);
        RequestManager cache = new RequestManager(
            bootstrapList,
            retryBackoffMs,
            requestTimeoutMs,
            random
        );

        assertFalse(cache.hasAnyInflightRequest(time.milliseconds()));

        // Send a request and check state
        cache.onRequestSent(otherNode, 11, time.milliseconds());
        assertTrue(cache.hasAnyInflightRequest(time.milliseconds()));

        // Wait until the request times out
        time.sleep(requestTimeoutMs);
        assertFalse(cache.hasAnyInflightRequest(time.milliseconds()));

        // Send another request and fail it
        cache.onRequestSent(otherNode, 12, time.milliseconds());
        cache.onResponseResult(otherNode, 12, false, time.milliseconds());
        assertFalse(cache.hasAnyInflightRequest(time.milliseconds()));

        // Send another request and mark it successful
        cache.onRequestSent(otherNode, 12, time.milliseconds());
        cache.onResponseResult(otherNode, 12, true, time.milliseconds());
        assertFalse(cache.hasAnyInflightRequest(time.milliseconds()));
    }

    private List<Node> makeBootstrapList(int numberOfNodes) {
        return IntStream.iterate(-2, id -> id - 1)
            .limit(numberOfNodes)
            .mapToObj(id -> new Node(id, String.format("mock-boot-host%d", id), 1234))
            .collect(Collectors.toList());
    }
}
