/*
 * 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.bookkeeper.metadata.etcd;

import static java.nio.charset.StandardCharsets.UTF_8;
import static org.apache.bookkeeper.metadata.etcd.EtcdUtils.getBookiesEndPath;
import static org.apache.bookkeeper.metadata.etcd.EtcdUtils.getBookiesPath;
import static org.apache.bookkeeper.metadata.etcd.EtcdUtils.getBucketsPath;
import static org.apache.bookkeeper.metadata.etcd.EtcdUtils.getClusterInstanceIdPath;
import static org.apache.bookkeeper.metadata.etcd.EtcdUtils.getCookiePath;
import static org.apache.bookkeeper.metadata.etcd.EtcdUtils.getCookiesPath;
import static org.apache.bookkeeper.metadata.etcd.EtcdUtils.getLayoutKey;
import static org.apache.bookkeeper.metadata.etcd.EtcdUtils.getLedgersPath;
import static org.apache.bookkeeper.metadata.etcd.EtcdUtils.getReadonlyBookiePath;
import static org.apache.bookkeeper.metadata.etcd.EtcdUtils.getReadonlyBookiesPath;
import static org.apache.bookkeeper.metadata.etcd.EtcdUtils.getScopeEndKey;
import static org.apache.bookkeeper.metadata.etcd.EtcdUtils.getUnderreplicationPath;
import static org.apache.bookkeeper.metadata.etcd.EtcdUtils.getWritableBookiePath;
import static org.apache.bookkeeper.metadata.etcd.EtcdUtils.getWritableBookiesPath;
import static org.apache.bookkeeper.metadata.etcd.EtcdUtils.msResult;

import com.google.common.annotations.VisibleForTesting;
import com.google.common.util.concurrent.UncheckedExecutionException;
import io.etcd.jetcd.ByteSequence;
import io.etcd.jetcd.Client;
import io.etcd.jetcd.KV;
import io.etcd.jetcd.KeyValue;
import io.etcd.jetcd.Txn;
import io.etcd.jetcd.Watch;
import io.etcd.jetcd.Watch.Watcher;
import io.etcd.jetcd.kv.DeleteResponse;
import io.etcd.jetcd.kv.GetResponse;
import io.etcd.jetcd.kv.TxnResponse;
import io.etcd.jetcd.op.Cmp;
import io.etcd.jetcd.op.Cmp.Op;
import io.etcd.jetcd.op.CmpTarget;
import io.etcd.jetcd.options.DeleteOption;
import io.etcd.jetcd.options.GetOption;
import io.etcd.jetcd.options.PutOption;
import io.etcd.jetcd.options.WatchOption;
import io.etcd.jetcd.watch.WatchEvent;
import io.etcd.jetcd.watch.WatchEvent.EventType;
import java.util.UUID;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException;
import lombok.AccessLevel;
import lombok.Getter;
import lombok.extern.slf4j.Slf4j;
import org.apache.bookkeeper.bookie.BookieException;
import org.apache.bookkeeper.bookie.BookieException.BookieIllegalOpException;
import org.apache.bookkeeper.bookie.BookieException.CookieNotFoundException;
import org.apache.bookkeeper.bookie.BookieException.MetadataStoreException;
import org.apache.bookkeeper.discover.BookieServiceInfo;
import org.apache.bookkeeper.discover.RegistrationManager;
import org.apache.bookkeeper.meta.LedgerLayout;
import org.apache.bookkeeper.net.BookieId;
import org.apache.bookkeeper.versioning.LongVersion;
import org.apache.bookkeeper.versioning.Version;
import org.apache.bookkeeper.versioning.Versioned;

/**
 * Etcd registration manager.
 */
@Slf4j
class EtcdRegistrationManager implements RegistrationManager {

    private final String scope;
    @Getter(AccessLevel.PACKAGE)
    private final Client client;
    private final boolean ownClient;
    private final KV kvClient;
    @Getter(AccessLevel.PACKAGE)
    private final EtcdBookieRegister bkRegister;

    @VisibleForTesting
    EtcdRegistrationManager(Client client,
                            String scope) {
        this(client, scope, 60);
    }

    @VisibleForTesting
    EtcdRegistrationManager(Client client,
                            String scope,
                            long ttlSeconds) {
        this(client, scope, ttlSeconds, () -> {});
    }

