/*
 * Licensed 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.stream.cluster;

import static com.google.common.base.Preconditions.checkArgument;
import static com.google.common.base.Preconditions.checkNotNull;
import static org.apache.bookkeeper.common.concurrent.FutureUtils.result;
import static org.apache.bookkeeper.stream.protocol.ProtocolConstants.DEFAULT_STREAM_CONF;

import com.google.common.collect.Lists;
import java.io.File;
import java.io.IOException;
import java.net.BindException;
import java.util.List;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.Future;
import java.util.stream.Collectors;
import lombok.extern.slf4j.Slf4j;
import org.apache.bookkeeper.clients.StorageClientBuilder;
import org.apache.bookkeeper.clients.admin.StorageAdminClient;
import org.apache.bookkeeper.clients.config.StorageClientSettings;
import org.apache.bookkeeper.clients.exceptions.ClientException;
import org.apache.bookkeeper.clients.exceptions.NamespaceNotFoundException;
import org.apache.bookkeeper.clients.utils.NetUtils;
import org.apache.bookkeeper.common.component.AbstractLifecycleComponent;
import org.apache.bookkeeper.common.component.LifecycleComponent;
import org.apache.bookkeeper.common.net.ServiceURI;
import org.apache.bookkeeper.conf.ServerConfiguration;
import org.apache.bookkeeper.discover.RegistrationManager;
import org.apache.bookkeeper.meta.MetadataDrivers;
import org.apache.bookkeeper.shims.zk.ZooKeeperServerShim;
import org.apache.bookkeeper.stats.NullStatsLogger;
import org.apache.bookkeeper.stream.proto.NamespaceConfiguration;
import org.apache.bookkeeper.stream.proto.NamespaceProperties;
import org.apache.bookkeeper.stream.proto.common.Endpoint;
import org.apache.bookkeeper.stream.server.StorageServer;
import org.apache.bookkeeper.stream.storage.conf.StorageConfiguration;
import org.apache.bookkeeper.stream.storage.exceptions.StorageRuntimeException;
import org.apache.bookkeeper.stream.storage.impl.cluster.ZkClusterInitializer;
import org.apache.commons.lang3.StringUtils;
import org.apache.commons.lang3.tuple.Pair;
import org.apache.distributedlog.LocalDLMEmulator;

/**
 * A Cluster that runs a few storage nodes.
 */
