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

import org.apache.flink.api.common.JobID;
import org.apache.flink.api.common.serialization.SerializerConfigImpl;
import org.apache.flink.api.common.state.StateTtlConfig;
import org.apache.flink.api.common.state.ValueState;
import org.apache.flink.api.common.state.ValueStateDescriptor;
import org.apache.flink.api.common.typeinfo.TypeInformation;
import org.apache.flink.api.common.typeutils.TypeSerializer;
import org.apache.flink.api.common.typeutils.base.IntSerializer;
import org.apache.flink.api.java.typeutils.GenericTypeInfo;
import org.apache.flink.changelog.fs.ChangelogStorageMetricGroup;
import org.apache.flink.changelog.fs.FsStateChangelogStorage;
import org.apache.flink.changelog.fs.TaskChangelogRegistry;
import org.apache.flink.core.fs.CloseableRegistry;
import org.apache.flink.core.fs.Path;
import org.apache.flink.metrics.MetricGroup;
import org.apache.flink.metrics.groups.UnregisteredMetricsGroup;
import org.apache.flink.runtime.checkpoint.CheckpointOptions;
import org.apache.flink.runtime.execution.Environment;
import org.apache.flink.runtime.mailbox.SyncMailboxExecutor;
import org.apache.flink.runtime.metrics.groups.UnregisteredMetricGroups;
import org.apache.flink.runtime.query.TaskKvStateRegistry;
import org.apache.flink.runtime.state.CheckpointStateOutputStream;
import org.apache.flink.runtime.state.CheckpointStateToolset;
import org.apache.flink.runtime.state.CheckpointStorageAccess;
import org.apache.flink.runtime.state.CheckpointStorageLocation;
import org.apache.flink.runtime.state.CheckpointStorageLocationReference;
import org.apache.flink.runtime.state.CheckpointStreamFactory;
import org.apache.flink.runtime.state.CheckpointableKeyedStateBackend;
import org.apache.flink.runtime.state.CompletedCheckpointStorageLocation;
import org.apache.flink.runtime.state.KeyGroupRange;
import org.apache.flink.runtime.state.KeyGroupedInternalPriorityQueue;
import org.apache.flink.runtime.state.KeyedStateBackendParametersImpl;
import org.apache.flink.runtime.state.KeyedStateHandle;
import org.apache.flink.runtime.state.SharedStateRegistry;
import org.apache.flink.runtime.state.SharedStateRegistryImpl;
import org.apache.flink.runtime.state.StateBackend;
import org.apache.flink.runtime.state.StateBackendTestBase;
import org.apache.flink.runtime.state.TestLocalRecoveryConfig;
import org.apache.flink.runtime.state.TestTaskStateManager;
import org.apache.flink.runtime.state.VoidNamespace;
import org.apache.flink.runtime.state.VoidNamespaceSerializer;
import org.apache.flink.runtime.state.ttl.TtlTimeProvider;
import org.apache.flink.runtime.testutils.statemigration.TestType;
import org.apache.flink.state.common.ChangelogMaterializationMetricGroup;
import org.apache.flink.state.common.PeriodicMaterializationManager;
import org.apache.flink.util.CloseableIterator;
import org.apache.flink.util.IOUtils;
import org.apache.flink.util.concurrent.Executors;

import javax.annotation.Nullable;

import java.io.File;
import java.io.IOException;
import java.util.ArrayList;
import java.util.Collection;
import java.util.Collections;
import java.util.List;
import java.util.concurrent.CompletableFuture;

import static org.apache.flink.runtime.state.StateBackendTestBase.runSnapshot;
import static org.hamcrest.MatcherAssert.assertThat;
import static org.hamcrest.Matchers.containsInAnyOrder;
import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertFalse;
import static org.junit.Assert.assertTrue;

/** Test Utilities for Changelog StateBackend. */
public class ChangelogStateBackendTestUtils {

    public static <K> CheckpointableKeyedStateBackend<K> createKeyedBackend(
            StateBackend stateBackend,
            TypeSerializer<K> keySerializer,
            int numberOfKeyGroups,
            KeyGroupRange keyGroupRange,
            Environment env)
            throws Exception {

        return createKeyedBackend(
                stateBackend,
                keySerializer,
                numberOfKeyGroups,
                keyGroupRange,
                new UnregisteredMetricsGroup(),
                null,
                env);
    }

