/*
 * 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.meta;

import static com.google.common.base.Preconditions.checkArgument;
import static com.google.common.base.Preconditions.checkNotNull;

import com.google.common.annotations.VisibleForTesting;
import com.google.common.collect.Sets;
import com.google.common.util.concurrent.UncheckedExecutionException;
import edu.umd.cs.findbugs.annotations.SuppressFBWarnings;
import java.net.URI;
import java.util.Collections;
import java.util.Optional;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.ScheduledExecutorService;
import java.util.function.Function;
import lombok.AccessLevel;
import lombok.Getter;
import lombok.NoArgsConstructor;
import lombok.ToString;
import lombok.extern.slf4j.Slf4j;
import org.apache.bookkeeper.common.util.ReflectionUtils;
import org.apache.bookkeeper.conf.ClientConfiguration;
import org.apache.bookkeeper.conf.ServerConfiguration;
import org.apache.bookkeeper.discover.RegistrationManager;
import org.apache.bookkeeper.meta.exceptions.Code;
import org.apache.bookkeeper.meta.exceptions.MetadataException;
import org.apache.bookkeeper.stats.NullStatsLogger;
import org.apache.commons.configuration2.ex.ConfigurationException;
import org.apache.commons.lang3.StringUtils;

/**
 * A driver manager for managing a set of metadata drivers.
 *
 */
@NoArgsConstructor(access = AccessLevel.PRIVATE)
@Slf4j
public final class MetadataDrivers {

    static final String ZK_CLIENT_DRIVER_CLASS = "org.apache.bookkeeper.meta.zk.ZKMetadataClientDriver";
    static final String ZK_BOOKIE_DRIVER_CLASS = "org.apache.bookkeeper.meta.zk.ZKMetadataBookieDriver";
    static final String BK_METADATA_CLIENT_DRIVERS_PROPERTY = "bookkeeper.metadata.client.drivers";
    static final String BK_METADATA_BOOKIE_DRIVERS_PROPERTY = "bookkeeper.metadata.bookie.drivers";

    @ToString
    static class MetadataClientDriverInfo {

        final Class<? extends MetadataClientDriver> driverClass;
        final String driverClassName;

        MetadataClientDriverInfo(Class<? extends MetadataClientDriver> driverClass) {
            this.driverClass = driverClass;
            this.driverClassName = driverClass.getName();
        }

    }

    @ToString
    static class MetadataBookieDriverInfo {

        final Class<? extends MetadataBookieDriver> driverClass;
        final String driverClassName;

        MetadataBookieDriverInfo(Class<? extends MetadataBookieDriver> driverClass) {
            this.driverClass = driverClass;
            this.driverClassName = driverClass.getName();
        }

    }

    @Getter(AccessLevel.PACKAGE)
    private static final ConcurrentMap<String, MetadataClientDriverInfo> clientDrivers;
    @Getter(AccessLevel.PACKAGE)
    private static final ConcurrentMap<String, MetadataBookieDriverInfo> bookieDrivers;

    static {
        clientDrivers = new ConcurrentHashMap<>();
        bookieDrivers = new ConcurrentHashMap<>();
        loadInitialDrivers();
    }

    @VisibleForTesting
    static void loadInitialDrivers() {
        loadInitialClientDrivers();
        loadInitialBookieDrivers();
        log.info("BookKeeper metadata driver manager initialized");
    }

    private static void loadInitialClientDrivers() {
        Set<String> driverList = Sets.newHashSet();

        // add default zookeeper based driver
        driverList.add(ZK_CLIENT_DRIVER_CLASS);

        // load drivers from system property
        String driversStr = System.getProperty(BK_METADATA_CLIENT_DRIVERS_PROPERTY);
        if (null != driversStr) {
            String[] driversArray = StringUtils.split(driversStr, ':');
            Collections.addAll(driverList, driversArray);
        }

        // initialize the drivers
        for (String driverClsName : driverList) {
            try {
                MetadataClientDriver driver =
                    ReflectionUtils.newInstance(driverClsName, MetadataClientDriver.class);
                MetadataClientDriverInfo driverInfo =
                    new MetadataClientDriverInfo(driver.getClass());
                clientDrivers.put(driver.getScheme().toLowerCase(), driverInfo);
            } catch (Exception e) {
                log.warn("Failed to load metadata client driver {}", driverClsName, e);
            }
        }
    }

