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

import org.apache.kafka.common.Uuid;
import org.apache.kafka.streams.KafkaStreams;
import org.apache.kafka.streams.processor.api.FixedKeyProcessor;
import org.apache.kafka.streams.processor.api.FixedKeyProcessorSupplier;
import org.apache.kafka.streams.processor.api.Processor;
import org.apache.kafka.streams.processor.api.ProcessorSupplier;
import org.apache.kafka.streams.processor.api.ProcessorWrapper;
import org.apache.kafka.streams.processor.api.WrappedFixedKeyProcessorSupplier;
import org.apache.kafka.streams.processor.api.WrappedProcessorSupplier;
import org.apache.kafka.streams.processor.internals.StoreFactory;
import org.apache.kafka.streams.state.StoreBuilder;
import org.apache.kafka.streams.utils.TestUtils.RecordingProcessorWrapper.WrapperRecorder;

import org.junit.jupiter.api.TestInfo;
import org.mockito.Mockito;

import java.lang.reflect.Method;
import java.time.Duration;
import java.util.Collections;
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 org.apache.kafka.streams.StreamsConfig.APPLICATION_ID_CONFIG;
import static org.apache.kafka.streams.StreamsConfig.BOOTSTRAP_SERVERS_CONFIG;
import static org.apache.kafka.test.TestUtils.retryOnExceptionWithTimeout;
import static org.hamcrest.MatcherAssert.assertThat;

public class TestUtils {

    public static final String PROCESSOR_WRAPPER_COUNTER_CONFIG = "wrapped.counter";

