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

import org.apache.flink.annotation.Internal;
import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.api.java.tuple.Tuple3;
import org.apache.flink.runtime.state.KeyGroupRange;
import org.apache.flink.runtime.state.KeyedStateHandle;
import org.apache.flink.runtime.state.SharedStateRegistry;
import org.apache.flink.runtime.state.SharedStateRegistryKey;
import org.apache.flink.runtime.state.StreamStateHandle;
import org.apache.flink.runtime.state.filesystem.FileStateHandle;
import org.apache.flink.runtime.state.memory.ByteStreamStateHandle;
import org.apache.flink.util.CloseableIterator;

import javax.annotation.Nullable;

import java.io.IOException;
import java.util.Collections;
import java.util.List;
import java.util.stream.Collectors;

/** {@link ChangelogStateHandle} implementation based on {@link StreamStateHandle}. */
@Internal
public final class ChangelogStateHandleStreamImpl implements ChangelogStateHandle {

    private static final long serialVersionUID = -8070326169926626355L;

    private final KeyGroupRange keyGroupRange;
    /** NOTE: order is important as it reflects the order of changes. */
    private final List<Tuple2<StreamStateHandle, Long>> handlesAndOffsets;

    private transient SharedStateRegistry stateRegistry;
    private final long size;

    public ChangelogStateHandleStreamImpl(
            List<Tuple2<StreamStateHandle, Long>> handlesAndOffsets,
            KeyGroupRange keyGroupRange,
            long size) {
        this.handlesAndOffsets = handlesAndOffsets;
        this.keyGroupRange = keyGroupRange;
        this.size = size;
    }

    public ChangelogStateHandleStreamImpl(
            List<Tuple3<StreamStateHandle, Long, Long>> sorted, KeyGroupRange keyGroupRange) {
        this(
                sorted.stream().map(t -> Tuple2.of(t.f0, t.f1)).collect(Collectors.toList()),
                keyGroupRange,
                sorted.stream().mapToLong(t1 -> t1.f2).sum());
    }

    @Override
    public void registerSharedStates(SharedStateRegistry stateRegistry) {
        this.stateRegistry = stateRegistry;
        handlesAndOffsets.forEach(
                handleAndOffset ->
                        stateRegistry.registerReference(
                                getKey(handleAndOffset.f0), handleAndOffset.f0));
    }

    @Override
    public KeyGroupRange getKeyGroupRange() {
        return keyGroupRange;
    }

    @Nullable
    @Override
    public KeyedStateHandle getIntersection(KeyGroupRange keyGroupRange) {
        KeyGroupRange offsets = keyGroupRange.getIntersection(keyGroupRange);
        if (offsets.getNumberOfKeyGroups() == 0) {
            return null;
        }
        return new ChangelogStateHandleStreamImpl(handlesAndOffsets, offsets, 0L /* unknown */);
    }

    @Override
    public void discardState() {
        handlesAndOffsets.forEach(
                handleAndOffset -> stateRegistry.unregisterReference(getKey(handleAndOffset.f0)));
    }

    @Override
    public long getStateSize() {
        return size;
    }

    private static SharedStateRegistryKey getKey(StreamStateHandle stateHandle) {
        // StateHandle key used in SharedStateRegistry should only be based on the file name
        // and not on backend UUID or keygroup (multiple handles can refer to the same file and
        // making keys unique will effectively disable sharing)
        if (stateHandle instanceof FileStateHandle) {
            return new SharedStateRegistryKey(
                    ((FileStateHandle) stateHandle).getFilePath().toString());
        } else if (stateHandle instanceof ByteStreamStateHandle) {
            return new SharedStateRegistryKey(
                    ((ByteStreamStateHandle) stateHandle).getHandleName());
        } else {
            return new SharedStateRegistryKey(
                    Integer.toString(System.identityHashCode(stateHandle)));
        }
    }

    public interface StateChangeStreamReader {
        CloseableIterator<StateChange> read(StreamStateHandle handle, long offset)
                throws IOException;
    }

    public List<Tuple2<StreamStateHandle, Long>> getHandlesAndOffsets() {
        return Collections.unmodifiableList(handlesAndOffsets);
    }
}

@@ -942,55 +942,8 @@
 e2;%0A
-import org.apache.flink.api.java.tuple.Tuple3;%0A
 impo
@@ -991,16 +991,16 @@
 pRange;%0A
+
 import o
@@ -1362,56 +1362,8 @@
 dle;
-%0Aimport org.apache.flink.util.CloseableIterator;
 %0A%0Aim
@@ -1399,36 +1399,8 @@
 e;%0A%0A
-import java.io.IOException;%0A
 impo
@@ -1451,44 +1451,8 @@
 ist;
-%0Aimport java.util.stream.Collectors;
 %0A%0A/*
@@ -2298,355 +2298,8 @@
  %7D%0A%0A
-    public ChangelogStateHandleStreamImpl(%0A            List%3CTuple3%3CStreamStateHandle, Long, Long%3E%3E sorted, KeyGroupRange keyGroupRange) %7B%0A        this(%0A                sorted.stream().map(t -%3E Tuple2.of(t.f0, t.f1)).collect(Collectors.toList()),%0A                keyGroupRange,%0A                sorted.stream().mapToLong(t1 -%3E t1.f2).sum());%0A    %7D%0A%0A
     
@@ -4215,24 +4215,24 @@
 eHandle)));%0A
+
         %7D%0A  
@@ -4240,181 +4240,8 @@
  %7D%0A%0A
-    public interface StateChangeStreamReader %7B%0A        CloseableIterator%3CStateChange%3E read(StreamStateHandle handle, long offset)%0A                throws IOException;%0A    %7D%0A%0A
     

Drop unused handle constructor and reader interface