    private static void loadInitialBookieDrivers() {
        Set<String> driverList = Sets.newHashSet();

        // add default zookeeper based driver
        driverList.add(ZK_BOOKIE_DRIVER_CLASS);

        // load drivers from system property
        String driversStr = System.getProperty(BK_METADATA_BOOKIE_DRIVERS_PROPERTY);
        if (null != driversStr) {
            String[] driversArray = StringUtils.split(driversStr, ':');
            Collections.addAll(driverList, driversArray);
        }

        // initialize the drivers
        for (String driverClsName : driverList) {
            try {
                MetadataBookieDriver driver =
                    ReflectionUtils.newInstance(driverClsName, MetadataBookieDriver.class);
                MetadataBookieDriverInfo driverInfo =
                    new MetadataBookieDriverInfo(driver.getClass());
                bookieDrivers.put(driver.getScheme().toLowerCase(), driverInfo);
            } catch (Exception e) {
                log.warn("Failed to load metadata bookie driver {}", driverClsName, e);
            }
        }
    }

    /**
     * Register the metadata client {@code driver}.
     *
     * @param metadataBackendScheme scheme of metadata backend.
     * @param driver metadata client driver
     */
    public static void registerClientDriver(String metadataBackendScheme,
                                            Class<? extends MetadataClientDriver> driver) {
        registerClientDriver(metadataBackendScheme, driver, false);
    }

    @VisibleForTesting
    public static void registerClientDriver(String metadataBackendScheme,
                                            Class<? extends MetadataClientDriver> driver,
                                            boolean allowOverride) {
        String scheme = metadataBackendScheme.toLowerCase();
        MetadataClientDriverInfo oldDriverInfo = clientDrivers.get(scheme);
        if (null != oldDriverInfo && !allowOverride) {
            return;
        }
        MetadataClientDriverInfo newDriverInfo = new MetadataClientDriverInfo(driver);
        oldDriverInfo = clientDrivers.putIfAbsent(scheme, newDriverInfo);
        if (null != oldDriverInfo) {
            if (log.isDebugEnabled()) {
                log.debug("Metadata client driver for {} is already there.", scheme);
            }
            if (allowOverride) {
                if (log.isDebugEnabled()) {
                    log.debug("Overriding client driver for {}", scheme);
                }
                clientDrivers.put(scheme, newDriverInfo);
            }
        }
    }

    /**
     * Register the metadata bookie {@code driver}.
     *
     * @param metadataBackendScheme scheme of metadata backend.
     * @param driver metadata bookie driver
     */
    public static void registerBookieDriver(String metadataBackendScheme,
                                            Class<? extends MetadataBookieDriver> driver) {
        registerBookieDriver(metadataBackendScheme, driver, false);
    }

    @VisibleForTesting
    public static void registerBookieDriver(String metadataBackendScheme,
                                            Class<? extends MetadataBookieDriver> driver,
                                            boolean allowOverride) {
        String scheme = metadataBackendScheme.toLowerCase();
        MetadataBookieDriverInfo oldDriverInfo = bookieDrivers.get(scheme);
        if (null != oldDriverInfo && !allowOverride) {
            return;
        }
        MetadataBookieDriverInfo newDriverInfo = new MetadataBookieDriverInfo(driver);
        oldDriverInfo = bookieDrivers.putIfAbsent(scheme, newDriverInfo);
        if (null != oldDriverInfo) {
            if (log.isDebugEnabled()) {
                log.debug("Metadata bookie driver for {} is already there.", scheme);
            }
            if (allowOverride) {
                if (log.isDebugEnabled()) {
                    log.debug("Overriding bookie driver for {}", scheme);
                }
                bookieDrivers.put(scheme, newDriverInfo);
            }
        }
    }

