/*
 * 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.fluss.flink.tiering.source.enumerator;

import org.apache.fluss.annotation.VisibleForTesting;
import org.apache.fluss.client.Connection;
import org.apache.fluss.client.ConnectionFactory;
import org.apache.fluss.client.admin.Admin;
import org.apache.fluss.client.metadata.MetadataUpdater;
import org.apache.fluss.config.Configuration;
import org.apache.fluss.flink.metrics.FlinkMetricRegistry;
import org.apache.fluss.flink.tiering.event.FailedTieringEvent;
import org.apache.fluss.flink.tiering.event.FinishedTieringEvent;
import org.apache.fluss.flink.tiering.event.TieringReachMaxDurationEvent;
import org.apache.fluss.flink.tiering.source.split.TieringSplit;
import org.apache.fluss.flink.tiering.source.split.TieringSplitGenerator;
import org.apache.fluss.flink.tiering.source.state.TieringSourceEnumeratorState;
import org.apache.fluss.metadata.TableInfo;
import org.apache.fluss.metadata.TablePath;
import org.apache.fluss.rpc.GatewayClientProxy;
import org.apache.fluss.rpc.RpcClient;
import org.apache.fluss.rpc.gateway.CoordinatorGateway;
import org.apache.fluss.rpc.messages.LakeTieringHeartbeatRequest;
import org.apache.fluss.rpc.messages.LakeTieringHeartbeatResponse;
import org.apache.fluss.rpc.messages.PbHeartbeatReqForTable;
import org.apache.fluss.rpc.messages.PbLakeTieringTableInfo;
import org.apache.fluss.rpc.metrics.ClientMetricGroup;
import org.apache.fluss.utils.MapUtils;

import org.apache.flink.api.connector.source.ReaderInfo;
import org.apache.flink.api.connector.source.SourceEvent;
import org.apache.flink.api.connector.source.SplitEnumerator;
import org.apache.flink.api.connector.source.SplitEnumeratorContext;
import org.apache.flink.api.java.tuple.Tuple3;
import org.apache.flink.metrics.groups.SplitEnumeratorMetricGroup;
import org.apache.flink.util.FlinkRuntimeException;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.annotation.Nullable;

import java.io.IOException;
import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.TreeSet;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;
import java.util.stream.Collectors;

import static org.apache.fluss.flink.tiering.source.enumerator.TieringSourceEnumerator.HeartBeatHelper.basicHeartBeat;
import static org.apache.fluss.flink.tiering.source.enumerator.TieringSourceEnumerator.HeartBeatHelper.failedTableHeartBeat;
import static org.apache.fluss.flink.tiering.source.enumerator.TieringSourceEnumerator.HeartBeatHelper.heartBeatWithRequestNewTieringTable;
import static org.apache.fluss.flink.tiering.source.enumerator.TieringSourceEnumerator.HeartBeatHelper.tieringTableHeartBeat;
import static org.apache.fluss.flink.tiering.source.enumerator.TieringSourceEnumerator.HeartBeatHelper.waitHeartbeatResponse;

/**
 * An implementation of {@link SplitEnumerator} used to request {@link TieringSplit} from Fluss
 * Cluster.
 *
 * <p>The enumerator is responsible for:
 *
 * <ul>
 *   <li>Register the Tiering Service job that the current TieringSourceEnumerator belongs to with
 *       the Fluss Cluster when the Flink Tiering job starts up.
 *   <li>Request Fluss table splits from Fluss Cluster and assigns to SourceReader to tier.
 *   <li>Un-Register the Tiering Service job that the current TieringSourceEnumerator belongs to
 *       with the Fluss Cluster when the Flink Tiering job shutdown as much as possible.
 * </ul>
 */
