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

import static java.nio.charset.StandardCharsets.UTF_8;

import com.google.common.base.Stopwatch;
import java.io.IOException;
import java.util.HashSet;
import java.util.List;
import java.util.Set;
import java.util.concurrent.CopyOnWriteArraySet;
import java.util.concurrent.TimeUnit;
import org.apache.bookkeeper.stats.NullStatsLogger;
import org.apache.bookkeeper.stats.StatsLogger;
import org.apache.bookkeeper.zookeeper.BoundExponentialBackoffRetryPolicy;
import org.apache.bookkeeper.zookeeper.RetryPolicy;
import org.apache.distributedlog.util.FailpointUtils;
import org.apache.distributedlog.zk.ZKWatcherManager;
import org.apache.zookeeper.KeeperException;
import org.apache.zookeeper.WatchedEvent;
import org.apache.zookeeper.Watcher;
import org.apache.zookeeper.Watcher.Event.EventType;
import org.apache.zookeeper.Watcher.Event.KeeperState;
import org.apache.zookeeper.ZooDefs;
import org.apache.zookeeper.ZooKeeper;
import org.apache.zookeeper.data.ACL;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;


/**
 * ZooKeeper Client wrapper over {@link org.apache.bookkeeper.zookeeper.ZooKeeperClient}.
 * It handles retries on session expires and provides a watcher manager {@link ZKWatcherManager}.
 *
 * <h3>Metrics</h3>
 * <ul>
 * <li> zookeeper operation stats are exposed under scope <code>zk</code> by
 * {@link org.apache.bookkeeper.zookeeper.ZooKeeperClient}
 * <li> stats on zookeeper watched events are exposed under scope <code>watcher</code> by
 * {@link org.apache.bookkeeper.zookeeper.ZooKeeperWatcherBase}
 * <li> stats about {@link ZKWatcherManager} are exposed under scope <code>watcher_manager</code>
 * </ul>
 */
public class ZooKeeperClient {

    /**
     * interface used to authenticate zk client.
     */
    public interface Credentials {

        Credentials NONE = new Credentials() {
            @Override
            public void authenticate(ZooKeeper zooKeeper) {
                // noop
            }
        };

        void authenticate(ZooKeeper zooKeeper);
    }
    /**
     * interface impl used to authenticate zk client.
     */
    public static class DigestCredentials implements Credentials {

        String username;
        String password;

        public DigestCredentials(String username, String password) {
            this.username = username;
            this.password = password;
        }

        @Override
        public void authenticate(ZooKeeper zooKeeper) {
            zooKeeper.addAuthInfo("digest", String.format("%s:%s", username, password).getBytes(UTF_8));
        }
    }

    /**
     * Notify a zk session expire event.
     */
    public interface ZooKeeperSessionExpireNotifier {
        void notifySessionExpired();
    }

    /**
     * Indicates an error connecting to a zookeeper cluster.
     */
    public static class ZooKeeperConnectionException extends IOException {
        private static final long serialVersionUID = 6682391687004819361L;

        public ZooKeeperConnectionException(String message) {
            super(message);
        }

        public ZooKeeperConnectionException(String message, Throwable cause) {
            super(message, cause);
        }
    }

    private static final Logger LOG = LoggerFactory.getLogger(ZooKeeperClient.class.getName());

    private final String name;
    private final int sessionTimeoutMs;
    private final int defaultConnectionTimeoutMs;
    private final String zooKeeperServers;
    // GuardedBy "this", but still volatile for tests, where we want to be able to see writes
    // made from within long synchronized blocks.
    private volatile ZooKeeper zooKeeper = null;
    private final RetryPolicy retryPolicy;
    private final StatsLogger statsLogger;
    private final int retryThreadCount;
    private final double requestRateLimit;
    private final Credentials credentials;
    private volatile boolean authenticated = false;
    private Stopwatch disconnectedStopwatch = null;

    private boolean closed = false;

    final Set<Watcher> watchers = new CopyOnWriteArraySet<Watcher>();

    // watcher manager to manage watchers
    private final ZKWatcherManager watcherManager;

    /**
     * Creates an unconnected client that will lazily attempt to connect on the first call to
     * {@link #get}.  All successful connections will be authenticated with the given
     * {@code credentials}.
     *
     * @param sessionTimeoutMs
     *          ZK session timeout in milliseconds
     * @param connectionTimeoutMs
     *          ZK connection timeout in milliseconds
     * @param zooKeeperServers
     *          the set of servers forming the ZK cluster
     */
    ZooKeeperClient(int sessionTimeoutMs, int connectionTimeoutMs, String zooKeeperServers) {
        this("default", sessionTimeoutMs, connectionTimeoutMs, zooKeeperServers, null, NullStatsLogger.INSTANCE, 1, 0,
             Credentials.NONE);
    }

