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

import org.apache.kafka.common.metrics.MetricConfig;
import org.apache.kafka.common.metrics.Metrics;
import org.apache.kafka.common.metrics.Sensor;
import org.apache.kafka.common.serialization.Serde;
import org.apache.kafka.common.utils.Time;
import org.apache.kafka.streams.StreamsConfig;
import org.apache.kafka.streams.StreamsMetrics;
import org.apache.kafka.streams.Topology;
import org.apache.kafka.streams.TopologyTestDriver;
import org.apache.kafka.streams.kstream.Transformer;
import org.apache.kafka.streams.kstream.ValueTransformer;
import org.apache.kafka.streams.processor.Cancellable;
import org.apache.kafka.streams.processor.CommitCallback;
import org.apache.kafka.streams.processor.PunctuationType;
import org.apache.kafka.streams.processor.Punctuator;
import org.apache.kafka.streams.processor.StateRestoreCallback;
import org.apache.kafka.streams.processor.StateStore;
import org.apache.kafka.streams.processor.StateStoreContext;
import org.apache.kafka.streams.processor.TaskId;
import org.apache.kafka.streams.processor.internals.ClientUtils;
import org.apache.kafka.streams.processor.internals.RecordCollector;
import org.apache.kafka.streams.processor.internals.metrics.StreamsMetricsImpl;
import org.apache.kafka.streams.processor.internals.metrics.TaskMetrics;
import org.apache.kafka.streams.state.internals.InMemoryKeyValueStore;

import java.io.File;
import java.time.Duration;
import java.util.HashMap;
import java.util.LinkedList;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.Optional;
import java.util.Properties;

import static org.apache.kafka.common.utils.Utils.mkEntry;
import static org.apache.kafka.common.utils.Utils.mkMap;
import static org.apache.kafka.common.utils.Utils.mkProperties;

/**
 * {@link MockProcessorContext} is a mock of {@link ProcessorContext} for users to test their {@link Processor},
 * {@link Transformer}, and {@link ValueTransformer} implementations.
 * <p>
 * The tests for this class (org.apache.kafka.streams.MockProcessorContextTest) include several behavioral
 * tests that serve as example usage.
 * <p>
 * Note that this class does not take any automated actions (such as firing scheduled punctuators).
 * It simply captures any data it witnesses.
 * If you require more automated tests, we recommend wrapping your {@link Processor} in a minimal source-processor-sink
 * {@link Topology} and using the {@link TopologyTestDriver}.
 */
public class MockProcessorContext<KForward, VForward> implements ProcessorContext<KForward, VForward>, RecordCollector.Supplier {
    // Immutable fields ================================================
    private final StreamsMetricsImpl metrics;
    private final TaskId taskId;
    private final StreamsConfig config;
    private final File stateDir;

    // settable record metadata ================================================
    private MockRecordMetadata recordMetadata;
    private Long currentSystemTimeMs;
    private Long currentStreamTimeMs;

    // mocks ================================================
    private final Map<String, StateStore> stateStores = new HashMap<>();
    private final List<CapturedPunctuator> punctuators = new LinkedList<>();
    private final List<CapturedForward<? extends KForward, ? extends VForward>> capturedForwards = new LinkedList<>();
    private boolean committed = false;

    private static final class MockRecordMetadata implements RecordMetadata {
        private final String topic;
        private final int partition;
        private final long offset;

        private MockRecordMetadata(final String topic, final int partition, final long offset) {
            this.topic = topic;
            this.partition = partition;
            this.offset = offset;
        }

        @Override
        public String topic() {
            return topic;
        }

        @Override
        public int partition() {
            return partition;
        }

        @Override
        public long offset() {
            return offset;
        }
    }

    /**
     * {@link CapturedPunctuator} holds captured punctuators, along with their scheduling information.
     */
    public static final class CapturedPunctuator {
        private final Duration interval;
        private final PunctuationType type;
        private final Punctuator punctuator;
        private boolean cancelled = false;

        private CapturedPunctuator(final Duration interval, final PunctuationType type, final Punctuator punctuator) {
            this.interval = interval;
            this.type = type;
            this.punctuator = punctuator;
        }