    @VisibleForTesting
    EtcdRegistrationManager(Client client,
                            String scope,
                            long ttlSeconds,
                            RegistrationListener listener) {
        this(
            client,
            scope,
            new EtcdBookieRegister(
                client.getLeaseClient(),
                ttlSeconds
            ).addRegistrationListener(listener).start(),
            true);
    }

    EtcdRegistrationManager(Client client,
                            String scope,
                            EtcdBookieRegister bkRegister) {
        this(client, scope, bkRegister, false);
    }

    private EtcdRegistrationManager(Client client,
                                    String scope,
                                    EtcdBookieRegister bkRegister,
                                    boolean ownClient) {
        this.scope = scope;
        this.client = client;
        this.kvClient = client.getKVClient();
        this.bkRegister = bkRegister;
        this.ownClient = ownClient;
    }

    @Override
    public void close() {
        if (ownClient) {
            log.info("Closing registration manager under scope '{}'", scope);
            bkRegister.close();
            client.close();
            log.info("Successfully closed registration manager under scope '{}'", scope);
        }
    }

    @Override
    public void registerBookie(BookieId bookieId, boolean readOnly,
                               BookieServiceInfo bookieServiceInfo) throws BookieException {
        if (readOnly) {
            doRegisterReadonlyBookie(bookieId, bkRegister.get());
        } else {
            doRegisterBookie(getWritableBookiePath(scope, bookieId), bkRegister.get());
        }
    }

    private boolean checkRegNodeAndWaitExpired(String regPath, long leaseId)
            throws MetadataStoreException {
        ByteSequence regPathBs = ByteSequence.from(regPath, UTF_8);
        GetResponse getResp = msResult(kvClient.get(regPathBs));
        if (getResp.getCount() <= 0) {
            // key doesn't exist anymore
            return false;
        } else {
            return waitUntilRegNodeExpired(regPath, leaseId);
        }
    }

    private boolean waitUntilRegNodeExpired(String regPath, long leaseId)
            throws MetadataStoreException {
        ByteSequence regPathBs = ByteSequence.from(regPath, UTF_8);
        // check regPath again
        GetResponse getResp = msResult(kvClient.get(regPathBs));
        if (getResp.getCount() <= 0) {
            // key disappears after watching it
            return false;
        } else {
            KeyValue kv = getResp.getKvs().get(0);
            if (kv.getLease() != leaseId) {
                Watch watchClient = client.getWatchClient();
                final CompletableFuture<Void> watchFuture = new CompletableFuture<>();
                Watcher watcher = watchClient.watch(
                    regPathBs,
                    WatchOption.newBuilder()
                        .withRevision(getResp.getHeader().getRevision() + 1)
                        .build(),
                    response -> {
                        for (WatchEvent event : response.getEvents()) {
                            log.info("Received watch event on '{}' : EventType = {}, lease {}",
                                regPath, event.getEventType(), leaseId);
                            if (EventType.DELETE == event.getEventType()) {
                                watchFuture.complete(null);
                                return;
                            }
                        }
                    },
                    exception -> {
                        log.warn("Exception in keepAlive for watch event on {}, lease {}",
                                regPath, leaseId, exception);
                        watchFuture.completeExceptionally(new UncheckedExecutionException(
                                    "Interrupted at waiting previous registration under "
                                    + regPath + " (lease = " + kv.getLease() + ") to be expired", exception));
                    }
                );
                log.info("Previous bookie registration (lease = {}) still exists at {}, "
                                + "so new lease '{}' will be waiting previous lease for {} seconds to be expired",
                        kv.getLease(), regPath, leaseId, bkRegister.getTtlSeconds());

                try {
                    msResult(watchFuture, 2 * bkRegister.getTtlSeconds(), TimeUnit.SECONDS);
                    return false;
                } catch (TimeoutException e) {
                    watchFuture.cancel(true);
                    throw new MetadataStoreException(
                        "Previous bookie registration still exists at "
                            + regPath + " (lease = " + kv.getLease() + ") after "
                            + (2 * bkRegister.getTtlSeconds()) + " seconds elapsed");
                } catch (UncheckedExecutionException uee) {
                    throw new MetadataStoreException(uee.getMessage(), uee.getCause());
                } finally {
                    watcher.close();
                }
            } else {
                // key exists with same lease
                return true;
            }
        }
    }

