/*
 * 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.hudi.sink.partitioner.index;

import org.apache.hudi.client.common.HoodieFlinkEngineContext;
import org.apache.hudi.common.data.HoodieListData;
import org.apache.hudi.common.data.HoodiePairData;
import org.apache.hudi.common.model.HoodieRecordGlobalLocation;
import org.apache.hudi.common.table.HoodieTableMetaClient;
import org.apache.hudi.common.util.VisibleForTesting;
import org.apache.hudi.common.util.collection.Pair;
import org.apache.hudi.configuration.FlinkOptions;
import org.apache.hudi.exception.HoodieException;
import org.apache.hudi.metadata.HoodieTableMetadata;
import org.apache.hudi.sink.event.Correspondent;
import org.apache.hudi.util.StreamerUtil;

import lombok.Getter;
import lombok.extern.slf4j.Slf4j;
import org.apache.flink.configuration.Configuration;

import java.io.IOException;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * Global record-level-index backend backed by the metadata table.
 *
 * <p>The backend serves the global RLI path: lookups are keyed by record key only and return
 * {@link HoodieRecordGlobalLocation}. A checkpoint-aware {@link RecordIndexCache} keeps recently
 * accessed locations locally and falls back to metadata table reads for cache misses.
 */
@Slf4j
public class GlobalRecordLevelIndexBackend implements MinibatchIndexBackend {
  @VisibleForTesting
  @Getter
  private final RecordIndexCache recordIndexCache;
  private final Configuration conf;
  private final HoodieTableMetaClient metaClient;
  private HoodieTableMetadata metadataTable;

  /**
   * Creates a global RLI backend with a checkpoint-aware cache.
   *
   * @param conf Flink write configuration
   * @param initCheckpointId restored checkpoint id used to seed cache eviction state, or {@code -1}
   */
  public GlobalRecordLevelIndexBackend(Configuration conf, long initCheckpointId) {
    this.metaClient = StreamerUtil.createMetaClient(conf);
    this.conf = conf;
    this.recordIndexCache = new RecordIndexCache(conf, initCheckpointId);
    reloadMetadataTable();
  }

  @Override
  public HoodieRecordGlobalLocation get(String recordKey) throws IOException {
    throw new UnsupportedOperationException(this.getClass().getSimpleName() + " doesn't support lookup with a single key.");
  }

  @Override
  public void update(String recordKey, HoodieRecordGlobalLocation recordGlobalLocation) {
    recordIndexCache.update(recordKey, recordGlobalLocation);
  }

  @Override
  public Map<String, HoodieRecordGlobalLocation> get(List<String> recordKeys) throws IOException {
    Map<String, HoodieRecordGlobalLocation> keysAndLocations = new HashMap<>();
    List<String> missedKeys = new ArrayList<>();
    for (String key: recordKeys) {
      HoodieRecordGlobalLocation location = recordIndexCache.get(key);
      if (location == null) {
        missedKeys.add(key);
      } else {
        keysAndLocations.put(key, location);
      }
    }
    if (!missedKeys.isEmpty()) {
      HoodiePairData<String, HoodieRecordGlobalLocation> recordIndexData =
          metadataTable.readRecordIndexLocationsWithKeys(HoodieListData.eager(missedKeys));
      recordIndexData.forEach(keyAndLocation -> {
        recordIndexCache.update(keyAndLocation.getKey(), keyAndLocation.getValue());
        keysAndLocations.put(keyAndLocation.getKey(), keyAndLocation.getValue());
      });
    }
    return keysAndLocations;
  }

  @Override
  public void update(List<Pair<String, HoodieRecordGlobalLocation>> recordKeysAndLocations) throws IOException {
    recordKeysAndLocations.forEach(keyAndLocation -> recordIndexCache.update(keyAndLocation.getKey(), keyAndLocation.getValue()));
  }

  /**
   * Starts a new checkpoint generation in the local record-index cache.
   *
   * @param checkpointId checkpoint id from the bucket assign operator
   */
  @Override
  public void onCheckpoint(long checkpointId) {
    recordIndexCache.addCheckpointCache(checkpointId);
  }

  /**
   * Marks older cache generations evictable after a checkpoint completes.
   *
   * <p>The minimum inflight checkpoint is used as the retention boundary. If no instant is inflight,
   * the completed checkpoint id is safe because the writer coordinator has already advanced past it.
   *
   * @param correspondent writer coordinator correspondent used to query inflight instants
   * @param completedCheckpointID completed checkpoint id reported by Flink
   */
  @Override
  public void onCheckpointComplete(Correspondent correspondent, long completedCheckpointID) {
    Map<Long, String> inflightInstants = correspondent.requestInflightInstants();
    log.info("Inflight instants and the corresponding checkpoints: {}, notified completed checkpoints: {}",
        inflightInstants, completedCheckpointID);
    // if there are no inflight instants,
    // the latest completed checkpoint id is used as the minimum checkpoint id,
    // since the streaming write operator always uses previous checkpoint id to request the new instant.
    recordIndexCache.markAsEvictable(inflightInstants.keySet().stream().min(Long::compareTo).orElse(completedCheckpointID));
    this.metaClient.reloadActiveTimeline();
    reloadMetadataTable();
  }

  private void reloadMetadataTable() {
    this.metadataTable = metaClient.getTableFormat().getMetadataFactory().create(
        HoodieFlinkEngineContext.DEFAULT,
        metaClient.getStorage(),
        StreamerUtil.metadataConfig(conf),
        conf.get(FlinkOptions.PATH));
  }

  @Override
  public void close() throws IOException {
    this.recordIndexCache.close();
    if (this.metadataTable == null) {
      return;
    }
    try {
      this.metadataTable.close();
    } catch (Exception e) {
      throw new HoodieException("Exception happened during close metadata table.", e);
    }
  }
}