    public static <K> CheckpointableKeyedStateBackend<K> createKeyedBackend(
            StateBackend stateBackend,
            TypeSerializer<K> keySerializer,
            int numberOfKeyGroups,
            KeyGroupRange keyGroupRange,
            MetricGroup metricGroup,
            Environment env)
            throws Exception {

        return createKeyedBackend(
                stateBackend,
                keySerializer,
                numberOfKeyGroups,
                keyGroupRange,
                metricGroup,
                null,
                env);
    }

    public static <K> CheckpointableKeyedStateBackend<K> createKeyedBackend(
            StateBackend stateBackend,
            TypeSerializer<K> keySerializer,
            int numberOfKeyGroups,
            KeyGroupRange keyGroupRange,
            MetricGroup metricGroup,
            KeyedStateHandle state,
            Environment env)
            throws Exception {

        JobID jobID = new JobID();
        TaskKvStateRegistry kvStateRegistry = env.getTaskKvStateRegistry();
        Collection<KeyedStateHandle> stateHandles =
                state == null ? Collections.emptyList() : Collections.singletonList(state);
        return stateBackend.createKeyedStateBackend(
                new KeyedStateBackendParametersImpl<>(
                        env,
                        jobID,
                        "test_op",
                        keySerializer,
                        numberOfKeyGroups,
                        keyGroupRange,
                        kvStateRegistry,
                        TtlTimeProvider.DEFAULT,
                        metricGroup,
                        stateHandles,
                        new CloseableRegistry()));
    }

    public static CheckpointableKeyedStateBackend<Integer> createKeyedBackend(
            StateBackend stateBackend, Environment env) throws Exception {

        return createKeyedBackend(
                stateBackend,
                IntSerializer.INSTANCE,
                10,
                new KeyGroupRange(0, 9),
                new UnregisteredMetricsGroup(),
                env);
    }

    public static CheckpointableKeyedStateBackend<Integer> createKeyedBackend(
            StateBackend stateBackend, Environment env, MetricGroup metricGroup) throws Exception {

        return createKeyedBackend(
                stateBackend,
                IntSerializer.INSTANCE,
                10,
                new KeyGroupRange(0, 9),
                metricGroup,
                env);
    }

    private static CheckpointableKeyedStateBackend<Integer> restoreKeyedBackend(
            StateBackend stateBackend, KeyedStateHandle state, Environment env) throws Exception {

        return createKeyedBackend(
                stateBackend,
                IntSerializer.INSTANCE,
                10,
                new KeyGroupRange(0, 9),
                new UnregisteredMetricsGroup(),
                state,
                env);
    }

    public static TestTaskStateManager createTaskStateManager(File changelogStoragePath)
            throws IOException {
        return TestTaskStateManager.builder()
                .setStateChangelogStorage(
                        new FsStateChangelogStorage(
                                JobID.generate(),
                                Path.fromLocalFile(changelogStoragePath),
                                false,
                                1024,
                                new ChangelogStorageMetricGroup(
                                        UnregisteredMetricGroups
                                                .createUnregisteredTaskManagerJobMetricGroup()),
                                TaskChangelogRegistry.NO_OP,
                                TestLocalRecoveryConfig.disabled()))
                .build();
    }