    private void doRegisterBookie(String regPath, long leaseId) throws MetadataStoreException {
        if (checkRegNodeAndWaitExpired(regPath, leaseId)) {
            // the bookie is already registered under `${regPath}` with `${leaseId}`.
            return;
        }

        ByteSequence regPathBs = ByteSequence.from(regPath, UTF_8);
        Txn txn = kvClient.txn()
            .If(new Cmp(
                regPathBs,
                Op.GREATER,
                CmpTarget.createRevision(0)))
            .Then(io.etcd.jetcd.op.Op.get(regPathBs, GetOption.DEFAULT))
            .Else(io.etcd.jetcd.op.Op.put(
                regPathBs,
                ByteSequence.from(new byte[0]),
                PutOption.newBuilder()
                    .withLeaseId(bkRegister.get())
                    .build()
            ));
        TxnResponse txnResp = msResult(txn.commit());
        if (txnResp.isSucceeded()) {
            // the key already exists
            GetResponse getResp = txnResp.getGetResponses().get(0);
            if (getResp.getCount() <= 0) {
                throw new MetadataStoreException(
                    "Failed to register bookie under '" + regPath
                        + "', but no bookie is registered there.");
            } else {
                KeyValue kv = getResp.getKvs().get(0);
                throw new MetadataStoreException("Another bookie already registered under '"
                    + regPath + "': lease = " + kv.getLease());
            }
        } else {
            log.info("Successfully registered bookie at {}", regPath);
        }
    }

    private void doRegisterReadonlyBookie(BookieId bookieId, long leaseId) throws MetadataStoreException {
        String readonlyRegPath = getReadonlyBookiePath(scope, bookieId);
        doRegisterBookie(readonlyRegPath, leaseId);
        String writableRegPath = getWritableBookiePath(scope, bookieId);
        msResult(kvClient.delete(ByteSequence.from(writableRegPath, UTF_8)));
    }

    @Override
    public void unregisterBookie(BookieId bookieId, boolean readOnly) throws BookieException {
        String regPath;
        if (readOnly) {
            regPath = getReadonlyBookiePath(scope, bookieId);
        } else {
            regPath = getWritableBookiePath(scope, bookieId);
        }
        DeleteResponse delResp = msResult(kvClient.delete(ByteSequence.from(regPath, UTF_8)));
        if (delResp.getDeleted() > 0) {
            log.info("Successfully unregistered bookie {} from {}", bookieId, regPath);
        } else {
            log.info("Bookie disappeared from {} before unregistering", regPath);
        }
    }

    @Override
    public boolean isBookieRegistered(BookieId bookieId) throws BookieException {
        CompletableFuture<GetResponse> getWritableFuture = kvClient.get(
            ByteSequence.from(getWritableBookiePath(scope, bookieId), UTF_8),
            GetOption.newBuilder()
                .withCountOnly(true)
                .build());
        CompletableFuture<GetResponse> getReadonlyFuture = kvClient.get(
            ByteSequence.from(getReadonlyBookiePath(scope, bookieId), UTF_8),
            GetOption.newBuilder()
                .withCountOnly(true)
                .build());

        return msResult(getWritableFuture).getCount() > 0
            || msResult(getReadonlyFuture).getCount() > 0;
    }

    @Override
    public void writeCookie(BookieId bookieId, Versioned<byte[]> cookieData) throws BookieException {
        ByteSequence cookiePath = ByteSequence.from(getCookiePath(scope, bookieId), UTF_8);
        Txn txn = kvClient.txn();
        if (Version.NEW == cookieData.getVersion()) {
            txn.If(new Cmp(
                cookiePath,
                Op.GREATER,
                CmpTarget.createRevision(0L))
            )
            // if key not exists, create one.
            .Else(io.etcd.jetcd.op.Op.put(
                cookiePath,
                ByteSequence.from(cookieData.getValue()),
                PutOption.DEFAULT)
            );
        } else {
            if (!(cookieData.getVersion() instanceof LongVersion)) {
                throw new BookieIllegalOpException("Invalid version type, expected it to be LongVersion");
            }
            txn.If(new Cmp(
                cookiePath,
                Op.EQUAL,
                CmpTarget.modRevision(((LongVersion) cookieData.getVersion()).getLongVersion()))
            )
            .Then(io.etcd.jetcd.op.Op.put(
                cookiePath,
                ByteSequence.from(cookieData.getValue()),
                PutOption.DEFAULT)
            );
        }
        TxnResponse response = msResult(txn.commit());
        if (response.isSucceeded() != (Version.NEW != cookieData.getVersion())) {
            throw new MetadataStoreException(
                "Conflict on writing cookie for bookie " + bookieId);
        }
    }

