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

import org.apache.kafka.common.serialization.Serde;
import org.apache.kafka.common.utils.Bytes;
import org.apache.kafka.common.utils.Time;
import org.apache.kafka.streams.errors.ProcessorStateException;
import org.apache.kafka.streams.kstream.internals.WrappingNullableUtils;
import org.apache.kafka.streams.processor.StateStore;
import org.apache.kafka.streams.processor.StateStoreContext;
import org.apache.kafka.streams.processor.internals.ProcessorContextUtils;
import org.apache.kafka.streams.processor.internals.SerdeGetter;
import org.apache.kafka.streams.query.KeyQuery;
import org.apache.kafka.streams.query.MultiVersionedKeyQuery;
import org.apache.kafka.streams.query.Position;
import org.apache.kafka.streams.query.PositionBound;
import org.apache.kafka.streams.query.Query;
import org.apache.kafka.streams.query.QueryConfig;
import org.apache.kafka.streams.query.QueryResult;
import org.apache.kafka.streams.query.RangeQuery;
import org.apache.kafka.streams.query.ResultOrder;
import org.apache.kafka.streams.query.VersionedKeyQuery;
import org.apache.kafka.streams.query.internals.InternalQueryResultUtil;
import org.apache.kafka.streams.state.KeyValueStore;
import org.apache.kafka.streams.state.StateSerdes;
import org.apache.kafka.streams.state.TimestampedKeyValueStore;
import org.apache.kafka.streams.state.ValueAndTimestamp;
import org.apache.kafka.streams.state.VersionedBytesStore;
import org.apache.kafka.streams.state.VersionedKeyValueStore;
import org.apache.kafka.streams.state.VersionedRecord;
import org.apache.kafka.streams.state.VersionedRecordIterator;
import org.apache.kafka.streams.state.internals.StoreQueryUtils.QueryHandler;

import java.time.Instant;
import java.util.Map;
import java.util.Objects;

import static org.apache.kafka.common.utils.Utils.mkEntry;
import static org.apache.kafka.common.utils.Utils.mkMap;
import static org.apache.kafka.streams.processor.internals.metrics.StreamsMetricsImpl.maybeMeasureLatency;

/**
 * A metered {@link VersionedKeyValueStore} wrapper that is used for recording operation
 * metrics, and hence its inner {@link VersionedBytesStore} implementation does not need to provide
 * its own metrics collecting functionality. The inner {@code VersionedBytesStore} of this class
 * is a {@link KeyValueStore} of type &lt;Bytes,byte[]&gt;, so we use {@link Serde}s
 * to convert from &lt;K,ValueAndTimestamp&lt;V&gt&gt; to &lt;Bytes,byte[]&gt;.
 *
 * @param <K> The key type
 * @param <V> The (raw) value type
 */