    /**
     * Retrieve the client driver for {@code scheme}.
     *
     * @param scheme the scheme for the metadata client driver
     * @return the metadata client driver
     * @throws NullPointerException when scheme is null
     */
    public static MetadataClientDriver getClientDriver(String scheme) {
        checkNotNull(scheme, "Client Driver Scheme is null");
        MetadataClientDriverInfo driverInfo = clientDrivers.get(scheme.toLowerCase());
        if (null == driverInfo) {
            throw new IllegalArgumentException("Unknown backend " + scheme);
        }
        return ReflectionUtils.newInstance(driverInfo.driverClass);
    }

    /**
     * Retrieve the client driver for {@code uri}.
     *
     * @param uri the metadata service uri
     * @return the metadata client driver for {@code uri}
     * @throws NullPointerException if the metadata service {@code uri} is null or doesn't have scheme
     *          or there is no namespace driver registered for the scheme
     * @throws IllegalArgumentException if the metadata service {@code uri} scheme is illegal
     */
    public static MetadataClientDriver getClientDriver(URI uri) {
        // Validate the uri and load the backend according to scheme
        checkNotNull(uri, "Metadata service uri is null");
        String scheme = uri.getScheme();
        checkNotNull(scheme, "Invalid metadata service uri : " + uri);
        scheme = scheme.toLowerCase();
        String[] schemeParts = StringUtils.split(scheme, '+');
        checkArgument(schemeParts.length > 0,
                "Invalid metadata service scheme found : " + uri);
        return getClientDriver(schemeParts[0]);
    }

    /**
     * Retrieve the bookie driver for {@code scheme}.
     *
     * @param scheme the scheme for the metadata bookie driver
     * @return the metadata bookie driver
     * @throws NullPointerException when scheme is null
     */
    public static MetadataBookieDriver getBookieDriver(String scheme) {
        checkNotNull(scheme, "Bookie Driver Scheme is null");
        MetadataBookieDriverInfo driverInfo = bookieDrivers.get(scheme.toLowerCase());
        if (null == driverInfo) {
            throw new IllegalArgumentException("Unknown backend " + scheme);
        }
        return ReflectionUtils.newInstance(driverInfo.driverClass);
    }

    /**
     * Retrieve the bookie driver for {@code uri}.
     *
     * @param uri the metadata service uri
     * @return the metadata bookie driver for {@code uri}
     * @throws NullPointerException if the metadata service {@code uri} is null or doesn't have scheme
     *          or there is no namespace driver registered for the scheme
     * @throws IllegalArgumentException if the metadata service {@code uri} scheme is illegal
     */
    public static MetadataBookieDriver getBookieDriver(URI uri) {
        // Validate the uri and load the backend according to scheme
        checkNotNull(uri, "Metadata service uri is null");
        String scheme = uri.getScheme();
        checkNotNull(scheme, "Invalid metadata service uri : " + uri);
        scheme = scheme.toLowerCase();
        String[] schemeParts = StringUtils.split(scheme, '+');
        checkArgument(schemeParts.length > 0,
                "Invalid metadata service scheme found : " + uri);
        return getBookieDriver(schemeParts[0]);
    }

