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

import org.apache.kafka.clients.CommonClientConfigs;
import org.apache.kafka.clients.admin.Admin;
import org.apache.kafka.clients.admin.LogDirDescription;
import org.apache.kafka.clients.admin.MockAdminClient;
import org.apache.kafka.clients.admin.NewTopic;
import org.apache.kafka.common.Node;
import org.apache.kafka.common.TopicPartition;
import org.apache.kafka.common.test.ClusterInstance;
import org.apache.kafka.common.test.api.ClusterTest;

import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.ObjectMapper;

import org.junit.jupiter.api.Test;

import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.Set;

import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
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 class LogDirsCommandTest {
    private static final String TOPIC = "test-log-dirs-topic";

    @ClusterTest(brokers = 3)
    public void testLogDirsWithoutBrokers(ClusterInstance clusterInstance) {
        createTopic(clusterInstance, TOPIC);
        try (Admin admin = Admin.create(Map.of(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, clusterInstance.bootstrapServers()))) {
            String output = assertDoesNotThrow(() -> execute(fromArgsToOptions("--bootstrap-server", clusterInstance.bootstrapServers(), "--describe"), admin));

            // check all brokers are present
            clusterInstance.brokerIds().forEach(brokerId -> assertTrue(output.contains("\"broker\":" + brokerId)));

            // check all log dirs and topic partitions are present
            Map<Integer, Map<String, LogDirDescription>> logDirs = assertDoesNotThrow(() -> admin.describeLogDirs(clusterInstance.brokerIds()).allDescriptions().get());
            assertFalse(logDirs.isEmpty());
            logDirs.forEach((brokerId, logDirInfo) ->
                logDirInfo.forEach((logDir, logDirInfoValue) -> {
                    assertTrue(output.contains("\"logDir\":\"" + logDir + "\""));
                    logDirInfoValue.replicaInfos().forEach((topicPartition, replicaInfo) -> assertTrue(output.contains("\"partition\":\"" + topicPartition + "\"")));
                }));
        }
    }

    @ClusterTest(brokers = 3)
    public void testLogDirsWithBrokers(ClusterInstance clusterInstance) {
        createTopic(clusterInstance, TOPIC);
        try (Admin admin = Admin.create(Map.of(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, clusterInstance.bootstrapServers()))) {
            int brokerId = 0;
            String output = assertDoesNotThrow(() -> execute(fromArgsToOptions("--bootstrap-server", clusterInstance.bootstrapServers(), "--broker-list", String.valueOf(brokerId), "--describe"), admin));

            // check only broker in list is present
            assertTrue(output.contains("\"broker\":" + brokerId));
            clusterInstance.brokerIds().stream().filter(id -> id != brokerId).forEach(id -> assertFalse(output.contains("\"broker\":" + id)));

            // check log dir and topic partition are present
            Map<Integer, Map<String, LogDirDescription>> logDirs = assertDoesNotThrow(() -> admin.describeLogDirs(Set.of(brokerId)).allDescriptions().get());
            assertEquals(1, logDirs.size());
            logDirs.forEach((brokerIdValue, logDirInfo) -> {
                assertFalse(logDirInfo.isEmpty());
                logDirInfo.forEach((logDir, logDirInfoValue) -> {
                    assertTrue(output.contains("\"logDir\":\"" + logDir + "\""));
                    Optional<TopicPartition> topicPartition = logDirInfoValue.replicaInfos().keySet().stream().filter(tp -> tp.topic().equals(TOPIC)).findFirst();
                    assertTrue(topicPartition.isPresent());
                    assertTrue(output.contains("\"partition\":\"" + topicPartition.get() + "\""));
                });
            });
        }
    }

    @ClusterTest
    public void testLogDirsWithNonExistentTopic(ClusterInstance clusterInstance) {
        try (Admin admin = Admin.create(Map.of(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, clusterInstance.bootstrapServers()))) {
            String output = assertDoesNotThrow(() -> execute(fromArgsToOptions("--bootstrap-server", clusterInstance.bootstrapServers(), "--topic-list", TOPIC, "--describe"), admin));
            // check all brokers are present
            clusterInstance.brokerIds().forEach(brokerId -> assertTrue(output.contains("\"broker\":" + brokerId)));

            // check log dir is present, but topic partition is not
            Map<Integer, Map<String, LogDirDescription>> logDirs = assertDoesNotThrow(() -> admin.describeLogDirs(clusterInstance.brokerIds()).allDescriptions().get());
            assertFalse(logDirs.isEmpty());
            logDirs.forEach((brokerId, logDirInfo) ->
                logDirInfo.forEach((logDir, logDirInfoValue) -> {
                    assertTrue(output.contains("\"logDir\":\"" + logDir + "\""));
                    logDirInfoValue.replicaInfos().forEach((topicPartition, replicaInfo) ->
                        assertFalse(output.contains("\"partition\":\"" + topicPartition + "\""))
                    );
                }));
        }
    }

    @ClusterTest
    public void testLogDirsWithSpecificTopic(ClusterInstance clusterInstance) {
        createTopic(clusterInstance, TOPIC);
        createTopic(clusterInstance, "other-topic");
        try (Admin admin = Admin.create(Map.of(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, clusterInstance.bootstrapServers()))) {
            String output = assertDoesNotThrow(() -> execute(fromArgsToOptions("--bootstrap-server", clusterInstance.bootstrapServers(), "--topic-list", TOPIC, "--describe"), admin));
            // check all brokers are present
            clusterInstance.brokerIds().forEach(brokerId -> assertTrue(output.contains("\"broker\":" + brokerId)));

            // check log dir is present and only one topic partition is present
            Map<Integer, Map<String, LogDirDescription>> logDirs = assertDoesNotThrow(() -> admin.describeLogDirs(clusterInstance.brokerIds()).allDescriptions().get());
            assertFalse(logDirs.isEmpty());
            assertTrue(output.contains("\"partition\":\"" + new TopicPartition(TOPIC, 0) + "\""));
            assertFalse(output.contains("\"partition\":\"" + new TopicPartition("other-topic", 0) + "\"")); // other-topic should not be present
            logDirs.forEach((brokerId, logDirInfo) ->
                logDirInfo.forEach((logDir, logDirInfoValue) -> {
                    assertTrue(output.contains("\"logDir\":\"" + logDir + "\""));
                    logDirInfoValue.replicaInfos().keySet().stream().filter(tp -> !tp.topic().equals(TOPIC)).forEach(tp ->
                        assertFalse(output.contains("\"partition\":\"" + tp + "\""))
                    );
                }));
        }
    }

    @Test
    public void shouldThrowWhenQueryingNonExistentBrokers() {
        Node broker = new Node(1, "hostname", 9092);
        try (MockAdminClient adminClient = new MockAdminClient(List.of(broker), broker)) {
            RuntimeException exception = assertThrows(RuntimeException.class, () -> execute(fromArgsToOptions("--bootstrap-server", "EMPTY", "--broker-list", "0,1,2", "--describe"), adminClient));
            assertNotNull(exception.getCause());
            assertEquals(TerseException.class, exception.getCause().getClass());
            assertEquals("ERROR: The given brokers do not exist from --broker-list: 0,2. Current existent brokers: 1", exception.getCause().getMessage());
        }
    }

    @Test
    @SuppressWarnings("unchecked")
    public void shouldNotThrowWhenDuplicatedBrokers() throws JsonProcessingException {
        Node broker = new Node(1, "hostname", 9092);
        try (MockAdminClient adminClient = new MockAdminClient(List.of(broker), broker)) {
            String standardOutput = execute(fromArgsToOptions("--bootstrap-server", "EMPTY", "--broker-list", "1,1", "--describe"), adminClient);
            String[] standardOutputLines = standardOutput.split("\n");
            assertEquals(3, standardOutputLines.length);
            Map<String, Object> information = new ObjectMapper().readValue(standardOutputLines[2], HashMap.class);
            List<Object> brokersInformation = (List<Object>) information.get("brokers");
            Integer brokerId = (Integer) ((HashMap<String, Object>) brokersInformation.get(0)).get("broker");
            assertEquals(1, brokersInformation.size());
            assertEquals(1, brokerId);
        }
    }

    @Test
    @SuppressWarnings("unchecked")
    public void shouldQueryAllBrokersIfNonSpecified() throws JsonProcessingException {
        Node brokerOne = new Node(1, "hostname", 9092);
        Node brokerTwo = new Node(2, "hostname", 9092);
        try (MockAdminClient adminClient = new MockAdminClient(List.of(brokerTwo, brokerOne), brokerOne)) {
            String standardOutput = execute(fromArgsToOptions("--bootstrap-server", "EMPTY", "--describe"), adminClient);
            String[] standardOutputLines = standardOutput.split("\n");
            assertEquals(3, standardOutputLines.length);
            Map<String, Object> information = new ObjectMapper().readValue(standardOutputLines[2], HashMap.class);
            List<Object> brokersInformation = (List<Object>) information.get("brokers");
            Set<Integer> brokerIds = new HashSet<>() {{
                    add((Integer) ((HashMap<String, Object>) brokersInformation.get(0)).get("broker"));
                    add((Integer) ((HashMap<String, Object>) brokersInformation.get(1)).get("broker"));
                }};
            assertEquals(2, brokersInformation.size());
            assertEquals(Set.of(2, 1), brokerIds);
        }
    }

    @Test
    @SuppressWarnings("unchecked")
    public void shouldQuerySpecifiedBroker() throws JsonProcessingException {
        Node brokerOne = new Node(1, "hostname", 9092);
        Node brokerTwo = new Node(2, "hostname", 9092);
        try (MockAdminClient adminClient = new MockAdminClient(List.of(brokerOne, brokerTwo), brokerOne)) {
            String standardOutput = execute(fromArgsToOptions("--bootstrap-server", "EMPTY", "--broker-list", "1", "--describe"), adminClient);
            String[] standardOutputLines = standardOutput.split("\n");
            assertEquals(3, standardOutputLines.length);
            Map<String, Object> information = new ObjectMapper().readValue(standardOutputLines[2], HashMap.class);
            List<Object> brokersInformation = (List<Object>) information.get("brokers");
            Integer brokerId = (Integer) ((HashMap<String, Object>) brokersInformation.get(0)).get("broker");
            assertEquals(1, brokersInformation.size());
            assertEquals(1, brokerId);
        }
    }

    private LogDirsCommand.LogDirsCommandOptions fromArgsToOptions(String... args) {
        return new LogDirsCommand.LogDirsCommandOptions(args);
    }

    private String execute(LogDirsCommand.LogDirsCommandOptions options, Admin adminClient) {
        Runnable runnable = () -> {
            try {
                LogDirsCommand.execute(options, adminClient);
            } catch (Exception e) {
                throw new RuntimeException(e);
            }
        };
        return ToolsTestUtils.captureStandardOut(runnable);
    }

    private void createTopic(ClusterInstance clusterInstance, String topic) {
        try (Admin admin = Admin.create(Map.of(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, clusterInstance.bootstrapServers()))) {
            assertDoesNotThrow(() -> admin.createTopics(List.of(new NewTopic(topic, Map.of(0, List.of(0))))).topicId(topic).get());
            assertDoesNotThrow(() -> clusterInstance.waitTopicCreation(topic, 1));
        }
    }
}