public class MeteredVersionedKeyValueStore<K, V>
    extends WrappedStateStore<VersionedBytesStore, K, V>
    implements VersionedKeyValueStore<K, V> {

    private final MeteredVersionedKeyValueStoreInternal internal;

    MeteredVersionedKeyValueStore(final VersionedBytesStore inner,
                                  final String metricScope,
                                  final Time time,
                                  final Serde<K> keySerde,
                                  final Serde<V> valueSerde) {
        super(inner);
        internal = new MeteredVersionedKeyValueStoreInternal(inner, metricScope, time, keySerde, valueSerde);
    }

    /**
     * Conceptually, {@link MeteredVersionedKeyValueStore} should {@code extend}
     * {@link MeteredKeyValueStore}, but due to type conflicts, we cannot do this. (Specifically,
     * the first needs to be {@link VersionedKeyValueStore} while the second is {@link KeyValueStore}
     * and the two interfaces conflict.) Thus, we use an internal <it>instance</it> of
     * {@code MeteredKeyValueStore} to mimic inheritance instead.
     * <p>
     * It's not ideal because it requires an extra step to translate between the APIs of
     * {@link VersionedKeyValueStore} in {@link MeteredVersionedKeyValueStore} and
     * the APIs of {@link TimestampedKeyValueStore} in {@link MeteredVersionedKeyValueStoreInternal}.
     * This extra step is all that the methods of {@code MeteredVersionedKeyValueStoreInternal} do.
     * <p>
     * Note that the addition of {@link #get(Object, long)} and {@link #delete(Object, long)} in
     * this class are to match the interface of {@link VersionedKeyValueStore}.
     */
    private class MeteredVersionedKeyValueStoreInternal
        extends MeteredKeyValueStore<K, ValueAndTimestamp<V>> {

        private final VersionedBytesStore inner;
        private final Serde<V> plainValueSerde;
        private StateSerdes<K, V> plainValueSerdes;

        private final Map<Class<?>, QueryHandler> queryHandlers =
            mkMap(
                mkEntry(
                    RangeQuery.class,
                    (query, positionBound, config, store) -> runRangeQuery(query, positionBound, config)
                ),
                mkEntry(
                    KeyQuery.class,
                    (query, positionBound, config, store) -> runKeyQuery(query, positionBound, config)
                ),
                mkEntry(
                    VersionedKeyQuery.class,
                    (query, positionBound, config, store) -> runVersionedKeyQuery(query, positionBound, config)
                ),
                mkEntry(
                    MultiVersionedKeyQuery.class,
                    (query, positionBound, config, store) -> runMultiVersionedKeyQuery(query, positionBound, config)
                )
            );

        MeteredVersionedKeyValueStoreInternal(final VersionedBytesStore inner,
                                              final String metricScope,
                                              final Time time,
                                              final Serde<K> keySerde,
                                              final Serde<V> valueSerde) {
            super(
                inner,
                metricScope,
                time,
                keySerde,
                valueSerde == null
                    ? null
                    : new ValueAndTimestampSerde<>(valueSerde)
            );
            this.inner = inner;
            this.plainValueSerde = valueSerde;
        }

        public long put(final K key, final V value, final long timestamp) {
            Objects.requireNonNull(key, "key cannot be null");
            try {
                final long validTo = maybeMeasureLatency(() -> inner.put(keyBytes(key), plainValueSerdes.rawValue(value), timestamp), time, putSensor);
                maybeRecordE2ELatency();
                return validTo;
            } catch (final ProcessorStateException e) {
                final String message = String.format(e.getMessage(), key, value);
                throw new ProcessorStateException(message, e);
            }
        }

        public ValueAndTimestamp<V> get(final K key, final long asOfTimestamp) {
            Objects.requireNonNull(key, "key cannot be null");
            try {
                return maybeMeasureLatency(() -> outerValue(inner.get(keyBytes(key), asOfTimestamp)), time, getSensor);
            } catch (final ProcessorStateException e) {
                final String message = String.format(e.getMessage(), key);
                throw new ProcessorStateException(message, e);
            }
        }

        public ValueAndTimestamp<V> delete(final K key, final long timestamp) {
            Objects.requireNonNull(key, "key cannot be null");
            try {
                return maybeMeasureLatency(() -> outerValue(inner.delete(keyBytes(key), timestamp)), time, deleteSensor);
            } catch (final ProcessorStateException e) {
                final String message = String.format(e.getMessage(), key);
                throw new ProcessorStateException(message, e);
            }
        }

        @SuppressWarnings("unchecked")
        @Override
        public <R> QueryResult<R> query(final Query<R> query, final PositionBound positionBound, final QueryConfig config) {

            final long start = time.nanoseconds();
            final QueryResult<R> result;

            final QueryHandler handler = queryHandlers.get(query.getClass());
            if (handler == null) {
                result = wrapped().query(query, positionBound, config);
                if (config.isCollectExecutionInfo()) {
                    result.addExecutionInfo(
                        "Handled in " + getClass() + " in " + (time.nanoseconds() - start) + "ns");
                }
            } else {
                result = (QueryResult<R>) handler.apply(
                    query,
                    positionBound,
                    config,
                    this
                );
                if (config.isCollectExecutionInfo()) {
                    result.addExecutionInfo(
                        "Handled in " + getClass() + " with serdes "
                            + serdes + " in " + (time.nanoseconds() - start) + "ns");
                }
            }
            return result;
        }

        @SuppressWarnings("unused")
        private <R> QueryResult<R> runRangeQuery(final Query<R> query,
                                                 final PositionBound positionBound,
                                                 final QueryConfig config) {
            // throw exception for now to reserve the ability to implement this in the future
            // without clashing with users' custom implementations in the meantime
            throw new UnsupportedOperationException("Versioned stores do not support RangeQuery queries at this time.");
        }

        @SuppressWarnings("unused")
        private <R> QueryResult<R> runKeyQuery(final Query<R> query,
                                               final PositionBound positionBound,
                                               final QueryConfig config) {
            // throw exception for now to reserve the ability to implement this in the future
            // without clashing with users' custom implementations in the meantime
            throw new UnsupportedOperationException("Versioned stores do not support KeyQuery queries at this time.");
        }

        @SuppressWarnings("unchecked")
        private <R> QueryResult<R> runVersionedKeyQuery(final Query<R> query,
                                                          final PositionBound positionBound,
                                                          final QueryConfig config) {
            final QueryResult<R> result;
            final VersionedKeyQuery<K, V> typedKeyQuery = (VersionedKeyQuery<K, V>) query;
            VersionedKeyQuery<Bytes, byte[]> rawKeyQuery = VersionedKeyQuery.withKey(keyBytes(typedKeyQuery.key()));
            if (typedKeyQuery.asOfTimestamp().isPresent()) {
                rawKeyQuery = rawKeyQuery.asOf(typedKeyQuery.asOfTimestamp().get());
            }
            final QueryResult<VersionedRecord<byte[]>> rawResult =
                wrapped().query(rawKeyQuery, positionBound, config);
            if (rawResult.isSuccess() && rawResult.getResult() != null) {
                final VersionedRecord<V> versionedRecord = StoreQueryUtils.deserializeVersionedRecord(plainValueSerdes, rawResult.getResult());
                final QueryResult<VersionedRecord<V>> typedQueryResult =
                    InternalQueryResultUtil.copyAndSubstituteDeserializedResult(rawResult, versionedRecord);
                result = (QueryResult<R>) typedQueryResult;
            } else {
                // the generic type doesn't matter, since failed queries have no result set.
                result = (QueryResult<R>) rawResult;
            }
            return result;
        }

        @SuppressWarnings("unchecked")
        private <R> QueryResult<R> runMultiVersionedKeyQuery(final Query<R> query, final PositionBound positionBound, final QueryConfig config) {
            final QueryResult<R> result;
            final MultiVersionedKeyQuery<K, V> typedKeyQuery = (MultiVersionedKeyQuery<K, V>) query;

            final Instant fromTime = typedKeyQuery.fromTime().orElse(Instant.ofEpochMilli(Long.MIN_VALUE));
            final Instant toTime = typedKeyQuery.toTime().orElse(Instant.ofEpochMilli(Long.MAX_VALUE));
            if (fromTime.compareTo(toTime) > 0) {
                throw new IllegalArgumentException("The `fromTime` timestamp must be smaller than the `toTime` timestamp.");
            }
            MultiVersionedKeyQuery<Bytes, byte[]> rawKeyQuery = MultiVersionedKeyQuery.withKey(keyBytes(typedKeyQuery.key()));
            rawKeyQuery = rawKeyQuery.fromTime(fromTime).toTime(toTime);
            if (typedKeyQuery.resultOrder().equals(ResultOrder.DESCENDING)) {
                rawKeyQuery = rawKeyQuery.withDescendingTimestamps();
            } else if (typedKeyQuery.resultOrder().equals(ResultOrder.ASCENDING)) {
                rawKeyQuery = rawKeyQuery.withAscendingTimestamps();
            }

            final QueryResult<VersionedRecordIterator<byte[]>> rawResult = wrapped().query(rawKeyQuery, positionBound, config);
            if (rawResult.isSuccess()) {
                final MeteredMultiVersionedKeyQueryIterator<V> typedResult =
                        new MeteredMultiVersionedKeyQueryIterator<>(
                            rawResult.getResult(),
                            iteratorDurationSensor,
                            time,
                            StoreQueryUtils.deserializeValue(plainValueSerdes),
                            numOpenIterators,
                            openIterators
                        );
                final QueryResult<MeteredMultiVersionedKeyQueryIterator<V>> typedQueryResult =
                        InternalQueryResultUtil.copyAndSubstituteDeserializedResult(rawResult, typedResult);
                result = (QueryResult<R>) typedQueryResult;
            } else {
                // the generic type doesn't matter, since failed queries have no result set.
                result = (QueryResult<R>) rawResult;
            }
            return result;
        }

        @SuppressWarnings("unchecked")
        @Override
        protected Serde<ValueAndTimestamp<V>> prepareValueSerdeForStore(
            final Serde<ValueAndTimestamp<V>> valueSerde,
            final SerdeGetter getter
        ) {
            if (valueSerde == null) {
                return new ValueAndTimestampSerde<>((Serde<V>) getter.valueSerde());
            } else {
                return super.prepareValueSerdeForStore(valueSerde, getter);
            }
        }

        @Override
        protected void initStoreSerde(final StateStoreContext context) {
            super.initStoreSerde(context);

            // additionally init raw value serde
            final String storeName = super.name();
            final String changelogTopic = ProcessorContextUtils.changelogFor(context, storeName, Boolean.FALSE);
            plainValueSerdes = StoreSerdeInitializer.prepareStoreSerde(
                context, storeName, changelogTopic, keySerde, plainValueSerde, WrappingNullableUtils::prepareValueSerde);
        }
    }

    @Override
    public long put(final K key, final V value, final long timestamp) {
        return internal.put(key, value, timestamp);
    }

    @Override
    public VersionedRecord<V> delete(final K key, final long timestamp) {
        final ValueAndTimestamp<V> valueAndTimestamp = internal.delete(key, timestamp);
        return valueAndTimestamp == null
            ? null
            : new VersionedRecord<>(valueAndTimestamp.value(), valueAndTimestamp.timestamp());
    }

    @Override
    public VersionedRecord<V> get(final K key) {
        final ValueAndTimestamp<V> valueAndTimestamp = internal.get(key);
        return valueAndTimestamp == null
            ? null
            : new VersionedRecord<>(valueAndTimestamp.value(), valueAndTimestamp.timestamp());
    }

    @Override
    public VersionedRecord<V> get(final K key, final long asOfTimestamp) {
        final ValueAndTimestamp<V> valueAndTimestamp = internal.get(key, asOfTimestamp);
        return valueAndTimestamp == null
            ? null
            : new VersionedRecord<>(valueAndTimestamp.value(), valueAndTimestamp.timestamp());
    }

    @Override
    public String name() {
        return internal.name();
    }

    @Override
    public void init(final StateStoreContext stateStoreContext, final StateStore root) {
        internal.init(stateStoreContext, root);
    }

    @Override
    public void flush() {
        internal.flush();
    }

    @Override
    public void close() {
        internal.close();
    }

    @Override
    public boolean persistent() {
        return internal.persistent();
    }

    @Override
    public boolean isOpen() {
        return internal.isOpen();
    }

    @Override
    public <R> QueryResult<R> query(final Query<R> query, final PositionBound positionBound, final QueryConfig config) {
        return internal.query(query, positionBound, config);
    }

    @Override
    public Position getPosition() {
        return internal.getPosition();
    }
}