/*
 * 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.pulsar.broker.service.schema;

import static com.google.common.collect.Lists.newArrayList;
import static com.google.protobuf.ByteString.copyFrom;
import static java.util.Objects.isNull;
import static java.util.concurrent.CompletableFuture.completedFuture;
import static org.apache.pulsar.broker.service.schema.BookkeeperSchemaStorage.Functions.newSchemaEntry;
import static org.apache.pulsar.metadata.api.MetadataStoreException.AlreadyExistsException;
import static org.apache.pulsar.metadata.api.MetadataStoreException.BadVersionException;
import com.google.common.annotations.VisibleForTesting;
import java.io.IOException;
import java.nio.ByteBuffer;
import java.util.ArrayList;
import java.util.Collections;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.Set;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.CompletionException;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;
import java.util.concurrent.ExecutionException;
import java.util.function.Function;
import java.util.stream.Collectors;
import org.apache.bookkeeper.client.AsyncCallback;
import org.apache.bookkeeper.client.BKException;
import org.apache.bookkeeper.client.BookKeeper;
import org.apache.bookkeeper.client.LedgerEntry;
import org.apache.bookkeeper.client.LedgerHandle;
import org.apache.bookkeeper.mledger.impl.LedgerMetadataUtils;
import org.apache.commons.lang3.tuple.Pair;
import org.apache.pulsar.broker.PulsarService;
import org.apache.pulsar.broker.ServiceConfiguration;
import org.apache.pulsar.broker.service.schema.SchemaStorageFormat.IndexEntry;
import org.apache.pulsar.broker.service.schema.SchemaStorageFormat.SchemaLocator;
import org.apache.pulsar.broker.service.schema.exceptions.IncompatibleSchemaException;
import org.apache.pulsar.broker.service.schema.exceptions.SchemaException;
import org.apache.pulsar.common.policies.data.SchemaMetadata;
import org.apache.pulsar.common.protocol.schema.SchemaStorage;
import org.apache.pulsar.common.protocol.schema.SchemaVersion;
import org.apache.pulsar.common.protocol.schema.StoredSchema;
import org.apache.pulsar.common.schema.LongSchemaVersion;
import org.apache.pulsar.common.util.FutureUtil;
import org.apache.pulsar.metadata.api.MetadataCache;
import org.apache.pulsar.metadata.api.MetadataSerde;
import org.apache.pulsar.metadata.api.MetadataStoreException;
import org.apache.pulsar.metadata.api.Stat;
import org.apache.pulsar.metadata.api.extended.MetadataStoreExtended;
import org.jspecify.annotations.NonNull;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

public class BookkeeperSchemaStorage implements SchemaStorage {
    private static final Logger log = LoggerFactory.getLogger(BookkeeperSchemaStorage.class);

    private static final String SchemaPath = "/schemas";
    private static final byte[] LedgerPassword = "".getBytes();

    private final MetadataStoreExtended store;
    private final PulsarService pulsar;
    private final MetadataCache<SchemaStorageFormat.SchemaLocator> locatorEntryCache;

    private final ServiceConfiguration config;
    private BookKeeper bookKeeper;

    private final ConcurrentMap<String, CompletableFuture<StoredSchema>> readSchemaOperations =
            new ConcurrentHashMap<>();

    @VisibleForTesting
    BookkeeperSchemaStorage(PulsarService pulsar) {
        this.pulsar = pulsar;
        this.store = pulsar.getLocalMetadataStore();
        this.config = pulsar.getConfiguration();
        this.locatorEntryCache = store.getMetadataCache(new MetadataSerde<SchemaStorageFormat.SchemaLocator>() {
            @Override
            public byte[] serialize(String path, SchemaStorageFormat.SchemaLocator value) {
                return value.toByteArray();
            }

            @Override
            public SchemaStorageFormat.SchemaLocator deserialize(String path, byte[] content, Stat stat)
                    throws IOException {
                return SchemaStorageFormat.SchemaLocator.parseFrom(content);
            }
        });
    }

    @Override
    public void start() throws IOException {
        this.bookKeeper = pulsar.getBookKeeperClientFactory().create(
            pulsar.getConfiguration(),
            store,
            pulsar.getIoEventLoopGroup(),
            Optional.empty(),
            null
        ).join();
    }

    @Override
    public CompletableFuture<SchemaVersion> put(String key, byte[] value, byte[] hash) {
        return putSchema(key, value, hash).thenApply(LongSchemaVersion::new);
    }

    @Override
    public CompletableFuture<SchemaVersion> put(String key,
            Function<CompletableFuture<List<CompletableFuture<StoredSchema>>>,
                    CompletableFuture<Pair<byte[], byte[]>>> fn) {
        CompletableFuture<SchemaVersion> promise = new CompletableFuture<>();
        put(key, fn, promise);
        return promise;
    }

    private void put(String key,
             Function<CompletableFuture<List<CompletableFuture<StoredSchema>>>,
             CompletableFuture<Pair<byte[], byte[]>>> fn,
             CompletableFuture<SchemaVersion> promise) {
        CompletableFuture<Pair<Optional<LocatorEntry>, List<CompletableFuture<StoredSchema>>>> schemasWithLocator =
                getAllWithLocator(key);
        schemasWithLocator.thenCompose(pair ->
                fn.apply(completedFuture(pair.getRight())).thenCompose(p -> {
                    // The schema is existed
                    if (p == null) {
                        return CompletableFuture.completedFuture(null);
                    }
                    return putSchema(key, p.getLeft(), p.getRight(), pair.getLeft());
                }).thenApply(version -> {
                    return version != null ? new LongSchemaVersion(version) : null;
                })).whenComplete((v, ex) -> {
                    if (ex == null) {
                        promise.complete(v);
                    } else {
                        Throwable cause = FutureUtil.unwrapCompletionException(ex);
                        if (cause instanceof AlreadyExistsException || cause instanceof BadVersionException) {
                            put(key, fn, promise);
                        } else {
                            promise.completeExceptionally(ex);
                        }
                    }
        });
    }

    @Override
    public CompletableFuture<StoredSchema> get(String key, SchemaVersion version) {
        if (version == SchemaVersion.Latest) {
            return getSchema(key);
        } else {
            LongSchemaVersion longVersion = (LongSchemaVersion) version;
            return getSchema(key, longVersion.getVersion());
        }
    }

    @Override
    public CompletableFuture<List<CompletableFuture<StoredSchema>>> getAll(String key) {
        return getAllWithLocator(key).thenApply(Pair::getRight);
    }

    private CompletableFuture<Pair<Optional<LocatorEntry>, List<CompletableFuture<StoredSchema>>>> getAllWithLocator(
            String key) {
        return getLocator(key).thenApply(locator -> {
            if (log.isDebugEnabled()) {
                log.debug("[{}] Get all schemas - locator: {}", key, locator);
            }

            if (locator.isEmpty()) {
                return Pair.of(locator, Collections.emptyList());
            }

            SchemaStorageFormat.SchemaLocator schemaLocator = locator.get().locator;
            List<CompletableFuture<StoredSchema>> list = new ArrayList<>();
            schemaLocator.getIndexList().forEach(indexEntry -> list.add(readSchemaEntry(indexEntry.getPosition())
                    .thenApply(entry -> new StoredSchema
                            (
                                    entry.getSchemaData().toByteArray(),
                                    new LongSchemaVersion(indexEntry.getVersion())
                            )
                    )
            ));
            return Pair.of(locator, list);
        });
    }

    CompletableFuture<Optional<LocatorEntry>> getLocator(String key) {
        return getSchemaLocator(getSchemaPath(key));
    }

    public List<Long> getSchemaLedgerList(String key) throws IOException {
        Optional<LocatorEntry> locatorEntry = null;
        try {
            locatorEntry = getLocator(key).get();
        } catch (Exception e) {
            log.warn("Failed to get list of schema-storage ledger for {}, the exception as follow: \n {}", key,
                    (e instanceof ExecutionException ? e.getCause() : e));
            throw new IOException("Failed to get schema ledger for" + key);
        }
        LocatorEntry entry = locatorEntry.orElse(null);
        return entry != null ? entry.locator.getIndexList().stream().map(i -> i.getPosition().getLedgerId())
                .collect(Collectors.toList()) : null;
    }

    @VisibleForTesting
    BookKeeper getBookKeeper() {
        return bookKeeper;
    }

    @Override
    public CompletableFuture<SchemaVersion> delete(String key, boolean forcefully) {
        return deleteSchema(key, forcefully).thenApply(version -> {
            if (version == null) {
                return null;
            }
            return new LongSchemaVersion(version);
        });
    }

    @Override
    public CompletableFuture<SchemaVersion> delete(String key) {
        return delete(key, false);
    }

    @NonNull
    private CompletableFuture<StoredSchema> getSchema(String schemaId) {
        // There's already a schema read operation in progress. Just piggyback on that
        return readSchemaOperations.computeIfAbsent(schemaId, key -> {
            if (log.isDebugEnabled()) {
                log.debug("[{}] Fetching schema from store", schemaId);
            }
            return getSchemaLocator(getSchemaPath(schemaId)).thenCompose(locator -> {
                if (log.isDebugEnabled()) {
                    log.debug("[{}] Got schema locator {}", schemaId, locator);
                }
                if (!locator.isPresent()) {
                    return completedFuture(null);
                }

                SchemaStorageFormat.SchemaLocator schemaLocator = locator.get().locator;

                return readSchemaEntry(schemaLocator.getInfo().getPosition())
                        .thenApply(entry -> new StoredSchema(entry.getSchemaData().toByteArray(),
                                new LongSchemaVersion(schemaLocator.getInfo().getVersion())));
            });
        }).whenComplete((res, ex) -> {
            if (log.isDebugEnabled()) {
                log.debug("[{}] Get operation completed. res={} -- ex={}", schemaId, res, ex);
            }
            readSchemaOperations.remove(schemaId);
        });
    }

    @Override
    public SchemaVersion versionFromBytes(byte[] version) {
        // The schema storage converts the schema from bytes to long
        // so it handles both cases 1) version is 64 bytes long pre 2.4.0;
        // 2) version is 8 bytes long post 2.4.0
        //
        // NOTE: if you are planning to change the logic here. you should consider
        //       both 64 bytes and 8 bytes cases.
        ByteBuffer bb = ByteBuffer.wrap(version);
        return new LongSchemaVersion(bb.getLong());
    }

    @Override
    public void close() throws Exception {
        if (bookKeeper != null) {
            bookKeeper.close();
        }
    }

    @NonNull
    private CompletableFuture<StoredSchema> getSchema(String schemaId, long version) {
        if (log.isDebugEnabled()) {
            log.debug("[{}] Get schema - version: {}", schemaId, version);
        }

        return getSchemaLocator(getSchemaPath(schemaId)).thenCompose(locator -> {
            if (log.isDebugEnabled()) {
                log.debug("[{}] Get schema - version: {} - locator: {}", schemaId, version, locator);
            }

            if (!locator.isPresent()) {
                return completedFuture(null);
            }

            SchemaStorageFormat.SchemaLocator schemaLocator = locator.get().locator;
            if (version > schemaLocator.getInfo().getVersion()) {
                return completedFuture(null);
            }

            return findSchemaEntryByVersion(schemaLocator.getIndexList(), version)
                .thenApply(entry ->
                        new StoredSchema(
                            entry.getSchemaData().toByteArray(),
                            new LongSchemaVersion(version)
                        )
                );
        });
    }

    @NonNull
    private CompletableFuture<Long> putSchema(String schemaId, byte[] data, byte[] hash) {
        return getSchemaLocator(getSchemaPath(schemaId)).thenCompose(optLocatorEntry ->
                putSchema(schemaId, data, hash, optLocatorEntry));
    }

    private CompletableFuture<Long> putSchema(String schemaId, byte[] data, byte[] hash,
                                              Optional<LocatorEntry> optLocatorEntry) {
        if (optLocatorEntry.isPresent()) {

            SchemaStorageFormat.SchemaLocator locator = optLocatorEntry.get().locator;

            if (log.isDebugEnabled()) {
                log.debug("[{}] findSchemaEntryByHash - hash={}", schemaId, hash);
            }

            //don't check the schema whether already exist
            return readSchemaEntry(locator.getIndexList().get(0).getPosition())
                    .thenCompose(schemaEntry -> addNewSchemaEntryToStore(schemaId,
                            locator.getIndexList(), data).thenCompose(
                            position -> updateSchemaLocator(schemaId, optLocatorEntry.get(), position, hash))
                    );
        } else {
            return createNewSchema(schemaId, data, hash);
        }
    }

    private CompletableFuture<Long> createNewSchema(String schemaId, byte[] data, byte[] hash) {
        SchemaStorageFormat.IndexEntry emptyIndex = SchemaStorageFormat.IndexEntry.newBuilder()
                        .setVersion(0)
                        .setHash(copyFrom(hash))
                        .setPosition(SchemaStorageFormat.PositionInfo.newBuilder()
                                .setEntryId(-1L)
                                .setLedgerId(-1L)
                        ).build();

        return addNewSchemaEntryToStore(schemaId, Collections.singletonList(emptyIndex), data).thenCompose(position -> {
            // The schema was stored in the ledger, now update the z-node with the pointer to it
            SchemaStorageFormat.IndexEntry info = SchemaStorageFormat.IndexEntry.newBuilder()
                    .setVersion(0)
                    .setPosition(position)
                    .setHash(copyFrom(hash))
                    .build();

            return createSchemaLocator(getSchemaPath(schemaId), SchemaStorageFormat.SchemaLocator.newBuilder()
                    .setInfo(info)
                    .addAllIndex(
                            newArrayList(info))
                    .build())
                            .thenApply(ignore -> 0L);
        });
    }

    @NonNull
    private CompletableFuture<Long> deleteSchema(String schemaId, boolean forcefully) {
        return (forcefully ? CompletableFuture.completedFuture(null)
                : ignoreUnrecoverableBKException(getSchema(schemaId))).thenCompose(schemaAndVersion -> {
            if (!forcefully && isNull(schemaAndVersion)) {
                return completedFuture(null);
            } else {
                // The version is only for the compatibility of the current interface
                final long version = -1;
                CompletableFuture<Long> future = new CompletableFuture<>();
                getLocator(schemaId).whenComplete((locator, ex) -> {
                    if (ex != null) {
                        future.completeExceptionally(ex);
                    } else {
                        if (!locator.isPresent()) {
                            future.complete(null);
                            return;
                        }
                        List<SchemaStorageFormat.IndexEntry> indexEntryList = locator.get().locator.getIndexList();
                        List<CompletableFuture<Void>> deleteFutures = new ArrayList<>(indexEntryList.size());
                        indexEntryList.forEach(indexEntry -> {
                            final long ledgerId = indexEntry.getPosition().getLedgerId();
                            CompletableFuture<Void> deleteFuture = new CompletableFuture<>();
                            deleteFutures.add(deleteFuture);
                            bookKeeper.asyncDeleteLedger(ledgerId, (int rc, Object cnx) -> {
                                if (rc != BKException.Code.OK) {
                                    // It's not a serious error, we didn't need call future.completeExceptionally()
                                    log.warn("Failed to delete ledger {} of {}: {}", ledgerId, schemaId, rc);
                                }
                                deleteFuture.complete(null);
                            }, null);
                        });
                        FutureUtil.waitForAll(deleteFutures).whenComplete((v, e) -> {
                            final String path = getSchemaPath(schemaId);
                            store.delete(path, Optional.empty())
                                    .thenRun(() -> {
                                        future.complete(version);
                                    }).exceptionally(zkException -> {
                                        if (zkException.getCause()
                                                instanceof MetadataStoreException.NotFoundException) {
                                            // The znode has been deleted by others.
                                            // In some cases, the program may enter this logic.
                                            // Since the znode is gone, we don’t need to deal with it.
                                            if (log.isDebugEnabled()) {
                                                log.debug("No node for schema path: {}", path);
                                            }
                                            future.complete(null);
                                        } else {
                                            future.completeExceptionally(zkException);
                                        }
                                        return null;
                            });
                        });
                    }
                });
                return future;
            }
        });
    }

    @NonNull
    private static String getSchemaPath(String schemaId) {
        return SchemaPath + "/" + schemaId;
    }

    @NonNull
    private CompletableFuture<SchemaStorageFormat.PositionInfo> addNewSchemaEntryToStore(
        String schemaId,
        List<SchemaStorageFormat.IndexEntry> index,
        byte[] data
    ) {
        SchemaStorageFormat.SchemaEntry schemaEntry = newSchemaEntry(index, data);
        return createLedger(schemaId).thenCompose(ledgerHandle -> {
            final long ledgerId = ledgerHandle.getId();
            return addEntry(ledgerHandle, schemaEntry)
                    .thenApply(entryId -> {
                        ledgerHandle.closeAsync();
                        return Functions.newPositionInfo(ledgerId, entryId);
                    });
        });
    }

    @NonNull
    private CompletableFuture<Long> updateSchemaLocator(
        String schemaId,
        LocatorEntry locatorEntry,
        SchemaStorageFormat.PositionInfo position,
        byte[] hash
    ) {
        long nextVersion = locatorEntry.locator.getInfo().getVersion() + 1;
        SchemaStorageFormat.SchemaLocator locator = locatorEntry.locator;
        SchemaStorageFormat.IndexEntry info =
            SchemaStorageFormat.IndexEntry.newBuilder()
                .setVersion(nextVersion)
                .setPosition(position)
                .setHash(copyFrom(hash))
                .build();

        final ArrayList<SchemaStorageFormat.IndexEntry> indexList = new ArrayList<>();
        indexList.addAll(locator.getIndexList());
        indexList.add(info);
        return updateSchemaLocator(getSchemaPath(schemaId),
                SchemaStorageFormat.SchemaLocator.newBuilder()
                        .setInfo(info)
                        .addAllIndex(indexList)
                        .build()
                , locatorEntry.version
        ).thenApply(ignore -> nextVersion).whenComplete((__, ex) -> {
            if (ex != null) {
                Throwable cause = FutureUtil.unwrapCompletionException(ex);
                log.warn("[{}] Failed to update schema locator with position {}", schemaId, position, cause);
                if (cause instanceof AlreadyExistsException || cause instanceof BadVersionException) {
                    bookKeeper.asyncDeleteLedger(position.getLedgerId(), new AsyncCallback.DeleteCallback() {
                        @Override
                        public void deleteComplete(int rc, Object ctx) {
                            if (rc != BKException.Code.OK) {
                                log.warn("[{}] Failed to delete ledger {} after updating schema locator failed, rc: {}",
                                    schemaId, position.getLedgerId(), rc);
                            }
                        }
                    }, null);
                }
            }
        });
    }

    @NonNull
    private CompletableFuture<SchemaStorageFormat.SchemaEntry> findSchemaEntryByVersion(
        List<SchemaStorageFormat.IndexEntry> index,
        long version
    ) {

        if (index.isEmpty()) {
            return completedFuture(null);
        }

        SchemaStorageFormat.IndexEntry lowest = index.get(0);
        if (version < lowest.getVersion()) {
            return readSchemaEntry(lowest.getPosition())
                    .thenCompose(entry -> findSchemaEntryByVersion(entry.getIndexList(), version));
        }

        for (SchemaStorageFormat.IndexEntry entry : index) {
            if (entry.getVersion() == version) {
                return readSchemaEntry(entry.getPosition());
            } else if (entry.getVersion() > version) {
                break;
            }
        }

        return completedFuture(null);
    }

    @NonNull
    private CompletableFuture<SchemaStorageFormat.SchemaEntry> readSchemaEntry(
        SchemaStorageFormat.PositionInfo position
    ) {
        if (log.isDebugEnabled()) {
            log.debug("Reading schema entry from {}", position);
        }

        return openLedger(position.getLedgerId())
            .thenCompose((ledger) ->
                Functions.getLedgerEntry(ledger, position.getEntryId(), config.isSchemaLedgerForceRecovery())
                    .thenCompose(entry -> closeLedger(ledger)
                        .thenApply(ignore -> entry)
                    )
            ).thenCompose(Functions::parseSchemaEntry);
    }

    @NonNull
    private CompletableFuture<Void> updateSchemaLocator(String id,
                                                        SchemaStorageFormat.SchemaLocator schema, long version) {
        return store.put(id, schema.toByteArray(), Optional.of(version)).thenApply(__ -> null);
    }

    @NonNull
    private CompletableFuture<LocatorEntry> createSchemaLocator(String id, SchemaStorageFormat.SchemaLocator locator) {
        return store.put(id, locator.toByteArray(), Optional.of(-1L))
                .thenApply(stat -> new LocatorEntry(locator, stat.getVersion()));
    }

    @NonNull
    private CompletableFuture<Optional<LocatorEntry>> getSchemaLocator(String schema) {
        return locatorEntryCache.getWithStats(schema)
                .thenApply(o ->
                        o.map(r -> new LocatorEntry(r.getValue(), r.getStat().getVersion())));
    }

    public CompletableFuture<SchemaMetadata> getSchemaMetadata(String schema) {
        return getLocator(schema).thenApply(locator -> {
            if (!locator.isPresent()) {
                return null;
            }
            SchemaLocator sl = locator.get().locator;
            SchemaMetadata metadata = new SchemaMetadata();
            IndexEntry info = sl.getInfo();
            metadata.info = new SchemaMetadata.Entry(info.getPosition().getLedgerId(), info.getPosition().getEntryId(),
                    info.getVersion());
            metadata.index = sl.getIndexList() == null ? null
                    : sl.getIndexList().stream().map(i -> new SchemaMetadata.Entry(i.getPosition().getLedgerId(),
                            i.getPosition().getEntryId(), i.getVersion())).collect(Collectors.toList());
            return metadata;
        });
    }

    @NonNull
    private CompletableFuture<Long> addEntry(LedgerHandle ledgerHandle, SchemaStorageFormat.SchemaEntry entry) {
        final CompletableFuture<Long> future = new CompletableFuture<>();
        ledgerHandle.asyncAddEntry(entry.toByteArray(),
            (rc, handle, entryId, ctx) -> {
                if (rc != BKException.Code.OK) {
                    future.completeExceptionally(bkException("Failed to add entry", rc, ledgerHandle.getId(), -1,
                            config.isSchemaLedgerForceRecovery()));
                } else {
                    future.complete(entryId);
                }
            }, null
        );
        return future;
    }

    @NonNull
    private CompletableFuture<LedgerHandle> createLedger(String schemaId) {
        Map<String, byte[]> metadata = LedgerMetadataUtils.buildMetadataForSchema(schemaId);
        final CompletableFuture<LedgerHandle> future = new CompletableFuture<>();
        try {
            bookKeeper.asyncCreateLedger(
                    config.getManagedLedgerDefaultEnsembleSize(),
                    config.getManagedLedgerDefaultWriteQuorum(),
                    config.getManagedLedgerDefaultAckQuorum(),
                    BookKeeper.DigestType.fromApiDigestType(config.getManagedLedgerDigestType()),
                    LedgerPassword,
                    (rc, handle, ctx) -> {
                        if (rc != BKException.Code.OK) {
                            future.completeExceptionally(bkException("Failed to create ledger", rc, -1, -1,
                                    config.isSchemaLedgerForceRecovery()));
                        } else {
                            future.complete(handle);
                        }
                    }, null, metadata);
        } catch (Throwable t) {
            log.error("[{}] Encountered unexpected error when creating schema ledger", schemaId, t);
            return FutureUtil.failedFuture(t);
        }
        return future;
    }

    @NonNull
    private CompletableFuture<LedgerHandle> openLedger(Long ledgerId) {
        final CompletableFuture<LedgerHandle> future = new CompletableFuture<>();
        bookKeeper.asyncOpenLedger(
            ledgerId,
            BookKeeper.DigestType.fromApiDigestType(config.getManagedLedgerDigestType()),
            LedgerPassword,
            (rc, handle, ctx) -> {
                if (rc != BKException.Code.OK) {
                    future.completeExceptionally(bkException("Failed to open ledger", rc, ledgerId, -1,
                            config.isSchemaLedgerForceRecovery()));
                } else {
                    future.complete(handle);
                }
            }, null
        );
        return future;
    }

    @NonNull
    private CompletableFuture<Void> closeLedger(LedgerHandle ledgerHandle) {
        CompletableFuture<Void> future = new CompletableFuture<>();
        ledgerHandle.asyncClose((rc, handle, ctx) -> {
            if (rc != BKException.Code.OK) {
                future.completeExceptionally(bkException("Failed to close ledger", rc, ledgerHandle.getId(), -1,
                        config.isSchemaLedgerForceRecovery()));
            } else {
                future.complete(null);
            }
        }, null);
        return future;
    }

    public CompletableFuture<List<Long>> getStoreLedgerIdsBySchemaId(String schemaId) {
        CompletableFuture<List<Long>> ledgerIdsFuture = new CompletableFuture<>();
        getSchemaLocator(getSchemaPath(schemaId)).thenAccept(locator -> {
            if (log.isDebugEnabled()) {
                log.debug("[{}] Get all store schema ledgerIds - locator: {}", schemaId, locator);
            }

            if (!locator.isPresent()) {
                ledgerIdsFuture.complete(Collections.emptyList());
                return;
            }
            Set<Long> ledgerIds = new HashSet<>();
            SchemaStorageFormat.SchemaLocator schemaLocator = locator.get().locator;
            schemaLocator.getIndexList().forEach(indexEntry -> ledgerIds.add(indexEntry.getPosition().getLedgerId()));
            ledgerIdsFuture.complete(new ArrayList<>(ledgerIds));
        }).exceptionally(e -> {
            ledgerIdsFuture.completeExceptionally(e);
            return null;
        });
        return ledgerIdsFuture;
    }

    interface Functions {
        static CompletableFuture<LedgerEntry> getLedgerEntry(LedgerHandle ledger, long entry,
                boolean forceRecovery) {
            final CompletableFuture<LedgerEntry> future = new CompletableFuture<>();
            ledger.asyncReadEntries(entry, entry,
                (rc, handle, entries, ctx) -> {
                    if (rc != BKException.Code.OK) {
                        future.completeExceptionally(bkException("Failed to read entry", rc, ledger.getId(), entry,
                                forceRecovery));
                    } else {
                        future.complete(entries.nextElement());
                    }
                }, null
            );
            return future;
        }

        static CompletableFuture<SchemaStorageFormat.SchemaEntry> parseSchemaEntry(LedgerEntry ledgerEntry) {
            CompletableFuture<SchemaStorageFormat.SchemaEntry> result = new CompletableFuture<>();
            try {
                result.complete(SchemaStorageFormat.SchemaEntry.parseFrom(ledgerEntry.getEntry()));
            } catch (IOException e) {
                result.completeExceptionally(e);
            }
            return result;
        }

        static SchemaStorageFormat.SchemaEntry newSchemaEntry(
            List<SchemaStorageFormat.IndexEntry> index,
            byte[] data
        ) {
            return SchemaStorageFormat.SchemaEntry.newBuilder()
                .setSchemaData(copyFrom(data))
                .addAllIndex(index)
                .build();
        }

        static SchemaStorageFormat.PositionInfo newPositionInfo(long ledgerId, long entryId) {
            return SchemaStorageFormat.PositionInfo.newBuilder()
                .setLedgerId(ledgerId)
                .setEntryId(entryId)
                .build();
        }
    }

    static class LocatorEntry {
        final SchemaStorageFormat.SchemaLocator locator;
        final long version;

        LocatorEntry(SchemaStorageFormat.SchemaLocator locator, long version) {
            this.locator = locator;
            this.version = version;
        }
    }

    public static Exception bkException(String operation, int rc, long ledgerId, long entryId,
            boolean forceRecovery) {
        String message = org.apache.bookkeeper.client.api.BKException.getMessage(rc)
                + " -  ledger=" + ledgerId + " - operation=" + operation;

        if (entryId != -1) {
            message += " - entry=" + entryId;
        }
        boolean recoverable = rc != BKException.Code.NoSuchLedgerExistsException
                && rc != BKException.Code.NoSuchEntryException
                && rc != BKException.Code.NoSuchLedgerExistsOnMetadataServerException
                // if force-recovery is enabled then made it non-recoverable exception
                // and force schema to skip this exception and recover immediately
                && !forceRecovery;
        return new SchemaException(recoverable, message);
    }

    public static <T> CompletableFuture<T> ignoreUnrecoverableBKException(CompletableFuture<T> source) {
        return source.exceptionally(t -> {
            if (t.getCause() != null
                    && (t.getCause() instanceof SchemaException)
                    && !(t.getCause() instanceof IncompatibleSchemaException)
                    && !((SchemaException) t.getCause()).isRecoverable()) {
                // Meeting NoSuchLedgerExistsException, NoSuchEntryException or
                // NoSuchLedgerExistsOnMetadataServerException when reading schemas in
                // bookkeeper. This also means that the data has already been deleted by other operations
                // in deleting schema.
                if (log.isDebugEnabled()) {
                    log.debug("Schema data in bookkeeper may be deleted by other operations.", t);
                }
                return null;
            }
            // rethrow other cases
            throw t instanceof CompletionException ? (CompletionException) t : new CompletionException(t);
        });
    }
}