    /**
     * Waits for the given {@link KafkaStreams} instances to all be in a specific {@link KafkaStreams.State}.
     * This method uses polling, which can be more error prone and slightly slower.
     *
     * @param streamsList the list of streams instances to run.
     * @param state the expected state that all the streams to be in within timeout
     * @param timeout the time to wait for the streams to all be in the specific state.
     *
     * @throws InterruptedException if the streams doesn't change to the expected state in time.
     */
    public static void waitForApplicationState(final List<KafkaStreams> streamsList,
                                               final KafkaStreams.State state,
                                               final Duration timeout) throws InterruptedException {
        retryOnExceptionWithTimeout(timeout.toMillis(), () -> {
            final Map<KafkaStreams, KafkaStreams.State> streamsToStates = streamsList
                .stream()
                .collect(Collectors.toMap(stream -> stream, KafkaStreams::state));

            final Map<KafkaStreams, KafkaStreams.State> wrongStateMap = streamsToStates.entrySet()
                .stream()
                .filter(entry -> entry.getValue() != state)
                .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue));

            final String reason = String.format(
                "Expected all streams instances in %s to be %s within %d ms, but the following were not: %s",
                streamsList,
                state,
                timeout.toMillis(),
                wrongStateMap
            );
            assertThat(reason, wrongStateMap.isEmpty());
        });
    }

    public static String safeUniqueTestName(final TestInfo testInfo) {
        final String methodName = testInfo.getTestMethod().map(Method::getName).orElse("unknownMethodName");
        return safeUniqueTestName(methodName);
    }

    public static String safeUniqueTestName(final String testName) {
        return sanitize(testName + Uuid.randomUuid().toString());
    }

    private static String sanitize(final String str) {
        return str
            // The `-` is used in Streams' thread name as a separator and some tests rely on this.
            .replace('-', '_')
            .replace(':', '_')
            .replace('.', '_')
            .replace('[', '_')
            .replace(']', '_')
            .replace(' ', '_')
            .replace('=', '_');
    }

    /**
     * Quick method of generating a config map prepopulated with the required
     * StreamsConfig properties
     */
    public static Map<Object, Object> dummyStreamsConfigMap() {
        final Map<Object, Object> baseConfigs = new HashMap<>();
        baseConfigs.put(APPLICATION_ID_CONFIG, "dummy-app-id");
        baseConfigs.put(BOOTSTRAP_SERVERS_CONFIG, "local");
        return baseConfigs;
    }

    public static StoreFactory mockStoreFactory(final String name) {
        final StoreFactory storeFactory = Mockito.mock(StoreFactory.class);
        Mockito.when(storeFactory.storeName()).thenReturn(name);
        return storeFactory;
    }

    /**
     * Simple pass-through processor wrapper that counts the number of processors
     * it wraps.
     * To retrieve the current count, pass an instance of AtomicInteger into the configs
     * alongside the wrapper itself. Use the config key defined with {@link #PROCESSOR_WRAPPER_COUNTER_CONFIG}
     */
    public static class RecordingProcessorWrapper implements ProcessorWrapper {

        private WrapperRecorder recorder;

        @SuppressWarnings("unchecked")
        @Override
        public void configure(final Map<String, ?> configs) {
            if (configs.containsKey(PROCESSOR_WRAPPER_COUNTER_CONFIG)) {
                recorder = (WrapperRecorder) configs.get(PROCESSOR_WRAPPER_COUNTER_CONFIG);
            } else {
                recorder = new WrapperRecorder();
            }
        }

        public static class WrapperRecorder {
            private final Set<String> uniqueStores = new HashSet<>();
            private final Set<String> processorStoresCounted = new HashSet<>();
            private final Set<String> wrappedProcessorNames = Collections.synchronizedSet(new HashSet<>());

            public void wrapProcessorSupplier(final String name) {
                wrappedProcessorNames.add(name);
            }

            public void wrapStateStore(final String processorName, final String storeName) {
                if (!uniqueStores.contains(storeName)) {
                    uniqueStores.add(storeName);
                }

                final String processorStoreKey = processorName + storeName;
                if (!processorStoresCounted.contains(processorStoreKey)) {
                    processorStoresCounted.add(processorStoreKey);
                }
            }

            public int numWrappedProcessors() {
                return wrappedProcessorNames.size();
            }

            // Number of unique state stores in the topology connected to their processors via the
            // ProcessorSupplier#stores method. State stores connected to more than one processor are
            // counted only once
            public int numUniqueStateStores() {
                return uniqueStores.size();
            }

            // Number of stores connected to a processor via the ProcessorSupplier#stores method (ie the size
            // of the set returned by #stores), summed across all processors in the topology.
            // Equal to the number of unique <processorName>-<storeName>
            // pairings. Will be greater than or equal to the value of #numUniqueStateStores, as this method
            // will double count any stores connected to more than one processor
            public int numConnectedStateStores() {
                return processorStoresCounted.size();
            }

            public Set<String> wrappedProcessorNames() {
                return wrappedProcessorNames;
            }

        }

        @Override
        public <KIn, VIn, KOut, VOut> WrappedProcessorSupplier<KIn, VIn, KOut, VOut> wrapProcessorSupplier(final String processorName,
                                                                                                           final ProcessorSupplier<KIn, VIn, KOut, VOut> processorSupplier) {

            return new CountingDelegatingProcessorSupplier<>(recorder, processorName, processorSupplier);
        }

        @Override
        public <KIn, VIn, VOut> WrappedFixedKeyProcessorSupplier<KIn, VIn, VOut> wrapFixedKeyProcessorSupplier(final String processorName,
                                                                                                               final FixedKeyProcessorSupplier<KIn, VIn, VOut> processorSupplier) {
            return new CountingDelegatingFixedKeyProcessorSupplier<>(recorder, processorName, processorSupplier);
        }
    }

    private static class CountingDelegatingProcessorSupplier<KIn, VIn, KOut, VOut>
        implements WrappedProcessorSupplier<KIn, VIn, KOut, VOut> {

        private final WrapperRecorder counter;
        private final String processorName;
        private final ProcessorSupplier<KIn, VIn, KOut, VOut> delegate;

        public CountingDelegatingProcessorSupplier(final WrapperRecorder counter,
                                                   final String processorName,
                                                   final ProcessorSupplier<KIn, VIn, KOut, VOut> processorSupplier) {
            this.counter = counter;
            this.processorName = processorName;
            this.delegate = processorSupplier;

            counter.wrapProcessorSupplier(processorName);
        }

        @Override
        public Set<StoreBuilder<?>> stores() {
            final Set<StoreBuilder<?>> stores = delegate.stores();
            if (stores != null) {
                for (final StoreBuilder<?> store : stores) {
                    counter.wrapStateStore(processorName, store.name());
                }
            }
            return stores;
        }

        @Override
        public Processor<KIn, VIn, KOut, VOut> get() {
            return delegate.get();
        }
    }

    private static class CountingDelegatingFixedKeyProcessorSupplier<KIn, VIn, VOut>
        implements WrappedFixedKeyProcessorSupplier<KIn, VIn, VOut> {

        private final WrapperRecorder counter;
        private final String processorName;
        private final FixedKeyProcessorSupplier<KIn, VIn, VOut> delegate;

        public CountingDelegatingFixedKeyProcessorSupplier(final WrapperRecorder counter,
                                                           final String processorName,
                                                           final FixedKeyProcessorSupplier<KIn, VIn, VOut> processorSupplier) {
            this.counter = counter;
            this.processorName = processorName;
            this.delegate = processorSupplier;

            counter.wrapProcessorSupplier(processorName);
        }

        @Override
        public Set<StoreBuilder<?>> stores() {
            final Set<StoreBuilder<?>> stores = delegate.stores();
            if (stores != null) {
                for (final StoreBuilder<?> store : stores) {
                    counter.wrapStateStore(processorName, store.name());
                }
            }
            return stores;
        }

        @Override
        public FixedKeyProcessor<KIn, VIn, VOut> get() {
            return delegate.get();
        }
    }
}