public class TieringSourceEnumerator
        implements SplitEnumerator<TieringSplit, TieringSourceEnumeratorState> {

    private static final Logger LOG = LoggerFactory.getLogger(TieringSourceEnumerator.class);

    private final Configuration flussConf;
    private final SplitEnumeratorContext<TieringSplit> context;
    private final ScheduledExecutorService timerService;
    private final SplitEnumeratorMetricGroup enumeratorMetricGroup;
    private final long pollTieringTableIntervalMs;
    private final List<TieringSplit> pendingSplits;
    private final Set<Integer> readersAwaitingSplit;

    private final Map<Long, Long> tieringTableEpochs;
    private final Map<Long, Long> failedTableEpochs;
    private final Map<Long, TieringFinishInfo> finishedTables;
    private final Set<Long> tieringReachMaxDurationsTables;

    // lazily instantiated
    private RpcClient rpcClient;
    private CoordinatorGateway coordinatorGateway;
    private Connection connection;
    private Admin flussAdmin;
    private TieringSplitGenerator splitGenerator;
    private int flussCoordinatorEpoch;

    private volatile boolean isFailOvering = false;

    private volatile boolean closed = false;

    public TieringSourceEnumerator(
            Configuration flussConf,
            SplitEnumeratorContext<TieringSplit> context,
            long pollTieringTableIntervalMs) {
        this.flussConf = flussConf;
        this.context = context;
        this.timerService =
                Executors.newSingleThreadScheduledExecutor(
                        r -> new Thread(r, "Tiering-Timer-Thread"));
        this.enumeratorMetricGroup = context.metricGroup();
        this.pollTieringTableIntervalMs = pollTieringTableIntervalMs;
        this.pendingSplits = Collections.synchronizedList(new ArrayList<>());
        this.readersAwaitingSplit = Collections.synchronizedSet(new TreeSet<>());
        this.tieringTableEpochs = MapUtils.newConcurrentHashMap();
        this.finishedTables = MapUtils.newConcurrentHashMap();
        this.failedTableEpochs = MapUtils.newConcurrentHashMap();
        this.tieringReachMaxDurationsTables = Collections.synchronizedSet(new TreeSet<>());
    }

    @Override
    public void start() {
        connection = ConnectionFactory.createConnection(flussConf);
        flussAdmin = connection.getAdmin();
        FlinkMetricRegistry metricRegistry = new FlinkMetricRegistry(enumeratorMetricGroup);
        ClientMetricGroup clientMetricGroup =
                new ClientMetricGroup(metricRegistry, "LakeTieringService");
        this.rpcClient = RpcClient.create(flussConf, clientMetricGroup, false);
        MetadataUpdater metadataUpdater = new MetadataUpdater(flussConf, rpcClient);
        this.coordinatorGateway =
                GatewayClientProxy.createGatewayProxy(
                        metadataUpdater::getCoordinatorServer, rpcClient, CoordinatorGateway.class);
        this.splitGenerator = new TieringSplitGenerator(flussAdmin);

        LOG.info("Starting register Tiering Service to Fluss Coordinator...");
        try {
            LakeTieringHeartbeatResponse heartbeatResponse =
                    waitHeartbeatResponse(
                            coordinatorGateway.lakeTieringHeartbeat(basicHeartBeat()));
            this.flussCoordinatorEpoch = heartbeatResponse.getCoordinatorEpoch();
            LOG.info(
                    "Register Tiering Service to Fluss Coordinator(epoch={}) success.",
                    flussCoordinatorEpoch);

        } catch (Exception e) {
            LOG.error("Register Tiering Service failed due to ", e);
            throw new FlinkRuntimeException("Register Tiering Service failed due to ", e);
        }
        this.context.callAsync(
                this::requestTieringTableSplitsViaHeartBeat,
                this::generateAndAssignSplits,
                0,
                pollTieringTableIntervalMs);
    }

    @Override
    public void handleSplitRequest(int subtaskId, @Nullable String requesterHostname) {
        if (!context.registeredReaders().containsKey(subtaskId)) {
            // reader may be failed, skip this request.
            return;
        }
        LOG.info("TieringSourceReader {} requests split.", subtaskId);
        readersAwaitingSplit.add(subtaskId);

        // If pending splits exist, assign them directly to the requesting reader
        if (!pendingSplits.isEmpty()) {
            assignSplits();
        } else {
            // Note: Ideally, only one table should be tiering at a time.
            // Here we block to request a tiering table synchronously to avoid multiple threads
            // requesting tiering tables concurrently, which would cause the enumerator to contain
            // multiple tiering tables simultaneously. This is not optimal for tiering performance.
            Tuple3<Long, Long, TablePath> tieringTable = null;
            Throwable throwable = null;
            try {
                tieringTable = this.requestTieringTableSplitsViaHeartBeat();
            } catch (Throwable t) {
                throwable = t;
            }
            this.generateAndAssignSplits(tieringTable, throwable);
        }
    }

    @Override
    public void addSplitsBack(List<TieringSplit> splits, int subtaskId) {
        readersAwaitingSplit.add(subtaskId);
        pendingSplits.addAll(splits);
        assignSplits();
    }

    @Override
    public void addReader(int subtaskId) {
        LOG.info("Adding reader: {} to Tiering Source enumerator.", subtaskId);
        Map<Integer, ReaderInfo> readerByAttempt =
                context.registeredReadersOfAttempts().get(subtaskId);
        if (readerByAttempt != null && !readerByAttempt.isEmpty()) {
            readersAwaitingSplit.add(subtaskId);
            int maxAttempt = max(readerByAttempt.keySet());
            if (maxAttempt >= 1) {
                if (isFailOvering) {
                    LOG.warn(
                            "Subtask {} (max attempt {}) registered during ongoing failover.",
                            subtaskId,
                            maxAttempt);
                } else {
                    LOG.warn(
                            "Detected failover: subtask {} has max attempt {} > 0. Triggering global failover handling.",
                            subtaskId,
                            maxAttempt);
                    // should be failover
                    isFailOvering = true;
                    handleSourceReaderFailOver();
                }

                // if registered readers equal to current parallelism, check whether all registered
                // readers have same max attempt
                if (context.registeredReadersOfAttempts().size() == context.currentParallelism()) {
                    // Check if all readers have the same max attempt number
                    Set<Integer> maxAttempts =
                            context.registeredReadersOfAttempts().values().stream()
                                    .map(_readerByAttempt -> max(_readerByAttempt.keySet()))
                                    .collect(Collectors.toSet());
                    int globalMaxAttempt = max(maxAttempts);
                    if (maxAttempts.size() == 1 && globalMaxAttempt >= 1) {
                        LOG.info(
                                "Failover completed. All {} subtasks reached the same attempt number {}. Current registered readers are {}",
                                context.currentParallelism(),
                                globalMaxAttempt,
                                context.registeredReadersOfAttempts());
                        isFailOvering = false;
                    }
                }
            }
        }
    }

    private int max(Set<Integer> integers) {
        return integers.stream().max(Integer::compareTo).orElse(-1);
    }

    @Override
    public void handleSourceEvent(int subtaskId, SourceEvent sourceEvent) {
        if (sourceEvent instanceof FinishedTieringEvent) {
            FinishedTieringEvent finishedTieringEvent = (FinishedTieringEvent) sourceEvent;
            long finishedTableId = finishedTieringEvent.getTableId();
            Long tieringEpoch = tieringTableEpochs.remove(finishedTableId);
            LOG.info("Got FinishedTieringEvent for tiering table {}. ", finishedTableId);
            if (tieringEpoch == null) {
                // shouldn't happen, warn it
                LOG.warn(
                        "The finished table {} is not in tiering table, won't report it to Fluss to mark as finished.",
                        finishedTableId);
            } else {
                boolean isForceFinished = tieringReachMaxDurationsTables.remove(finishedTableId);
                LOG.info("Before finishedTables table {}.", finishedTables);
                finishedTables.put(
                        finishedTableId, TieringFinishInfo.from(tieringEpoch, isForceFinished));
                LOG.info("After finishedTables table {}.", finishedTables);
            }
        }

        if (sourceEvent instanceof FailedTieringEvent) {
            FailedTieringEvent failedEvent = (FailedTieringEvent) sourceEvent;
            long failedTableId = failedEvent.getTableId();
            Long tieringEpoch = tieringTableEpochs.remove(failedTableId);
            LOG.info(
                    "Tiering table {} is failed, fail reason is {}.",
                    failedTableId,
                    failedEvent.failReason());
            if (tieringEpoch == null) {
                // shouldn't happen, warn it
                LOG.warn(
                        "The failed table {} is not in tiering table, won't report it to Fluss to mark as failed.",
                        failedTableId);
            } else {
                failedTableEpochs.put(failedTableId, tieringEpoch);
            }
        }

        if (!finishedTables.isEmpty() || !failedTableEpochs.isEmpty()) {
            // call one round of heartbeat to notify table has been finished or failed
            LOG.info("Finished tiering table {}.", finishedTables);
            this.context.callAsync(
                    this::requestTieringTableSplitsViaHeartBeat, this::generateAndAssignSplits);
        }
    }

    private void handleSourceReaderFailOver() {
        LOG.info(
                "Handling source reader fail over, mark current tiering table epoch {} as failed.",
                tieringTableEpochs);
        // we need to make all as failed
        failedTableEpochs.putAll(new HashMap<>(tieringTableEpochs));
        tieringTableEpochs.clear();
        tieringReachMaxDurationsTables.clear();
        // also clean all pending splits since we mark all as failed
        pendingSplits.clear();
        if (!failedTableEpochs.isEmpty()) {
            // call one round of heartbeat to notify table has been finished or failed
            this.context.callAsync(
                    this::requestTieringTableSplitsViaHeartBeat, this::generateAndAssignSplits);
        }
    }

    @VisibleForTesting
    protected void handleTableTieringReachMaxDuration(
            TablePath tablePath, long tableId, long tieringEpoch) {
        Long currentEpoch = tieringTableEpochs.get(tableId);
        if (currentEpoch != null && currentEpoch.equals(tieringEpoch)) {
            LOG.info("Table {}-{} reached max duration. Force completing.", tablePath, tableId);
            tieringReachMaxDurationsTables.add(tableId);

            for (TieringSplit tieringSplit : pendingSplits) {
                if (tieringSplit.getTableBucket().getTableId() == tableId) {
                    // mark this tiering split to skip the current round since the tiering for
                    // this table has timed out, so the tiering source reader can skip them directly
                    tieringSplit.skipCurrentRound();
                }
            }

            // broadcast the tiering reach max duration event to all readers,
            // we broadcast all for simplicity
            Set<Integer> readers = new HashSet<>(context.registeredReaders().keySet());
            for (int reader : readers) {
                TieringReachMaxDurationEvent tieringReachMaxDurationEvent =
                        new TieringReachMaxDurationEvent(tableId);
                LOG.info("Send {} to reader {}", tieringReachMaxDurationEvent, reader);
                context.sendEventToSourceReader(reader, tieringReachMaxDurationEvent);
            }
        }
    }

    private void generateAndAssignSplits(
            @Nullable Tuple3<Long, Long, TablePath> tieringTable, Throwable throwable) {
        if (throwable != null) {
            LOG.warn("Failed to request tiering table, will retry later.", throwable);
        }
        if (tieringTable != null) {
            generateTieringSplits(tieringTable);
        }
        assignSplits();
    }

    private void assignSplits() {
        // we don't assign splits during failover
        if (isFailOvering) {
            return;
        }
        if (!readersAwaitingSplit.isEmpty()) {
            final Integer[] readers = readersAwaitingSplit.toArray(new Integer[0]);
            for (Integer nextAwaitingReader : readers) {
                if (!context.registeredReaders().containsKey(nextAwaitingReader)) {
                    readersAwaitingSplit.remove(nextAwaitingReader);
                    continue;
                }
                if (!pendingSplits.isEmpty()) {
                    TieringSplit tieringSplit = pendingSplits.remove(0);
                    context.assignSplit(tieringSplit, nextAwaitingReader);
                    LOG.info("Assigning split {} to readers {}", tieringSplit, nextAwaitingReader);
                    readersAwaitingSplit.remove(nextAwaitingReader);
                }
            }
        }
    }

    private @Nullable Tuple3<Long, Long, TablePath> requestTieringTableSplitsViaHeartBeat() {
        if (closed) {
            return null;
        }
        Map<Long, TieringFinishInfo> currentFinishedTables = new HashMap<>(this.finishedTables);
        Map<Long, Long> currentFailedTableEpochs = new HashMap<>(this.failedTableEpochs);
        LakeTieringHeartbeatRequest tieringHeartbeatRequest =
                tieringTableHeartBeat(
                        basicHeartBeat(),
                        this.tieringTableEpochs,
                        currentFinishedTables,
                        currentFailedTableEpochs,
                        this.flussCoordinatorEpoch);

        Tuple3<Long, Long, TablePath> lakeTieringInfo = null;
        // report heartbeat with request table to fluss coordinator
        LOG.info(
                "currentFinishedTables: {}, currentFailedTableEpochs: {}, tieringTableEpochs: {}",
                currentFinishedTables,
                currentFailedTableEpochs,
                tieringTableEpochs);

        if (pendingSplits.isEmpty() && !readersAwaitingSplit.isEmpty()) {
            LakeTieringHeartbeatResponse heartbeatResponse =
                    waitHeartbeatResponse(
                            coordinatorGateway.lakeTieringHeartbeat(
                                    heartBeatWithRequestNewTieringTable(tieringHeartbeatRequest)));
            if (heartbeatResponse.hasTieringTable()) {
                PbLakeTieringTableInfo tieringTable = heartbeatResponse.getTieringTable();
                lakeTieringInfo =
                        Tuple3.of(
                                tieringTable.getTableId(),
                                tieringTable.getTieringEpoch(),
                                TablePath.of(
                                        tieringTable.getTablePath().getDatabaseName(),
                                        tieringTable.getTablePath().getTableName()));
                LOG.info("Tiering table {} has been requested.", lakeTieringInfo);
            } else {
                LOG.info("No available Tiering table found, will poll later.");
            }
        } else {
            // report heartbeat to fluss coordinator
            waitHeartbeatResponse(coordinatorGateway.lakeTieringHeartbeat(tieringHeartbeatRequest));
        }

        // if come to here, we can remove currentFinishedTables/failedTableEpochs to avoid send
        // in next round
        currentFinishedTables.forEach(finishedTables::remove);
        currentFailedTableEpochs.forEach(failedTableEpochs::remove);
        return lakeTieringInfo;
    }

    private void generateTieringSplits(Tuple3<Long, Long, TablePath> tieringTable)
            throws FlinkRuntimeException {
        if (tieringTable == null) {
            return;
        }
        long start = System.currentTimeMillis();
        LOG.info("Generate Tiering splits for table {}.", tieringTable.f2);
        try {
            TablePath tablePath = tieringTable.f2;
            final TableInfo tableInfo = flussAdmin.getTableInfo(tablePath).get();
            List<TieringSplit> tieringSplits =
                    populateNumberOfTieringSplits(splitGenerator.generateTableSplits(tableInfo));
            // shuffle tiering split to avoid splits tiering skew
            // after introduce tiering max duration
            Collections.shuffle(tieringSplits);
            LOG.info(
                    "Generate Tiering {} splits for table {} with cost {}ms.",
                    tieringSplits.size(),
                    tieringTable.f2,
                    System.currentTimeMillis() - start);
            if (tieringSplits.isEmpty()) {
                LOG.info(
                        "Generate Tiering splits for table {} is empty, no need to tier data.",
                        tieringTable.f2.getTableName());
                finishedTables.put(tieringTable.f0, TieringFinishInfo.from(tieringTable.f1));
            } else {
                tieringTableEpochs.put(tieringTable.f0, tieringTable.f1);
                pendingSplits.addAll(tieringSplits);

                timerService.schedule(
                        () ->
                                context.runInCoordinatorThread(
                                        () ->
                                                handleTableTieringReachMaxDuration(
                                                        tablePath,
                                                        tieringTable.f0,
                                                        tieringTable.f1)),

                        // for simplicity, we use the freshness as
                        tableInfo.getTableConfig().getDataLakeFreshness().toMillis(),
                        TimeUnit.MILLISECONDS);
            }
        } catch (Exception e) {
            LOG.warn("Fail to generate Tiering splits for table {}.", tieringTable.f2, e);
            failedTableEpochs.put(tieringTable.f0, tieringTable.f1);
        }
    }

    private List<TieringSplit> populateNumberOfTieringSplits(List<TieringSplit> tieringSplits) {
        int numberOfSplits = tieringSplits.size();
        return tieringSplits.stream()
                .map(split -> split.copy(numberOfSplits))
                .collect(Collectors.toList());
    }

    @Override
    public TieringSourceEnumeratorState snapshotState(long checkpointId) throws Exception {
        // do nothing, the downstream lake commiter will snapshot the state to Fluss Cluster
        return new TieringSourceEnumeratorState();
    }

    @Override
    public void close() throws IOException {
        closed = true;
        timerService.shutdownNow();
        if (rpcClient != null) {
            failedTableEpochs.putAll(tieringTableEpochs);
            tieringTableEpochs.clear();
            if (!failedTableEpochs.isEmpty()) {
                reportFailedTable(basicHeartBeat(), failedTableEpochs);
            }
            try {
                LOG.info("Closing Tiering Source Enumerator of at {}.", System.currentTimeMillis());
                rpcClient.close();
            } catch (Exception e) {
                LOG.error("Failed to close Tiering Source enumerator.", e);
            }
        }
        try {
            if (flussAdmin != null) {
                LOG.info("Closing Fluss Admin client...");
                flussAdmin.close();
            }
        } catch (Exception e) {
            LOG.error("Failed to close Fluss Admin client.", e);
        }
        try {
            if (connection != null) {
                LOG.info("Closing Fluss connection...");
                connection.close();
            }
        } catch (Exception e) {
            LOG.error("Failed to close Fluss connection.", e);
        }
    }

    /**
     * Report failed table to Fluss coordinator via HeartBeat, this method should be called when
     * {@link TieringSourceEnumerator} is closed or receives failed table from downstream lake
     * committer.
     */
    private void reportFailedTable(
            LakeTieringHeartbeatRequest heartbeatRequest, Map<Long, Long> failedTableEpochs)
            throws FlinkRuntimeException {
        try {
            waitHeartbeatResponse(
                    coordinatorGateway.lakeTieringHeartbeat(
                            failedTableHeartBeat(
                                    heartbeatRequest, failedTableEpochs, flussCoordinatorEpoch)));
            LOG.info("Report failed table to Fluss Coordinator success");

        } catch (Exception e) {
            LOG.error("Errors happens when report failed table to Fluss cluster.", e);
            throw new FlinkRuntimeException(
                    "Errors happens when report failed table to Fluss cluster.", e);
        }
    }

    /** A helper class to build heartbeat request. */
    @VisibleForTesting
    static class HeartBeatHelper {

        static LakeTieringHeartbeatRequest basicHeartBeat() {
            return new LakeTieringHeartbeatRequest();
        }

        static LakeTieringHeartbeatRequest heartBeatWithRequestNewTieringTable(
                LakeTieringHeartbeatRequest heartbeatRequest) {
            heartbeatRequest.setRequestTable(true);
            return heartbeatRequest;
        }

        static LakeTieringHeartbeatRequest tieringTableHeartBeat(
                LakeTieringHeartbeatRequest heartbeatRequest,
                Map<Long, Long> tieringTableEpochs,
                Map<Long, TieringFinishInfo> finishedTables,
                Map<Long, Long> failedTableEpochs,
                int coordinatorEpoch) {
            if (!tieringTableEpochs.isEmpty()) {
                heartbeatRequest.addAllTieringTables(
                        toPbHeartbeatReqForTable(tieringTableEpochs, coordinatorEpoch));
            }
            if (!finishedTables.isEmpty()) {
                Map<Long, Long> finishTieringEpochs = new HashMap<>();
                Set<Long> forceFinishedTables = new HashSet<>();
                finishedTables.forEach(
                        (tableId, tieringFinishInfo) -> {
                            finishTieringEpochs.put(tableId, tieringFinishInfo.tieringEpoch);
                            if (tieringFinishInfo.isForceFinished) {
                                forceFinishedTables.add(tableId);
                            }
                        });
                heartbeatRequest.addAllFinishedTables(
                        toPbHeartbeatReqForTable(finishTieringEpochs, coordinatorEpoch));
                for (long forceFinishedTableId : forceFinishedTables) {
                    heartbeatRequest.addForceFinishedTable(forceFinishedTableId);
                }
            }
            // add failed tiering table to heart beat request
            return failedTableHeartBeat(heartbeatRequest, failedTableEpochs, coordinatorEpoch);
        }

        private static Set<PbHeartbeatReqForTable> toPbHeartbeatReqForTable(
                Map<Long, Long> tableEpochs, int coordinatorEpoch) {
            return tableEpochs.entrySet().stream()
                    .map(
                            tieringTableEpoch ->
                                    new PbHeartbeatReqForTable()
                                            .setTableId(tieringTableEpoch.getKey())
                                            .setCoordinatorEpoch(coordinatorEpoch)
                                            .setTieringEpoch(tieringTableEpoch.getValue()))
                    .collect(Collectors.toSet());
        }

        /**
         * Report failed table to Fluss coordinator via HeartBeat, this method should be called when
         * {@link TieringSourceEnumerator} is closed or receives failed table from downstream lake
         * committer.
         */
        static LakeTieringHeartbeatRequest failedTableHeartBeat(
                LakeTieringHeartbeatRequest lakeTieringHeartbeatRequest,
                Map<Long, Long> failedTieringTableEpochs,
                int coordinatorEpoch) {
            if (!failedTieringTableEpochs.isEmpty()) {
                lakeTieringHeartbeatRequest.addAllFailedTables(
                        toPbHeartbeatReqForTable(failedTieringTableEpochs, coordinatorEpoch));
            }
            return lakeTieringHeartbeatRequest;
        }

        static LakeTieringHeartbeatResponse waitHeartbeatResponse(
                CompletableFuture<LakeTieringHeartbeatResponse> responseCompletableFuture) {
            try {
                return responseCompletableFuture.get(3, TimeUnit.MINUTES);
            } catch (Exception e) {
                LOG.error("Failed to wait heartbeat response due to ", e);
                throw new FlinkRuntimeException("Failed to wait heartbeat response due to ", e);
            }
        }
    }

    private static class TieringFinishInfo {
        /** The epoch of the tiering operation for this table. */
        long tieringEpoch;

        /**
         * Whether this table was force finished due to reaching the maximum tiering duration. When
         * a table's tiering operation exceeds the max duration (data lake freshness), it will be
         * force finished to prevent it from blocking other tables' tiering operations.
         */
        boolean isForceFinished;

        public static TieringFinishInfo from(long tieringEpoch) {
            return new TieringFinishInfo(tieringEpoch, false);
        }

        public static TieringFinishInfo from(long tieringEpoch, boolean isForceFinished) {
            return new TieringFinishInfo(tieringEpoch, isForceFinished);
        }

        private TieringFinishInfo(long tieringEpoch, boolean isForceFinished) {
            this.tieringEpoch = tieringEpoch;
            this.isForceFinished = isForceFinished;
        }
    }
}