        public Duration getInterval() {
            return interval;
        }

        public PunctuationType getType() {
            return type;
        }

        public Punctuator getPunctuator() {
            return punctuator;
        }

        public void cancel() {
            cancelled = true;
        }

        public boolean cancelled() {
            return cancelled;
        }
    }

    public static final class CapturedForward<K, V> {

        private final Record<K, V> record;
        private final Optional<String> childName;

        public CapturedForward(final Record<K, V> record) {
            this(record, Optional.empty());
        }

        public CapturedForward(final Record<K, V> record, final Optional<String> childName) {
            this.record = Objects.requireNonNull(record);
            this.childName = Objects.requireNonNull(childName);
        }

        /**
         * The child this data was forwarded to.
         *
         * @return If present, the child name the record was forwarded to.
         *         If empty, the forward was a broadcast.
         */
        public Optional<String> childName() {
            return childName;
        }

        /**
         * The record that was forwarded.
         *
         * @return The forwarded record. Not null.
         */
        public Record<K, V> record() {
            return record;
        }

        @Override
        public String toString() {
            return "CapturedForward{" +
                "record=" + record +
                ", childName=" + childName +
                '}';
        }

        @Override
        public boolean equals(final Object o) {
            if (this == o) return true;
            if (o == null || getClass() != o.getClass()) return false;
            final CapturedForward<?, ?> that = (CapturedForward<?, ?>) o;
            return Objects.equals(record, that.record) &&
                Objects.equals(childName, that.childName);
        }

        @Override
        public int hashCode() {
            return Objects.hash(record, childName);
        }
    }

    // constructors ================================================

    /**
     * Create a {@link MockProcessorContext} with dummy {@code config} and {@code taskId} and {@code null} {@code stateDir}.
     * Most unit tests using this mock won't need to know the taskId,
     * and most unit tests should be able to get by with the
     * {@link InMemoryKeyValueStore}, so the stateDir won't matter.
     */
    public MockProcessorContext() {
        this(
            mkProperties(mkMap(
                mkEntry(StreamsConfig.APPLICATION_ID_CONFIG, ""),
                mkEntry(StreamsConfig.BOOTSTRAP_SERVERS_CONFIG, "")
            )),
            new TaskId(0, 0),
            null
        );
    }

    /**
     * Create a {@link MockProcessorContext} with dummy {@code taskId} and {@code null} {@code stateDir}.
     * Most unit tests using this mock won't need to know the taskId,
     * and most unit tests should be able to get by with the
     * {@link InMemoryKeyValueStore}, so the stateDir won't matter.
     *
     * @param config a Properties object, used to configure the context and the processor.
     */
    public MockProcessorContext(final Properties config) {
        this(config, new TaskId(0, 0), null);
    }

    /**
     * Create a {@link MockProcessorContext} with a specified taskId and null stateDir.
     *
     * @param config   a {@link Properties} object, used to configure the context and the processor.
     * @param taskId   a {@link TaskId}, which the context makes available via {@link MockProcessorContext#taskId()}.
     * @param stateDir a {@link File}, which the context makes available viw {@link MockProcessorContext#stateDir()}.
     */
    public MockProcessorContext(final Properties config, final TaskId taskId, final File stateDir) {
        final Properties configCopy = new Properties();
        configCopy.putAll(config);
        configCopy.putIfAbsent(StreamsConfig.BOOTSTRAP_SERVERS_CONFIG, "dummy-bootstrap-host:0");
        configCopy.putIfAbsent(StreamsConfig.APPLICATION_ID_CONFIG, "dummy-mock-app-id");
        final StreamsConfig streamsConfig = new ClientUtils.QuietStreamsConfig(configCopy);
        this.taskId = taskId;
        this.config = streamsConfig;
        this.stateDir = stateDir;
        final MetricConfig metricConfig = new MetricConfig();
        metricConfig.recordLevel(Sensor.RecordingLevel.DEBUG);
        final String threadId = Thread.currentThread().getName();
        metrics = new StreamsMetricsImpl(
            new Metrics(metricConfig),
            threadId,
            streamsConfig.getString(StreamsConfig.BUILT_IN_METRICS_VERSION_CONFIG),
            Time.SYSTEM
        );
        TaskMetrics.droppedRecordsSensor(threadId, taskId.toString(), metrics);
    }