    /**
     * Process the provided <i>function</i> with metadata client driver resolved
     * from the metadata service uri returned by {@link ClientConfiguration#getMetadataServiceUri()}.
     *
     * @param conf client configuration
     * @param function function to apply with metadata client driver.
     * @param executorService executor service used by the metadata client driver.
     * @throws MetadataException when failed to access metadata store
     * @throws ExecutionException exception thrown when processing <tt>function</tt>.
     */
    @SuppressFBWarnings("RCN_REDUNDANT_NULLCHECK_WOULD_HAVE_BEEN_A_NPE")
    public static <T> T runFunctionWithMetadataClientDriver(ClientConfiguration conf,
                                                            Function<MetadataClientDriver, T> function,
                                                            ScheduledExecutorService executorService)
            throws MetadataException, ExecutionException {
        try (MetadataClientDriver driver = MetadataDrivers.getClientDriver(
            URI.create(conf.getMetadataServiceUri())
        )) {
            driver.initialize(conf, executorService, NullStatsLogger.INSTANCE, Optional.empty());
            try {
                return function.apply(driver);
            } catch (Exception uee) {
                if (uee.getCause() instanceof MetadataException) {
                    throw (MetadataException) uee.getCause();
                } else {
                    throw new ExecutionException(uee.getMessage(), uee.getCause());
                }
            }
        } catch (ConfigurationException e) {
            throw new MetadataException(Code.INVALID_METADATA_SERVICE_URI, e);
        }
    }

    /**
     * Process the provided <i>function</i> with metadata bookie driver resolved
     * from the metadata service uri returned by {@link ServerConfiguration#getMetadataServiceUri()}.
     *
     * @param conf server configuration
     * @param function function to apply with metadata bookie driver.
     * @throws MetadataException when failed to access metadata store
     * @throws ExecutionException exception thrown when processing <tt>function</tt>.
     */
    @SuppressFBWarnings("RCN_REDUNDANT_NULLCHECK_WOULD_HAVE_BEEN_A_NPE")
    public static <T> T runFunctionWithMetadataBookieDriver(ServerConfiguration conf,
                                                            Function<MetadataBookieDriver, T> function)
            throws MetadataException, ExecutionException {
        try (MetadataBookieDriver driver = MetadataDrivers.getBookieDriver(
            URI.create(conf.getMetadataServiceUri())
        )) {
            driver.initialize(conf, NullStatsLogger.INSTANCE);
            try {
                return function.apply(driver);
            } catch (Exception uee) {
                if (uee.getCause() instanceof MetadataException) {
                    throw (MetadataException) uee.getCause();
                } else {
                    throw new ExecutionException(uee.getMessage(), uee.getCause());
                }
            }
        } catch (ConfigurationException e) {
            throw new MetadataException(Code.INVALID_METADATA_SERVICE_URI, e);
        }
    }

    /**
     * Process the provided <i>function</i> with registration manager resolved
     * from the metadata service uri returned by {@link ServerConfiguration#getMetadataServiceUri()}.
     *
     * @param conf server configuration
     * @param function function to apply with registration manager.
     * @throws MetadataException when failed to access metadata store
     * @throws ExecutionException exception thrown when processing <tt>consumer</tt>.
     */
    public static <T> T runFunctionWithRegistrationManager(ServerConfiguration conf,
                                                           Function<RegistrationManager, T> function)
            throws MetadataException, ExecutionException {
        return runFunctionWithMetadataBookieDriver(
                conf, driver -> {
                    try (RegistrationManager rm = driver.createRegistrationManager()) {
                        return function.apply(rm);
                    }
                });
    }

    /**
     * Process the provided <i>function</i> with ledger manager factory resolved
     * from the metadata service uri returned by {@link ServerConfiguration#getMetadataServiceUri()}.
     *
     * @param conf server configuration
     * @param function function to apply with ledger manager factory.
     * @throws MetadataException when failed to access metadata store
     * @throws ExecutionException exception thrown when processing <tt>consumer</tt>.
     */
    public static <T> T runFunctionWithLedgerManagerFactory(ServerConfiguration conf,
                                                            Function<LedgerManagerFactory, T> function)
            throws MetadataException, ExecutionException {
        return runFunctionWithMetadataBookieDriver(conf, driver -> {
            try {
                return function.apply(driver.getLedgerManagerFactory());
            } catch (MetadataException me) {
                throw new UncheckedExecutionException(me.getMessage(), me);
            }
        });
    }

}