    ZooKeeperClient(String name,
                    int sessionTimeoutMs,
                    int connectionTimeoutMs,
                    String zooKeeperServers,
                    RetryPolicy retryPolicy,
                    StatsLogger statsLogger,
                    int retryThreadCount,
                    double requestRateLimit,
                    Credentials credentials) {
        this.name = name;
        this.sessionTimeoutMs = sessionTimeoutMs;
        this.zooKeeperServers = zooKeeperServers;
        this.defaultConnectionTimeoutMs = connectionTimeoutMs;
        this.retryPolicy = retryPolicy;
        this.statsLogger = statsLogger;
        this.retryThreadCount = retryThreadCount;
        this.requestRateLimit = requestRateLimit;
        this.credentials = credentials;
        this.watcherManager = ZKWatcherManager.newBuilder()
                .name(name)
                .zkc(this)
                .statsLogger(statsLogger.scope("watcher_manager"))
                .build();
    }

    public List<ACL> getDefaultACL() {
        if (Credentials.NONE == credentials) {
            return ZooDefs.Ids.OPEN_ACL_UNSAFE;
        } else {
            return DistributedLogConstants.EVERYONE_READ_CREATOR_ALL;
        }
    }

    public ZKWatcherManager getWatcherManager() {
        return watcherManager;
    }

    /**
     * Returns the current active ZK connection or establishes a new one if none has yet been
     * established or a previous connection was disconnected or had its session time out.
     *
     * @return a connected ZooKeeper client
     * @throws ZooKeeperConnectionException if there was a problem connecting to the ZK cluster
     * @throws InterruptedException if interrupted while waiting for a connection to be established
     * session timeout
     */
    public synchronized ZooKeeper get()
        throws ZooKeeperConnectionException, InterruptedException {

        try {
            FailpointUtils.checkFailPoint(FailpointUtils.FailPointName.FP_ZooKeeperConnectionLoss);
        } catch (IOException ioe) {
            throw new ZooKeeperConnectionException("Client "
                    + name + " failed on establishing zookeeper connection", ioe);
        }

        // This indicates that the client was explictly closed
        if (closed) {
            throw new ZooKeeperConnectionException("Client " + name + " has already been closed");
        }

        // the underneath zookeeper is retryable zookeeper
        if (zooKeeper != null && retryPolicy != null) {
            if (zooKeeper.getState().equals(ZooKeeper.States.CONNECTED)) {
                // the zookeeper client is connected
                disconnectedStopwatch = null;
            } else {
                if (disconnectedStopwatch == null) {
                    disconnectedStopwatch = Stopwatch.createStarted();
                } else {
                    long disconnectedMs = disconnectedStopwatch.elapsed(TimeUnit.MILLISECONDS);
                    if (disconnectedMs > defaultConnectionTimeoutMs) {
                        closeInternal();
                        authenticated = false;
                    }
                }
            }
        }

        if (zooKeeper == null) {
            zooKeeper = buildZooKeeper();
            disconnectedStopwatch = null;
        }

        // In case authenticate throws an exception, the caller can try to recover the client by
        // calling get again.
        if (!authenticated) {
            credentials.authenticate(zooKeeper);
            authenticated = true;
        }

        return zooKeeper;
    }