    @Override
    public String applicationId() {
        return config.getString(StreamsConfig.APPLICATION_ID_CONFIG);
    }

    @Override
    public TaskId taskId() {
        return taskId;
    }

    @Override
    public Map<String, Object> appConfigs() {
        final Map<String, Object> combined = new HashMap<>();
        combined.putAll(config.originals());
        combined.putAll(config.values());
        return combined;
    }

    @Override
    public Map<String, Object> appConfigsWithPrefix(final String prefix) {
        return config.originalsWithPrefix(prefix);
    }

    @Override
    public long currentSystemTimeMs() {
        if (currentSystemTimeMs == null) {
            throw new IllegalStateException("System time must be set before use via setCurrentSystemTimeMs().");
        }
        return currentSystemTimeMs;
    }

    @Override
    public long currentStreamTimeMs() {
        if (currentStreamTimeMs == null) {
            throw new IllegalStateException("Stream time must be set before use via setCurrentStreamTimeMs().");
        }
        return currentStreamTimeMs;
    }

    @Override
    public Serde<?> keySerde() {
        return config.defaultKeySerde();
    }

    @Override
    public Serde<?> valueSerde() {
        return config.defaultValueSerde();
    }

    @Override
    public File stateDir() {
        return Objects.requireNonNull(
            stateDir,
            "The stateDir constructor argument was needed (probably for a state store) but not supplied. " +
                "You can either reconfigure your test so that it doesn't need access to the disk " +
                "(such as using an in-memory store), or use the full MockProcessorContext constructor to supply " +
                "a non-null stateDir argument."
        );
    }

    @Override
    public StreamsMetrics metrics() {
        return metrics;
    }

    // settable record metadata ================================================

    /**
     * The context exposes these metadata for use in the processor. Normally, they are set by the Kafka Streams framework,
     * but for the purpose of driving unit tests, you can set them directly.
     *
     * @param topic     A topic name
     * @param partition A partition number
     * @param offset    A record offset
     */
    public void setRecordMetadata(final String topic,
                                  final int partition,
                                  final long offset) {
        recordMetadata = new MockRecordMetadata(topic, partition, offset);
    }

    public void setCurrentSystemTimeMs(final long currentSystemTimeMs) {
        this.currentSystemTimeMs = currentSystemTimeMs;
    }

    public void setCurrentStreamTimeMs(final long currentStreamTimeMs) {
        this.currentStreamTimeMs = currentStreamTimeMs;
    }

    @Override
    public Optional<RecordMetadata> recordMetadata() {
        return Optional.ofNullable(recordMetadata);
    }

    // mocks ================================================

    @SuppressWarnings("unchecked")
    @Override
    public <S extends StateStore> S getStateStore(final String name) {
        return (S) stateStores.get(name);
    }

    public <S extends StateStore> void addStateStore(final S stateStore) {
        stateStores.put(stateStore.name(), stateStore);
    }

    @Override
    public Cancellable schedule(final Duration interval,
                                final PunctuationType type,
                                final Punctuator callback) {
        final CapturedPunctuator capturedPunctuator = new CapturedPunctuator(interval, type, callback);

        punctuators.add(capturedPunctuator);

        return capturedPunctuator::cancel;
    }

    /**
     * Get the punctuators scheduled so far. The returned list is not affected by subsequent calls to {@code schedule(...)}.
     *
     * @return A list of captured punctuators.
     */
    public List<CapturedPunctuator> scheduledPunctuators() {
        return new LinkedList<>(punctuators);
    }

    @Override
    public <K extends KForward, V extends VForward> void forward(final Record<K, V> record) {
        forward(record, null);
    }

