/*
 * 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.Deserializer;
import org.apache.kafka.common.serialization.Serde;
import org.apache.kafka.common.serialization.Serdes;
import org.apache.kafka.common.serialization.Serializer;
import org.apache.kafka.common.utils.Bytes;
import org.apache.kafka.streams.KeyValue;
import org.apache.kafka.streams.kstream.Window;
import org.apache.kafka.streams.kstream.Windowed;
import org.apache.kafka.streams.kstream.WindowedSerdes;
import org.apache.kafka.streams.kstream.internals.SessionWindow;
import org.apache.kafka.streams.state.internals.PrefixedSessionKeySchemas.KeyFirstSessionKeySchema;
import org.apache.kafka.streams.state.internals.PrefixedSessionKeySchemas.TimeFirstSessionKeySchema;
import org.apache.kafka.streams.state.internals.SegmentedBytesStore.KeySchema;
import org.apache.kafka.test.KeyValueIteratorStub;

import org.junit.jupiter.api.AfterEach;
import org.junit.jupiter.params.ParameterizedTest;
import org.junit.jupiter.params.provider.EnumSource;

import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.function.Function;

import static java.util.Arrays.asList;
import static org.apache.kafka.common.utils.Utils.mkEntry;
import static org.apache.kafka.common.utils.Utils.mkMap;
import static org.hamcrest.MatcherAssert.assertThat;
import static org.hamcrest.core.IsEqual.equalTo;
import static org.junit.jupiter.api.Assertions.assertArrayEquals;
import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertNull;

public class SessionKeySchemaTest {
    private static final Map<SchemaType, KeySchema> SCHEMA_TYPE_MAP = mkMap(
        mkEntry(SchemaType.SessionKeySchema, new SessionKeySchema()),
        mkEntry(SchemaType.PrefixedKeyFirstSchema, new KeyFirstSessionKeySchema()),
        mkEntry(SchemaType.PrefixedTimeFirstSchema, new TimeFirstSessionKeySchema())
    );

    private static final Map<SchemaType, Function<Windowed<Bytes>, Bytes>> WINDOW_TO_STORE_BINARY_MAP = mkMap(
        mkEntry(SchemaType.SessionKeySchema, SessionKeySchema::toBinary),
        mkEntry(SchemaType.PrefixedKeyFirstSchema, KeyFirstSessionKeySchema::toBinary),
        mkEntry(SchemaType.PrefixedTimeFirstSchema, TimeFirstSessionKeySchema::toBinary)
    );

    private static final Map<SchemaType, Function<byte[], Long>> EXTRACT_END_TS_MAP = mkMap(
        mkEntry(SchemaType.SessionKeySchema, SessionKeySchema::extractEndTimestamp),
        mkEntry(SchemaType.PrefixedKeyFirstSchema, KeyFirstSessionKeySchema::extractEndTimestamp),
        mkEntry(SchemaType.PrefixedTimeFirstSchema, TimeFirstSessionKeySchema::extractEndTimestamp)
    );

    private static final Map<SchemaType, Function<byte[], Long>> EXTRACT_START_TS_MAP = mkMap(
        mkEntry(SchemaType.SessionKeySchema, SessionKeySchema::extractStartTimestamp),
        mkEntry(SchemaType.PrefixedKeyFirstSchema, KeyFirstSessionKeySchema::extractStartTimestamp),
        mkEntry(SchemaType.PrefixedTimeFirstSchema, TimeFirstSessionKeySchema::extractStartTimestamp)
    );

    @FunctionalInterface
    interface TriFunction<A, B, C, R> {
        R apply(A a, B b, C c);
    }

    private static final Map<SchemaType, TriFunction<Windowed<String>, Serializer<String>, String, byte[]>> SERDE_TO_STORE_BINARY_MAP = mkMap(
        mkEntry(SchemaType.SessionKeySchema, SessionKeySchema::toBinary),
        mkEntry(SchemaType.PrefixedKeyFirstSchema, KeyFirstSessionKeySchema::toBinary),
        mkEntry(SchemaType.PrefixedTimeFirstSchema, TimeFirstSessionKeySchema::toBinary)
    );

    private static final Map<SchemaType, TriFunction<byte[], Deserializer<String>, String, Windowed<String>>> SERDE_FROM_BYTES_MAP = mkMap(
        mkEntry(SchemaType.SessionKeySchema, SessionKeySchema::from),
        mkEntry(SchemaType.PrefixedKeyFirstSchema, KeyFirstSessionKeySchema::from),
        mkEntry(SchemaType.PrefixedTimeFirstSchema, TimeFirstSessionKeySchema::from)
    );

    private static final Map<SchemaType, Function<Bytes, Windowed<Bytes>>> FROM_BYTES_MAP = mkMap(
        mkEntry(SchemaType.SessionKeySchema, SessionKeySchema::from),
        mkEntry(SchemaType.PrefixedKeyFirstSchema, KeyFirstSessionKeySchema::from),
        mkEntry(SchemaType.PrefixedTimeFirstSchema, TimeFirstSessionKeySchema::from)
    );

    private static final Map<SchemaType, Function<byte[], Window>> EXTRACT_WINDOW = mkMap(
        mkEntry(SchemaType.SessionKeySchema, SessionKeySchema::extractWindow),
        mkEntry(SchemaType.PrefixedKeyFirstSchema, KeyFirstSessionKeySchema::extractWindow),
        mkEntry(SchemaType.PrefixedTimeFirstSchema, TimeFirstSessionKeySchema::extractWindow)
    );

    private static final Map<SchemaType, Function<byte[], byte[]>> EXTRACT_KEY_BYTES = mkMap(
        mkEntry(SchemaType.SessionKeySchema, SessionKeySchema::extractKeyBytes),
        mkEntry(SchemaType.PrefixedKeyFirstSchema, KeyFirstSessionKeySchema::extractKeyBytes),
        mkEntry(SchemaType.PrefixedTimeFirstSchema, TimeFirstSessionKeySchema::extractKeyBytes)
    );

    private final String key = "key";
    private final String topic = "topic";
    private final long startTime = 50L;
    private final long endTime = 100L;
    private final Serde<String> serde = Serdes.String();

    private final Window window = new SessionWindow(startTime, endTime);
    private final Windowed<String> windowedKey = new Windowed<>(key, window);
    private final Serde<Windowed<String>> keySerde = new WindowedSerdes.SessionWindowedSerde<>(serde);

    private KeySchema keySchema;
    private DelegatingPeekingKeyValueIterator<Bytes, Integer> iterator;
    private SchemaType schemaType;
    private Function<Windowed<Bytes>, Bytes> toBinary;
    private TriFunction<Windowed<String>, Serializer<String>, String, byte[]> serdeToBinary;
    private TriFunction<byte[], Deserializer<String>, String, Windowed<String>> serdeFromBytes;
    private Function<Bytes, Windowed<Bytes>> fromBytes;
    private Function<byte[], Long> extractStartTS;
    private Function<byte[], Long> extractEndTS;
    private Function<byte[], byte[]> extractKeyBytes;
    private Function<byte[], Window> extractWindow;

    private enum SchemaType {
        SessionKeySchema,
        PrefixedTimeFirstSchema,
        PrefixedKeyFirstSchema
    }

    private void setUp(final SchemaType type) {
        schemaType = type;
        keySchema = SCHEMA_TYPE_MAP.get(type);
        toBinary = WINDOW_TO_STORE_BINARY_MAP.get(schemaType);
        serdeToBinary = SERDE_TO_STORE_BINARY_MAP.get(schemaType);
        serdeFromBytes = SERDE_FROM_BYTES_MAP.get(schemaType);
        fromBytes = FROM_BYTES_MAP.get(schemaType);
        extractStartTS = EXTRACT_START_TS_MAP.get(schemaType);
        extractEndTS = EXTRACT_END_TS_MAP.get(schemaType);
        extractKeyBytes = EXTRACT_KEY_BYTES.get(schemaType);
        extractWindow = EXTRACT_WINDOW.get(schemaType);
        final List<KeyValue<Bytes, Integer>> keys = asList(KeyValue.pair(toBinary.apply(new Windowed<>(Bytes.wrap(new byte[]{0, 0}), new SessionWindow(0, 0))), 1),
                KeyValue.pair(toBinary.apply(new Windowed<>(Bytes.wrap(new byte[]{0}), new SessionWindow(0, 0))), 2),
                KeyValue.pair(toBinary.apply(new Windowed<>(Bytes.wrap(new byte[]{0, 0, 0}), new SessionWindow(0, 0))), 3),
                KeyValue.pair(toBinary.apply(new Windowed<>(Bytes.wrap(new byte[]{0}), new SessionWindow(10, 20))), 4),
                KeyValue.pair(toBinary.apply(new Windowed<>(Bytes.wrap(new byte[]{0, 0}), new SessionWindow(10, 20))), 5),
                KeyValue.pair(toBinary.apply(new Windowed<>(Bytes.wrap(new byte[]{0, 0, 0}), new SessionWindow(10, 20))), 6));
        iterator = new DelegatingPeekingKeyValueIterator<>("foo", new KeyValueIteratorStub<>(keys.iterator()));
    }

    @AfterEach
    public void after() {
        if (iterator != null) {
            iterator.close();
        }
    }

    @ParameterizedTest
    @EnumSource(SchemaType.class)
    public void shouldFetchExactKeysSkippingLongerKeys(final SchemaType type) {
        setUp(type);
        final Bytes key = Bytes.wrap(new byte[]{0});
        final List<Integer> result = getValues(keySchema.hasNextCondition(key, key, 0, Long.MAX_VALUE, true));
        assertThat(result, equalTo(asList(2, 4)));
    }

    @ParameterizedTest
    @EnumSource(SchemaType.class)
    public void shouldFetchExactKeySkippingShorterKeys(final SchemaType type) {
        setUp(type);
        final Bytes key = Bytes.wrap(new byte[]{0, 0});
        final HasNextCondition hasNextCondition = keySchema.hasNextCondition(key, key, 0, Long.MAX_VALUE, true);
        final List<Integer> results = getValues(hasNextCondition);
        assertThat(results, equalTo(asList(1, 5)));
    }

    @ParameterizedTest
    @EnumSource(SchemaType.class)
    public void shouldFetchAllKeysUsingNullKeys(final SchemaType type) {
        setUp(type);
        final HasNextCondition hasNextCondition = keySchema.hasNextCondition(null, null, 0, Long.MAX_VALUE, true);
        final List<Integer> results = getValues(hasNextCondition);
        assertThat(results, equalTo(asList(1, 2, 3, 4, 5, 6)));
    }
    
    @ParameterizedTest
    @EnumSource(SchemaType.class)
    public void testUpperBoundWithLargeTimestamps(final SchemaType type) {
        setUp(type);
        final Bytes upper = keySchema.upperRange(Bytes.wrap(new byte[]{0xA, 0xB, 0xC}), Long.MAX_VALUE);

        assertThat(
            "shorter key with max timestamp should be in range",
            upper.compareTo(toBinary.apply(
                new Windowed<>(
                    Bytes.wrap(new byte[]{0xA}),
                    new SessionWindow(Long.MAX_VALUE, Long.MAX_VALUE))
            )) >= 0
        );

        assertThat(
            "shorter key with max timestamp should be in range",
            upper.compareTo(toBinary.apply(
                new Windowed<>(
                    Bytes.wrap(new byte[]{0xA, 0xB}),
                    new SessionWindow(Long.MAX_VALUE, Long.MAX_VALUE))

            )) >= 0
        );

        if (schemaType == SchemaType.PrefixedTimeFirstSchema) {
            assertThat(upper, equalTo(toBinary.apply(
                new Windowed<>(Bytes.wrap(new byte[]{0xA, 0xB, 0xC}),
                    new SessionWindow(Long.MAX_VALUE, Long.MAX_VALUE))))
            );
        } else {
            assertThat(upper, equalTo(toBinary.apply(
                new Windowed<>(Bytes.wrap(new byte[]{0xA}),
                    new SessionWindow(Long.MAX_VALUE, Long.MAX_VALUE))))
            );
        }
    }

    @ParameterizedTest
    @EnumSource(SchemaType.class)
    public void testUpperBoundWithKeyBytesLargerThanFirstTimestampByte(final SchemaType type) {
        setUp(type);
        final Bytes upper = keySchema.upperRange(Bytes.wrap(new byte[]{0xA, (byte) 0x8F, (byte) 0x9F}), Long.MAX_VALUE);

        assertThat(
            "shorter key with max timestamp should be in range",
            upper.compareTo(toBinary.apply(
                new Windowed<>(
                    Bytes.wrap(new byte[]{0xA, (byte) 0x8F}),
                    new SessionWindow(Long.MAX_VALUE, Long.MAX_VALUE))
                )
            ) >= 0
        );

        assertThat(upper, equalTo(toBinary.apply(
            new Windowed<>(Bytes.wrap(new byte[]{0xA, (byte) 0x8F, (byte) 0x9F}), new SessionWindow(Long.MAX_VALUE, Long.MAX_VALUE))))
        );
    }

    @ParameterizedTest
    @EnumSource(SchemaType.class)
    public void testUpperBoundWithZeroTimestamp(final SchemaType type) {
        setUp(type);
        final Bytes upper = keySchema.upperRange(Bytes.wrap(new byte[]{0xA, 0xB, 0xC}), 0);
        final Function<Windowed<Bytes>, Bytes> toBinary = WINDOW_TO_STORE_BINARY_MAP.get(schemaType);

        if (schemaType == SchemaType.PrefixedTimeFirstSchema) {
            assertThat(upper, equalTo(toBinary.apply(
                new Windowed<>(Bytes.wrap(new byte[]{0xA, 0xB, 0xC}), new SessionWindow(0, Long.MAX_VALUE))))
            );
        } else {
            assertThat(upper, equalTo(toBinary.apply(
                new Windowed<>(Bytes.wrap(new byte[]{0xA}), new SessionWindow(0, Long.MAX_VALUE))))
            );
        }
    }

    @ParameterizedTest
    @EnumSource(SchemaType.class)
    public void testLowerBoundWithZeroTimestamp(final SchemaType type) {
        setUp(type);
        final Bytes lower = keySchema.lowerRange(Bytes.wrap(new byte[]{0xA, 0xB, 0xC}), 0);
        assertThat(lower, equalTo(toBinary.apply(new Windowed<>(Bytes.wrap(new byte[]{0xA, 0xB, 0xC}), new SessionWindow(0, 0)))));
    }

    @ParameterizedTest
    @EnumSource(SchemaType.class)
    public void testLowerBoundMatchesTrailingZeros(final SchemaType type) {
        setUp(type);
        final Bytes lower = keySchema.lowerRange(Bytes.wrap(new byte[]{0xA, 0xB, 0xC}), Long.MAX_VALUE);

        assertThat(
            "appending zeros to key should still be in range",
            lower.compareTo(toBinary.apply(
                new Windowed<>(
                    Bytes.wrap(new byte[]{0xA, 0xB, 0xC, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0}),
                    new SessionWindow(Long.MAX_VALUE, Long.MAX_VALUE))
            )) < 0
        );

        if (schemaType == SchemaType.PrefixedTimeFirstSchema) {
            assertThat(lower, equalTo(toBinary.apply(
                new Windowed<>(Bytes.wrap(new byte[]{0xA, 0xB, 0xC}), new SessionWindow(0, Long.MAX_VALUE)))));
        } else {
            assertThat(lower, equalTo(toBinary.apply(
                new Windowed<>(Bytes.wrap(new byte[]{0xA, 0xB, 0xC}), new SessionWindow(0, 0)))));
        }
    }

    @ParameterizedTest
    @EnumSource(SchemaType.class)
    public void shouldSerializeDeserialize(final SchemaType type) {
        setUp(type);
        final byte[] bytes = keySerde.serializer().serialize(topic, windowedKey);
        final Windowed<String> result = keySerde.deserializer().deserialize(topic, bytes);
        assertEquals(windowedKey, result);
    }

    @ParameterizedTest
    @EnumSource(SchemaType.class)
    public void shouldSerializeNullToNull(final SchemaType type) {
        setUp(type);
        assertNull(keySerde.serializer().serialize(topic, null));
    }

    @ParameterizedTest
    @EnumSource(SchemaType.class)
    public void shouldDeSerializeEmptyByteArrayToNull(final SchemaType type) {
        setUp(type);
        assertNull(keySerde.deserializer().deserialize(topic, new byte[0]));
    }

    @ParameterizedTest
    @EnumSource(SchemaType.class)
    public void shouldDeSerializeNullToNull(final SchemaType type) {
        setUp(type);
        assertNull(keySerde.deserializer().deserialize(topic, null));
    }

    @ParameterizedTest
    @EnumSource(SchemaType.class)
    public void shouldConvertToBinaryAndBack(final SchemaType type) {
        setUp(type);
        final byte[] serialized = serdeToBinary.apply(windowedKey, serde.serializer(), "dummy");
        final Windowed<String> result = serdeFromBytes.apply(serialized, Serdes.String().deserializer(), "dummy");
        assertEquals(windowedKey, result);
    }

    @ParameterizedTest
    @EnumSource(SchemaType.class)
    public void shouldExtractEndTimeFromBinary(final SchemaType type) {
        setUp(type);
        final byte[] serialized = serdeToBinary.apply(windowedKey, serde.serializer(), "dummy");
        assertEquals(endTime, (long) extractEndTS.apply(serialized));
    }

    @ParameterizedTest
    @EnumSource(SchemaType.class)
    public void shouldExtractStartTimeFromBinary(final SchemaType type) {
        setUp(type);
        final byte[] serialized = serdeToBinary.apply(windowedKey, serde.serializer(), "dummy");
        assertEquals(startTime, (long) extractStartTS.apply(serialized));
    }

    @ParameterizedTest
    @EnumSource(SchemaType.class)
    public void shouldExtractWindowFromBindary(final SchemaType type) {
        setUp(type);
        final byte[] serialized = serdeToBinary.apply(windowedKey, serde.serializer(), "dummy");
        assertEquals(window, extractWindow.apply(serialized));
    }

    @ParameterizedTest
    @EnumSource(SchemaType.class)
    public void shouldExtractKeyBytesFromBinary(final SchemaType type) {
        setUp(type);
        final byte[] serialized = serdeToBinary.apply(windowedKey, serde.serializer(), "dummy");
        assertArrayEquals(key.getBytes(), extractKeyBytes.apply(serialized));
    }

    @ParameterizedTest
    @EnumSource(SchemaType.class)
    public void shouldExtractKeyFromBinary(final SchemaType type) {
        setUp(type);
        final byte[] serialized = serdeToBinary.apply(windowedKey, serde.serializer(), "dummy");
        assertEquals(windowedKey, serdeFromBytes.apply(serialized, serde.deserializer(), "dummy"));
    }

    @ParameterizedTest
    @EnumSource(SchemaType.class)
    public void shouldExtractBytesKeyFromBinary(final SchemaType type) {
        setUp(type);
        final Bytes bytesKey = Bytes.wrap(key.getBytes());
        final Windowed<Bytes> windowedBytesKey = new Windowed<>(bytesKey, window);
        final Bytes serialized = toBinary.apply(windowedBytesKey);
        assertEquals(windowedBytesKey, fromBytes.apply(serialized));
    }

    private List<Integer> getValues(final HasNextCondition hasNextCondition) {
        final List<Integer> results = new ArrayList<>();
        while (hasNextCondition.hasNext(iterator)) {
            results.add(iterator.next().value);
        }
        return results;
    }

}