    private ZooKeeper buildZooKeeper()
        throws ZooKeeperConnectionException, InterruptedException {
        Watcher watcher = new Watcher() {
            @Override
            public void process(WatchedEvent event) {
                switch (event.getType()) {
                    case None:
                        switch (event.getState()) {
                            case Expired:
                                if (null == retryPolicy) {
                                    LOG.info("ZooKeeper {}' session expired. Event: {}", name, event);
                                    closeInternal();
                                }
                                authenticated = false;
                                break;
                            case Disconnected:
                                if (null == retryPolicy) {
                                    LOG.info("ZooKeeper {} is disconnected from zookeeper now,"
                                            + " but it is OK unless we received EXPIRED event.", name);
                                }
                                // Mark as not authenticated if expired or disconnected. In both cases
                                // we lose any attached auth info. Relying on Expired/Disconnected is
                                // sufficient since all Expired/Disconnected events are processed before
                                // all SyncConnected events, and the underlying member is not updated until
                                // SyncConnected is received.
                                authenticated = false;
                                break;
                            default:
                                break;
                        }
                }

                try {
                    for (Watcher watcher : watchers) {
                        try {
                            watcher.process(event);
                        } catch (Throwable t) {
                            LOG.warn("Encountered unexpected exception from watcher {} : ", watcher, t);
                        }
                    }
                } catch (Throwable t) {
                    LOG.warn("Encountered unexpected exception when firing watched event {} : ", event, t);
                }
            }
        };

        Set<Watcher> watchers = new HashSet<Watcher>();
        watchers.add(watcher);

        ZooKeeper zk;
        try {
            RetryPolicy opRetryPolicy = null == retryPolicy
                    ? new BoundExponentialBackoffRetryPolicy(sessionTimeoutMs, sessionTimeoutMs, 0) : retryPolicy;
            RetryPolicy connectRetryPolicy = null == retryPolicy
                    ? new BoundExponentialBackoffRetryPolicy(sessionTimeoutMs, sessionTimeoutMs, 0) :
                    new BoundExponentialBackoffRetryPolicy(sessionTimeoutMs, sessionTimeoutMs, Integer.MAX_VALUE);
            zk = org.apache.bookkeeper.zookeeper.ZooKeeperClient.newBuilder()
                    .connectString(zooKeeperServers)
                    .sessionTimeoutMs(sessionTimeoutMs)
                    .watchers(watchers)
                    .operationRetryPolicy(opRetryPolicy)
                    .connectRetryPolicy(connectRetryPolicy)
                    .statsLogger(statsLogger)
                    .retryThreadCount(retryThreadCount)
                    .requestRateLimit(requestRateLimit)
                    .build();
        } catch (KeeperException e) {
            throw new ZooKeeperConnectionException("Problem connecting to servers: " + zooKeeperServers, e);
        } catch (IOException e) {
            throw new ZooKeeperConnectionException("Problem connecting to servers: " + zooKeeperServers, e);
        }
        return zk;
    }

    /**
     * Clients that need to re-establish state after session expiration can register an
     * {@code onExpired} command to execute.
     *
     * @param onExpired the {@code Command} to register
     * @return the new {@link Watcher} which can later be passed to {@link #unregister} for
     *         removal.
     */
    public Watcher registerExpirationHandler(final ZooKeeperSessionExpireNotifier onExpired) {
        Watcher watcher = new Watcher() {
            @Override
            public void process(WatchedEvent event) {
                if (event.getType() == EventType.None && event.getState() == KeeperState.Expired) {
                    try {
                        onExpired.notifySessionExpired();
                    } catch (Exception exc) {
                        // do nothing
                    }
                }
            }
        };
        register(watcher);
        return watcher;
    }

    /**
     * Clients that need to register a top-level {@code Watcher} should do so using this method.  The
     * registered {@code watcher} will remain registered across re-connects and session expiration
     * events.
     *
     * @param watcher the {@code Watcher to register}
     */
    public void register(Watcher watcher) {
        if (null != watcher) {
            watchers.add(watcher);
        }
    }

    /**
     * Clients can attempt to unregister a top-level {@code Watcher} that has previously been
     * registered.
     *
     * @param watcher the {@code Watcher} to unregister as a top-level, persistent watch
     * @return whether the given {@code Watcher} was found and removed from the active set
     */
    public boolean unregister(Watcher watcher) {
        return null != watcher && watchers.remove(watcher);
    }

    /**
     * Closes the current connection if any expiring the current ZooKeeper session.  Any subsequent
     * calls to this method will no-op until the next successful {@link #get}.
     */
    public synchronized void closeInternal() {
        if (zooKeeper != null) {
            try {
                LOG.info("Closing zookeeper client {}.", name);
                zooKeeper.close();
                LOG.info("Closed zookeeper client {}.", name);
            } catch (InterruptedException e) {
                Thread.currentThread().interrupt();
                LOG.warn("Interrupted trying to close zooKeeper {} : ", name, e);
            } finally {
                zooKeeper = null;
            }
        }
    }

    /**
     * Closes the underlying zookeeper instance.
     * Subsequent attempts to {@link #get} will fail
     */
    public synchronized void close() {
        if (closed) {
            return;
        }
        LOG.info("Close zookeeper client {}.", name);
        closeInternal();
        // unregister gauges to prevent GC spiral
        this.watcherManager.unregisterGauges();
        closed = true;
    }
}