    @Override
    public <K extends KForward, V extends VForward> void forward(final Record<K, V> record, final String childName) {
        capturedForwards.add(new CapturedForward<>(record, Optional.ofNullable(childName)));
    }

    /**
     * Get all the forwarded data this context has observed. The returned list will not be
     * affected by subsequent interactions with the context. The data in the list is in the same order as the calls to
     * {@code forward(...)}.
     *
     * @return A list of records that were previously passed to the context.
     */
    public List<CapturedForward<? extends KForward, ? extends VForward>> forwarded() {
        return new LinkedList<>(capturedForwards);
    }

    /**
     * Get all the forwarded data this context has observed for a specific child by name.
     * The returned list will not be affected by subsequent interactions with the context.
     * The data in the list is in the same order as the calls to {@code forward(...)}.
     *
     * @param childName The child name to retrieve forwards for
     * @return A list of records that were previously passed to the context.
     */
    public List<CapturedForward<? extends KForward, ? extends VForward>> forwarded(final String childName) {
        final LinkedList<CapturedForward<? extends KForward, ? extends VForward>> result = new LinkedList<>();
        for (final CapturedForward<? extends KForward, ? extends VForward> capture : capturedForwards) {
            if (!capture.childName().isPresent() || capture.childName().equals(Optional.of(childName))) {
                result.add(capture);
            }
        }
        return result;
    }

    /**
     * Clear the captured forwarded data.
     */
    public void resetForwards() {
        capturedForwards.clear();
    }

    @Override
    public void commit() {
        committed = true;
    }

    /**
     * Whether {@link ProcessorContext#commit()} has been called in this context.
     *
     * @return {@code true} iff {@link ProcessorContext#commit()} has been called in this context since construction or reset.
     */
    public boolean committed() {
        return committed;
    }

    /**
     * Reset the commit capture to {@code false} (whether or not it was previously {@code true}).
     */
    public void resetCommit() {
        committed = false;
    }

    @Override
    public RecordCollector recordCollector() {
        // This interface is assumed by state stores that add change-logging.
        // Rather than risk a mysterious ClassCastException during unit tests, throw an explanatory exception.

        throw new UnsupportedOperationException(
            "MockProcessorContext does not provide record collection. " +
                "For processor unit tests, use an in-memory state store with change-logging disabled. " +
                "Alternatively, use the TopologyTestDriver for testing processor/store/topology integration."
        );
    }

    /**
     * Used to get a {@link StateStoreContext} for use with
     * {@link StateStore#init(StateStoreContext, StateStore)}
     * if you need to initialize a store for your tests.
     * @return a {@link StateStoreContext} that delegates to this ProcessorContext.
     */
    public StateStoreContext getStateStoreContext() {
        return new StateStoreContext() {
            @Override
            public String applicationId() {
                return MockProcessorContext.this.applicationId();
            }

            @Override
            public TaskId taskId() {
                return MockProcessorContext.this.taskId();
            }

            @Override
            public Optional<RecordMetadata> recordMetadata() {
                return MockProcessorContext.this.recordMetadata();
            }

            @Override
            public Serde<?> keySerde() {
                return MockProcessorContext.this.keySerde();
            }

            @Override
            public Serde<?> valueSerde() {
                return MockProcessorContext.this.valueSerde();
            }

            @Override
            public File stateDir() {
                return MockProcessorContext.this.stateDir();
            }

            @Override
            public StreamsMetrics metrics() {
                return MockProcessorContext.this.metrics();
            }

            @Override
            public void register(final StateStore store,
                                 final StateRestoreCallback stateRestoreCallback) {
                register(store, stateRestoreCallback, () -> { });
            }

            @Override
            public void register(final StateStore store,
                                 final StateRestoreCallback stateRestoreCallback,
                                 final CommitCallback checkpoint) {
                stateStores.put(store.name(), store);
            }

            @Override
            public Map<String, Object> appConfigs() {
                return MockProcessorContext.this.appConfigs();
            }

            @Override
            public Map<String, Object> appConfigsWithPrefix(final String prefix) {
                return MockProcessorContext.this.appConfigsWithPrefix(prefix);
            }
        };
    }
}