    public static void testMaterializedRestore(
            StateBackend stateBackend,
            StateTtlConfig stateTtlConfig,
            Environment env,
            CheckpointStreamFactory streamFactory)
            throws Exception {
        SharedStateRegistry sharedStateRegistry = new SharedStateRegistryImpl();

        TypeInformation<StateBackendTestBase.TestPojo> pojoType =
                new GenericTypeInfo<>(StateBackendTestBase.TestPojo.class);
        ValueStateDescriptor<StateBackendTestBase.TestPojo> kvId =
                new ValueStateDescriptor<>("id", pojoType);
        if (stateTtlConfig.isEnabled()) {
            kvId.enableTimeToLive(stateTtlConfig);
        }

        ChangelogKeyedStateBackend<Integer> keyedBackend =
                (ChangelogKeyedStateBackend<Integer>) createKeyedBackend(stateBackend, env);

        CompletableFuture<Void> asyncComplete = new CompletableFuture<>();
        PeriodicMaterializationManager periodicMaterializationManager =
                periodicMaterializationManager(keyedBackend, asyncComplete);

        try {
            ValueState<StateBackendTestBase.TestPojo> state =
                    keyedBackend.getPartitionedState(
                            VoidNamespace.INSTANCE, VoidNamespaceSerializer.INSTANCE, kvId);

            keyedBackend.setCurrentKey(1);
            state.update(new StateBackendTestBase.TestPojo("u1", 1));

            keyedBackend.setCurrentKey(2);
            state.update(new StateBackendTestBase.TestPojo("u2", 2));

            periodicMaterializationManager.triggerMaterialization();

            keyedBackend.setCurrentKey(2);
            state.update(new StateBackendTestBase.TestPojo("u2", 22));

            keyedBackend.setCurrentKey(3);
            state.update(new StateBackendTestBase.TestPojo("u3", 3));

            periodicMaterializationManager.triggerMaterialization();

            keyedBackend.setCurrentKey(4);
            state.update(new StateBackendTestBase.TestPojo("u4", 4));

            keyedBackend.setCurrentKey(2);
            state.update(new StateBackendTestBase.TestPojo("u2", 222));

            KeyedStateHandle snapshot =
                    runSnapshot(
                            keyedBackend.snapshot(
                                    682375462378L,
                                    2,
                                    streamFactory,
                                    CheckpointOptions.forCheckpointWithDefaultLocation()),
                            sharedStateRegistry);

            IOUtils.closeQuietly(keyedBackend);
            keyedBackend.dispose();

            // make sure the asycn phase completes successfully
            if (asyncComplete.isCompletedExceptionally()) {
                asyncComplete.get();
            }

            // ============================ restore snapshot ===============================

            ((SerializerConfigImpl) env.getExecutionConfig().getSerializerConfig())
                    .registerKryoType(StateBackendTestBase.TestPojo.class);

            keyedBackend =
                    (ChangelogKeyedStateBackend<Integer>)
                            restoreKeyedBackend(stateBackend, snapshot, env);
            snapshot.discardState();

            state =
                    keyedBackend.getPartitionedState(
                            VoidNamespace.INSTANCE, VoidNamespaceSerializer.INSTANCE, kvId);

            keyedBackend.setCurrentKey(1);
            assertEquals(new StateBackendTestBase.TestPojo("u1", 1), state.value());

            keyedBackend.setCurrentKey(2);
            assertEquals(new StateBackendTestBase.TestPojo("u2", 222), state.value());

            keyedBackend.setCurrentKey(3);
            assertEquals(new StateBackendTestBase.TestPojo("u3", 3), state.value());
        } finally {
            IOUtils.closeQuietly(keyedBackend);
            keyedBackend.dispose();
        }
    }

    public static void testMaterializedRestoreForPriorityQueue(
            StateBackend stateBackend, Environment env, CheckpointStreamFactory streamFactory)
            throws Exception {
        SharedStateRegistry sharedStateRegistry = new SharedStateRegistryImpl();
        String fieldName = "key-grouped-priority-queue";
        ChangelogKeyedStateBackend<Integer> keyedBackend =
                (ChangelogKeyedStateBackend<Integer>) createKeyedBackend(stateBackend, env);

        CompletableFuture<Void> asyncComplete = new CompletableFuture<>();
        PeriodicMaterializationManager periodicMaterializationManager =
                periodicMaterializationManager(keyedBackend, asyncComplete);

        try {
            KeyGroupedInternalPriorityQueue<TestType> priorityQueue =
                    keyedBackend.create(fieldName, new TestType.V1TestTypeSerializer());

            TestType elementA100 = new TestType("a", 100);
            TestType elementA10 = new TestType("a", 10);
            TestType elementA20 = new TestType("a", 20);

            assertTrue(priorityQueue.add(elementA100));
            assertTrue(priorityQueue.add(elementA10));
            assertFalse(priorityQueue.add(elementA20));
            assertFalse(priorityQueue.add(elementA10));

            List<TestType> actualList = new ArrayList<>();
            try (CloseableIterator<TestType> iterator = priorityQueue.iterator()) {
                iterator.forEachRemaining(actualList::add);
            }

            assertThat(actualList, containsInAnyOrder(elementA100, elementA10, elementA20));

            periodicMaterializationManager.triggerMaterialization();

            TestType elementB9 = new TestType("b", 9);
            assertTrue(priorityQueue.add(elementB9));

            periodicMaterializationManager.triggerMaterialization();

            TestType elementC9 = new TestType("c", 9);
            TestType elementC8 = new TestType("c", 8);
            assertFalse(priorityQueue.add(elementC9));
            assertTrue(priorityQueue.add(elementC8));

            KeyedStateHandle snapshot =
                    runSnapshot(
                            keyedBackend.snapshot(
                                    682375462378L,
                                    2,
                                    streamFactory,
                                    CheckpointOptions.forCheckpointWithDefaultLocation()),
                            sharedStateRegistry);
            IOUtils.closeQuietly(keyedBackend);
            keyedBackend.dispose();

            // make sure the asycn phase completes successfully
            if (asyncComplete.isCompletedExceptionally()) {
                asyncComplete.get();
            }

            // ============================ restore snapshot ===============================

            keyedBackend =
                    (ChangelogKeyedStateBackend<Integer>)
                            restoreKeyedBackend(stateBackend, snapshot, env);
            snapshot.discardState();

            KeyGroupedInternalPriorityQueue<TestType> priorityQueueRestored =
                    keyedBackend.create(fieldName, new TestType.V1TestTypeSerializer());

            List<TestType> actualListRestore = new ArrayList<>();
            try (CloseableIterator<TestType> iterator = priorityQueueRestored.iterator()) {
                iterator.forEachRemaining(actualListRestore::add);
            }

            assertThat(
                    actualListRestore,
                    containsInAnyOrder(
                            elementA100, elementA10, elementA20, elementB9, elementC9, elementC8));

            assertFalse(priorityQueueRestored.add(new TestType("d", 11)));
        } finally {
            IOUtils.closeQuietly(keyedBackend);
            keyedBackend.dispose();
        }
    }