    @Override
    public Versioned<byte[]> readCookie(BookieId bookieId) throws BookieException {
        ByteSequence cookiePath = ByteSequence.from(getCookiePath(scope, bookieId), UTF_8);
        GetResponse resp = msResult(kvClient.get(cookiePath));
        if (resp.getCount() <= 0) {
            throw new CookieNotFoundException(bookieId.toString());
        } else {
            KeyValue kv = resp.getKvs().get(0);
            return new Versioned<>(
                kv.getValue().getBytes(),
                new LongVersion(kv.getModRevision()));
        }
    }

    @Override
    public void removeCookie(BookieId bookieId, Version version) throws BookieException {
        ByteSequence cookiePath = ByteSequence.from(getCookiePath(scope, bookieId), UTF_8);
        Txn delTxn = kvClient.txn()
            .If(new Cmp(
                cookiePath,
                Op.EQUAL,
                CmpTarget.modRevision(((LongVersion) version).getLongVersion())
            ))
            .Then(io.etcd.jetcd.op.Op.delete(
                cookiePath,
                DeleteOption.DEFAULT
            ))
            .Else(io.etcd.jetcd.op.Op.get(
                cookiePath,
                GetOption.newBuilder().withCountOnly(true).build()
            ));
        TxnResponse txnResp = msResult(delTxn.commit());
        if (!txnResp.isSucceeded()) {
            GetResponse getResp = txnResp.getGetResponses().get(0);
            if (getResp.getCount() > 0) {
                throw new MetadataStoreException(
                    "Failed to remove cookie from " + cookiePath.toString(UTF_8)
                        + " for bookie " + bookieId + " : bad version '" + version + "'");
            } else {
                throw new CookieNotFoundException(bookieId.toString());
            }
        } else {
            log.info("Removed cookie from {} for bookie {}",
                cookiePath.toString(UTF_8), bookieId);
        }
    }

    @Override
    public String getClusterInstanceId() throws BookieException {
        GetResponse response = msResult(
            kvClient.get(ByteSequence.from(getClusterInstanceIdPath(scope), UTF_8)));
        if (response.getCount() <= 0) {
            log.error("BookKeeper metadata doesn't exist in Etcd. "
                + "Has the cluster been initialized? "
                + "Try running bin/bookkeeper shell initNewCluster");
            throw new MetadataStoreException("BookKeeper is not initialized under '" + scope + "' yet");
        } else {
            KeyValue kv = response.getKvs().get(0);
            return new String(kv.getValue().getBytes(), UTF_8);
        }
    }

    @Override
    public boolean prepareFormat() throws Exception {
        ByteSequence rootScopeKey = ByteSequence.from(scope, UTF_8);
        GetResponse resp = msResult(kvClient.get(rootScopeKey));
        return resp.getCount() > 0;
    }

    @Override
    public boolean initNewCluster() throws Exception {
        return initNewCluster(kvClient, scope);
    }