@Slf4j
public class StreamCluster
    extends AbstractLifecycleComponent<StorageConfiguration> {

    private static final int MAX_RETRIES = 20;

    /**
     * Build a stream cluster from the provided cluster {@code spec}.
     *
     * @param spec the spec to build stream cluster.
     * @return stream cluster spec.
     */
    public static StreamCluster build(StreamClusterSpec spec) {
        return new StreamCluster(spec);
    }

    private static ServerConfiguration newBookieConfiguration(ServiceURI serviceURI) {
        ServerConfiguration serverConf = new ServerConfiguration();
        serverConf.setMetadataServiceUri(serviceURI.getUri().toString());
        serverConf.setAllowLoopback(true);
        serverConf.setGcWaitTime(300000);
        serverConf.setDiskUsageWarnThreshold(0.9999f);
        serverConf.setDiskUsageThreshold(0.999999f);
        return serverConf;
    }

    private final StreamClusterSpec spec;
    private final List<Endpoint> rpcEndpoints;
    private ServiceURI metadataServiceUri;
    private int zkPort;
    private ZooKeeperServerShim zks;
    private List<LifecycleComponent> servers;
    private int nextBookiePort;
    private int nextGrpcPort;

    private StreamCluster(StreamClusterSpec spec) {
        super(
            "stream-cluster",
            new StorageConfiguration(spec.baseConf()),
            NullStatsLogger.INSTANCE);
        this.spec = spec;
        this.servers = Lists.newArrayListWithExpectedSize(spec.numServers());
        this.rpcEndpoints = Lists.newArrayListWithExpectedSize(spec.numServers());
        this.nextBookiePort = spec.initialBookiePort();
        this.nextGrpcPort = spec.initialGrpcPort();
    }

    public List<Endpoint> getRpcEndpoints() {
        return rpcEndpoints;
    }

    private void startZooKeeper() throws Exception {
        if (!spec.shouldStartZooKeeper()) {
            metadataServiceUri = checkNotNull(spec.metadataServiceUri,
                "No metadata service uri is configured while configuring not to start zookeeper");
            return;
        }

        File zkDir = new File(spec.storageRootDir(), "zookeeper");
        Pair<ZooKeeperServerShim, Integer> zkServerAndPort =
            LocalDLMEmulator.runZookeeperOnAnyPort(spec.zkPort(), zkDir);
        zks = zkServerAndPort.getLeft();
        zkPort = zkServerAndPort.getRight();
        log.info("Started zookeeper at port {}.", zkPort);
        metadataServiceUri = ServiceURI.create(
            "zk://127.0.0.1:" + zkPort + "/ledgers");
    }

    private void stopZooKeeper() {
        // stop the zookeeper server
        if (null != zks) {
            zks.stop();
        }
    }

    private void initializeCluster() throws Exception {
        checkArgument(ServiceURI.SERVICE_ZK.equals(metadataServiceUri.getServiceName()),
            "Only support zookeeper based metadata service now");
        String[] serviceHosts = metadataServiceUri.getServiceHosts();
        String metadataServers = StringUtils.join(serviceHosts, ',');

        new ZkClusterInitializer(metadataServers).initializeCluster(
            metadataServiceUri.getUri(),
            spec.numServers() * 2);

        // format the bookkeeper cluster
        MetadataDrivers.runFunctionWithMetadataBookieDriver(newBookieConfiguration(metadataServiceUri), driver -> {
                try (RegistrationManager rm = driver.createRegistrationManager()) {
                    boolean initialized = rm.initNewCluster();
                    if (initialized) {
                        log.info("Successfully initialized the segment storage");
                    } else {
                        log.info("The segment storage was already initialized");
                    }
                } catch (Exception e) {
                    throw new StorageRuntimeException("Failed to initialize the segment storage", e);
                }
                return null;
            });
    }

    private LifecycleComponent startServer() throws Exception {
        int bookiePort;
        int grpcPort;
        boolean success = false;
        int retries = 0;

        while (!success) {
            synchronized (this) {
                bookiePort = nextBookiePort++;
                grpcPort = nextGrpcPort++;
            }
            LifecycleComponent server = null;
            try {
                ServerConfiguration serverConf = newBookieConfiguration(metadataServiceUri);
                serverConf.loadConf(spec.baseConf());
                serverConf.setBookiePort(bookiePort);
                File bkDir = new File(spec.storageRootDir(), "bookie_" + bookiePort);
                serverConf.setJournalDirName(bkDir.getPath());
                serverConf.setLedgerDirNames(new String[]{bkDir.getPath()});

                File rangesStoreDir = new File(spec.storageRootDir(), "ranges_" + grpcPort);
                StorageConfiguration storageConf = new StorageConfiguration(serverConf);
                storageConf.setRangeStoreDirNames(new String[]{rangesStoreDir.getPath()});

                log.info("Attempting to start storage server at (bookie port = {}, grpc port = {})"
                        + " : bkDir = {}, rangesStoreDir = {}",
                    bookiePort, grpcPort, bkDir, rangesStoreDir);
                server = StorageServer.buildStorageServer(
                    serverConf,
                    grpcPort);
                server.start();
                log.info("Started storage server at (bookie port = {}, grpc port = {})",
                    bookiePort, grpcPort);
                this.rpcEndpoints.add(StorageServer.createLocalEndpoint(grpcPort, false));
                return server;
            } catch (Throwable e) {
                log.error("Failed to start storage server", e);
                if (null != server) {
                    server.stop();
                }
                if (e.getCause() instanceof BindException) {
                    retries++;
                    if (retries > MAX_RETRIES) {
                        throw (BindException) e.getCause();
                    }
                } else {
                    throw e;
                }
            }
        }
        throw new IOException("Failed to start any storage server.");
    }

    private void startServers() throws Exception {
        log.info("Starting {} storage servers.", spec.numServers());
        ExecutorService executor = Executors.newCachedThreadPool();
        List<Future<LifecycleComponent>> startFutures = Lists.newArrayList();
        for (int i = 0; i < spec.numServers(); i++) {
            Future<LifecycleComponent> future = executor.submit(() -> startServer());
            startFutures.add(future);
        }
        for (Future<LifecycleComponent> future : startFutures) {
            servers.add(future.get());
        }
        log.info("Started {} storage servers.", spec.numServers());
        executor.shutdown();
    }


    private void createDefaultNamespaces() throws Exception {
        String serviceUri = String.format(
            "bk://%s/",
            getRpcEndpoints().stream()
                .map(endpoint -> NetUtils.endpointToString(endpoint))
                .collect(Collectors.joining(",")));
        StorageClientSettings settings = StorageClientSettings.newBuilder()
            .serviceUri(serviceUri)
            .usePlaintext(true)
            .build();
        log.info("Service uri are : {}", serviceUri);
        String namespaceName = "default";
        try (StorageAdminClient admin = StorageClientBuilder.newBuilder()
            .withSettings(settings)
            .buildAdmin()) {

            boolean created = false;
            while (!created) {
                try {
                    NamespaceProperties nsProps = result(admin.getNamespace(namespaceName));
                    log.info("Namespace '{}':\n{}", namespaceName, nsProps);
                    created = true;
                } catch (NamespaceNotFoundException nnfe) {
                    log.info("Namespace '{}' is not found.", namespaceName);
                    log.info("Creating namespace '{}' ...", namespaceName);
                    try {
                        NamespaceProperties nsProps = result(
                            admin.createNamespace(
                                namespaceName,
                                NamespaceConfiguration.newBuilder()
                                    .setDefaultStreamConf(DEFAULT_STREAM_CONF)
                                    .build()));
                        log.info("Successfully created namespace '{}':", namespaceName);
                        log.info("{}", nsProps);
                    } catch (ClientException ce) {
                        // encountered exception, try to fetch the namespace again
                    }
                }
            }
        }
    }

    private void stopServers() {
        for (LifecycleComponent server : servers) {
            server.close();
        }
    }

    @Override
    protected void doStart() {
        try {
            // start zookeeper servers
            startZooKeeper();

            // initialize the cluster
            initializeCluster();

            // start servers
            startServers();

            // create default namespaces
            createDefaultNamespaces();
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }

    @Override
    protected void doStop() {
    }

    @Override
    protected void doClose() throws IOException {
        // stop the servers
        stopServers();
        // stop zookeeper
        stopZooKeeper();
    }
}