    private static PeriodicMaterializationManager periodicMaterializationManager(
            ChangelogKeyedStateBackend<Integer> keyedBackend,
            CompletableFuture<Void> asyncComplete) {
        return periodicMaterializationManager(
                keyedBackend,
                UnregisteredMetricGroups.createUnregisteredOperatorMetricGroup(),
                asyncComplete);
    }

    static PeriodicMaterializationManager periodicMaterializationManager(
            ChangelogKeyedStateBackend<Integer> keyedBackend,
            MetricGroup metricGroup,
            CompletableFuture<Void> asyncComplete) {
        return new PeriodicMaterializationManager(
                new SyncMailboxExecutor(),
                Executors.newDirectExecutorService(),
                "testTask",
                (message, exception) -> asyncComplete.completeExceptionally(exception),
                keyedBackend,
                new ChangelogMaterializationMetricGroup(metricGroup),
                true,
                10,
                1,
                "testTask");
    }

    /** Dummy {@link CheckpointStorageAccess}. */
    public static class DummyCheckpointingStorageAccess implements CheckpointStorageAccess {

        DummyCheckpointingStorageAccess() {}

        @Override
        public boolean supportsHighlyAvailableStorage() {
            return false;
        }

        @Override
        public boolean hasDefaultSavepointLocation() {
            return false;
        }

        @Override
        public CompletedCheckpointStorageLocation resolveCheckpoint(String externalPointer) {
            throw new UnsupportedOperationException(
                    "Checkpoints are not supported in a single key state backend");
        }

        @Override
        public void initializeBaseLocationsForCheckpoint() {}

        @Override
        public CheckpointStorageLocation initializeLocationForCheckpoint(long checkpointId) {
            throw new UnsupportedOperationException(
                    "Checkpoints are not supported in a single key state backend");
        }

        @Override
        public CheckpointStorageLocation initializeLocationForSavepoint(
                long checkpointId, @Nullable String externalLocationPointer) {
            throw new UnsupportedOperationException(
                    "Checkpoints are not supported in a single key state backend");
        }

        @Override
        public CheckpointStreamFactory resolveCheckpointStorageLocation(
                long checkpointId, CheckpointStorageLocationReference reference) {
            throw new UnsupportedOperationException(
                    "Checkpoints are not supported in a single key state backend");
        }

        @Override
        public CheckpointStateOutputStream createTaskOwnedStateStream() {
            throw new UnsupportedOperationException(
                    "Checkpoints are not supported in a single key state backend");
        }

        @Override
        public CheckpointStateToolset createTaskOwnedCheckpointStateToolset() {
            throw new UnsupportedOperationException(
                    "Checkpoints are not supported in a single key state backend");
        }
    }
}