    static boolean initNewCluster(KV kvClient, String scope) throws Exception {
        ByteSequence rootScopeKey = ByteSequence.from(scope, UTF_8);
        String instanceId = UUID.randomUUID().toString();
        LedgerLayout layout = new LedgerLayout(
            EtcdLedgerManagerFactory.class.getName(),
            EtcdLedgerManagerFactory.VERSION
        );
        Txn initTxn = kvClient.txn()
            .If(new Cmp(
                rootScopeKey,
                Op.GREATER,
                CmpTarget.createRevision(0L)
            ))
            // only put keys when root scope doesn't exist
            .Else(
                // `${scope}`
                io.etcd.jetcd.op.Op.put(
                    rootScopeKey,
                    EtcdConstants.EMPTY_BS,
                    PutOption.DEFAULT
                ),
                // `${scope}/layout`
                io.etcd.jetcd.op.Op.put(
                    ByteSequence.from(getLayoutKey(scope), UTF_8),
                    ByteSequence.from(layout.serialize()),
                    PutOption.DEFAULT
                ),
                // `${scope}/instanceid`
                io.etcd.jetcd.op.Op.put(
                    ByteSequence.from(getClusterInstanceIdPath(scope), UTF_8),
                    ByteSequence.from(instanceId, UTF_8),
                    PutOption.DEFAULT
                ),
                // `${scope}/cookies`
                io.etcd.jetcd.op.Op.put(
                    ByteSequence.from(getCookiesPath(scope), UTF_8),
                    EtcdConstants.EMPTY_BS,
                    PutOption.DEFAULT
                ),
                // `${scope}/bookies`
                io.etcd.jetcd.op.Op.put(
                    ByteSequence.from(getBookiesPath(scope), UTF_8),
                    EtcdConstants.EMPTY_BS,
                    PutOption.DEFAULT
                ),
                // `${scope}/bookies/writable`
                io.etcd.jetcd.op.Op.put(
                    ByteSequence.from(getWritableBookiesPath(scope), UTF_8),
                    EtcdConstants.EMPTY_BS,
                    PutOption.DEFAULT
                ),
                // `${scope}/bookies/readonly`
                io.etcd.jetcd.op.Op.put(
                    ByteSequence.from(getReadonlyBookiesPath(scope), UTF_8),
                    EtcdConstants.EMPTY_BS,
                    PutOption.DEFAULT
                ),
                // `${scope}/ledgers`
                io.etcd.jetcd.op.Op.put(
                    ByteSequence.from(getLedgersPath(scope), UTF_8),
                    EtcdConstants.EMPTY_BS,
                    PutOption.DEFAULT
                ),
                // `${scope}/buckets`
                io.etcd.jetcd.op.Op.put(
                    ByteSequence.from(getBucketsPath(scope), UTF_8),
                    EtcdConstants.EMPTY_BS,
                    PutOption.DEFAULT
                ),
                // `${scope}/underreplication`
                io.etcd.jetcd.op.Op.put(
                    ByteSequence.from(getUnderreplicationPath(scope), UTF_8),
                    EtcdConstants.EMPTY_BS,
                    PutOption.DEFAULT
                )
            );

        return !msResult(initTxn.commit()).isSucceeded();
    }

    @Override
    public boolean format() throws Exception {
        return format(kvClient, scope);
    }

    static boolean format(KV kvClient, String scope) throws Exception {
        ByteSequence rootScopeKey = ByteSequence.from(scope, UTF_8);
        GetResponse resp = msResult(kvClient.get(rootScopeKey));
        if (resp.getCount() <= 0) {
            // cluster doesn't exist
            return initNewCluster(kvClient, scope);
        } else if (nukeExistingCluster(kvClient, scope)) { // cluster exists and has successfully nuked it
            return initNewCluster(kvClient, scope);
        } else {
            return false;
        }
    }

    @Override
    public boolean nukeExistingCluster() throws Exception {
        return nukeExistingCluster(kvClient, scope);
    }

    @Override
    public void addRegistrationListener(RegistrationListener listener) {
        bkRegister.addRegistrationListener(listener);
    }

    static boolean nukeExistingCluster(KV kvClient, String scope) throws Exception {
        ByteSequence rootScopeKey = ByteSequence.from(scope, UTF_8);
        GetResponse resp = msResult(kvClient.get(rootScopeKey));
        if (resp.getCount() <= 0) {
            log.info("There is no existing cluster with under scope '{}' in Etcd, "
                + "so exiting nuke operation", scope);
            return true;
        }

        String bookiesPath = getBookiesPath(scope);
        String bookiesEndPath = getBookiesEndPath(scope);
        resp = msResult(kvClient.get(
            ByteSequence.from(bookiesPath, UTF_8),
            GetOption.newBuilder()
                .withRange(ByteSequence.from(bookiesEndPath, UTF_8))
                .withKeysOnly(true)
                .build()
        ));
        String writableBookiesPath = getWritableBookiesPath(scope);
        String readonlyBookiesPath = getReadonlyBookiesPath(scope);
        boolean hasBookiesAlive = false;
        for (KeyValue kv : resp.getKvs()) {
            String keyStr = new String(kv.getKey().getBytes(), UTF_8);
            if (keyStr.equals(bookiesPath)
                || keyStr.equals(writableBookiesPath)
                || keyStr.equals(readonlyBookiesPath)) {
                continue;
            } else {
                hasBookiesAlive = true;
                break;
            }
        }
        if (hasBookiesAlive) {
            log.error("Bookies are still up and connected to this cluster, "
                + "stop all bookies before nuking the cluster");
            return false;
        }
        DeleteResponse delResp = msResult(kvClient.delete(
            rootScopeKey,
            DeleteOption.newBuilder()
                .withRange(ByteSequence.from(getScopeEndKey(scope), UTF_8))
                .build()));
        log.info("Successfully nuked cluster under scope '{}' : {} kv pairs deleted",
            scope, delResp.getDeleted());
        return true;
    }
}
