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

import java.io.IOException;

import org.apache.hadoop.util.Progressable;
import org.apache.log4j.Logger;
import org.apache.zookeeper.KeeperException;
import org.apache.zookeeper.CreateMode;
import org.apache.zookeeper.data.ACL;
import org.apache.zookeeper.data.Stat;

import java.util.ArrayList;
import java.util.Collections;
import java.util.Comparator;
import java.util.List;

import org.apache.zookeeper.Watcher;
import org.apache.zookeeper.ZooKeeper;

/**
 * ZooKeeper provides only atomic operations.  ZooKeeperExt provides additional
 * non-atomic operations that are useful.  It also provides wrappers to
 * deal with ConnectionLossException.  All methods of this class
 * should be thread-safe.
 */
public class ZooKeeperExt {
  /** Length of the ZK sequence number */
  public static final int SEQUENCE_NUMBER_LENGTH = 10;
  /** Internal logger */
  private static final Logger LOG = Logger.getLogger(ZooKeeperExt.class);
  /** Internal ZooKeeper */
  private final ZooKeeper zooKeeper;
  /** Ensure we have progress */
  private final Progressable progressable;
  /** Number of max attempts to retry when failing due to connection loss */
  private final int maxRetryAttempts;
  /** Milliseconds to wait before trying again due to connection loss */
  private final int retryWaitMsecs;

  /**
   * Constructor to connect to ZooKeeper, does not make progress
   *
   * @param connectString Comma separated host:port pairs, each corresponding
   *        to a zk server. e.g.
   *        "127.0.0.1:3000,127.0.0.1:3001,127.0.0.1:3002" If the optional
   *        chroot suffix is used the example would look
   *        like: "127.0.0.1:3000,127.0.0.1:3001,127.0.0.1:3002/app/a"
   *        where the client would be rooted at "/app/a" and all paths
   *        would be relative to this root - ie getting/setting/etc...
   *        "/foo/bar" would result in operations being run on
   *        "/app/a/foo/bar" (from the server perspective).
   * @param sessionTimeout Session timeout in milliseconds
   * @param maxRetryAttempts Max retry attempts during connection loss
   * @param retryWaitMsecs Msecs to wait when retrying due to connection
   *        loss
   * @param watcher A watcher object which will be notified of state changes,
   *        may also be notified for node events
   * @throws IOException
   */
  public ZooKeeperExt(String connectString,
                      int sessionTimeout,
                      int maxRetryAttempts,
                      int retryWaitMsecs,
                      Watcher watcher) throws IOException {
    this(connectString, sessionTimeout, maxRetryAttempts,
        retryWaitMsecs, watcher, null);
  }

  /**
   * Constructor to connect to ZooKeeper, make progress
   *
   * @param connectString Comma separated host:port pairs, each corresponding
   *        to a zk server. e.g.
   *        "127.0.0.1:3000,127.0.0.1:3001,127.0.0.1:3002" If the optional
   *        chroot suffix is used the example would look
   *        like: "127.0.0.1:3000,127.0.0.1:3001,127.0.0.1:3002/app/a"
   *        where the client would be rooted at "/app/a" and all paths
   *        would be relative to this root - ie getting/setting/etc...
   *        "/foo/bar" would result in operations being run on
   *        "/app/a/foo/bar" (from the server perspective).
   * @param sessionTimeout Session timeout in milliseconds
   * @param maxRetryAttempts Max retry attempts during connection loss
   * @param retryWaitMsecs Msecs to wait when retrying due to connection
   *        loss
   * @param watcher A watcher object which will be notified of state changes,
   *        may also be notified for node events
   * @param progressable Makes progress for longer operations
   * @throws IOException
   */
  public ZooKeeperExt(String connectString,
      int sessionTimeout,
      int maxRetryAttempts,
      int retryWaitMsecs,
      Watcher watcher,
      Progressable progressable) throws IOException {
    this.zooKeeper = new ZooKeeper(connectString, sessionTimeout, watcher);
    this.progressable = progressable;
    this.maxRetryAttempts = maxRetryAttempts;
    this.retryWaitMsecs = retryWaitMsecs;
  }

  /**
   * Provides a possibility of a creating a path consisting of more than one
   * znode (not atomic).  If recursive is false, operates exactly the
   * same as create().
   *
   * @param path path to create
   * @param data data to set on the final znode
   * @param acl acls on each znode created
   * @param createMode only affects the final znode
   * @param recursive if true, creates all ancestors
   * @return Actual created path
   * @throws KeeperException
   * @throws InterruptedException
   */
  public String createExt(
      final String path,
      byte[] data,
      List<ACL> acl,
      CreateMode createMode,
      boolean recursive) throws KeeperException, InterruptedException {
    if (LOG.isDebugEnabled()) {
      LOG.debug("createExt: Creating path " + path);
    }

    int attempt = 0;
    while (attempt < maxRetryAttempts) {
      try {
        if (!recursive) {
          return zooKeeper.create(path, data, acl, createMode);
        }

        try {
          return zooKeeper.create(path, data, acl, createMode);
        } catch (KeeperException.NoNodeException e) {
          if (LOG.isDebugEnabled()) {
            LOG.debug("createExt: Cannot directly create node " + path);
          }
        }

        int pos = path.indexOf("/", 1);
        for (; pos != -1; pos = path.indexOf("/", pos + 1)) {
          try {
            if (progressable != null) {
              progressable.progress();
            }
            String filePath = path.substring(0, pos);
            if (zooKeeper.exists(filePath, false) == null) {
              zooKeeper.create(
                  filePath, null, acl, CreateMode.PERSISTENT);
            }
          } catch (KeeperException.NodeExistsException e) {
            if (LOG.isDebugEnabled()) {
              LOG.debug("createExt: Znode " + path.substring(0, pos) +
                  " already exists");
            }
          }
        }
        return zooKeeper.create(path, data, acl, createMode);
      } catch (KeeperException.ConnectionLossException e) {
        LOG.warn("createExt: Connection loss on attempt " + attempt + ", " +
            "waiting " + retryWaitMsecs + " msecs before retrying.", e);
      }
      ++attempt;
      Thread.sleep(retryWaitMsecs);
    }
    throw new IllegalStateException("createExt: Failed to create " + path  +
        " after " + attempt + " tries!");
  }

  /**
   * Data structure for handling the output of createOrSet()
   */
  public static class PathStat {
    /** Path to created znode (if any) */
    private String path;
    /** Stat from set znode (if any) */
    private Stat stat;

    /**
     * Put in results from createOrSet()
     *
     * @param path Path to created znode (or null)
     * @param stat Stat from set znode (if set)
     */
    public PathStat(String path, Stat stat) {
      this.path = path;
      this.stat = stat;
    }

    /**
     * Get the path of the created znode if it was created.
     *
     * @return Path of created znode or null if not created
     */
    public String getPath() {
      return path;
    }

    /**
     * Get the stat of the set znode if set
     *
     * @return Stat of set znode or null if not set
     */
    public Stat getStat() {
      return stat;
    }
  }

  /**
   * Create a znode.  Set the znode if the created znode already exists.
   *
   * @param path path to create
   * @param data data to set on the final znode
   * @param acl acls on each znode created
   * @param createMode only affects the final znode
   * @param recursive if true, creates all ancestors
   * @param version Version to set if setting
   * @return Path of created znode or Stat of set znode
   * @throws InterruptedException
   * @throws KeeperException
   */
  public PathStat createOrSetExt(final String path,
      byte[] data,
      List<ACL> acl,
      CreateMode createMode,
      boolean recursive,
      int version) throws KeeperException, InterruptedException {
    String createdPath = null;
    Stat setStat = null;
    try {
      createdPath = createExt(path, data, acl, createMode, recursive);
    } catch (KeeperException.NodeExistsException e) {
      if (LOG.isDebugEnabled()) {
        LOG.debug("createOrSet: Node exists on path " + path);
      }
      setStat = zooKeeper.setData(path, data, version);
    }
    return new PathStat(createdPath, setStat);
  }

  /**
   * Create a znode if there is no other znode there
   *
   * @param path path to create
   * @param data data to set on the final znode
   * @param acl acls on each znode created
   * @param createMode only affects the final znode
   * @param recursive if true, creates all ancestors
   * @return Path of created znode or Stat of set znode
   * @throws InterruptedException
   * @throws KeeperException
   */
  public PathStat createOnceExt(final String path,
      byte[] data,
      List<ACL> acl,
      CreateMode createMode,
      boolean recursive) throws KeeperException, InterruptedException {
    String createdPath = null;
    Stat setStat = null;
    try {
      createdPath = createExt(path, data, acl, createMode, recursive);
    } catch (KeeperException.NodeExistsException e) {
      if (LOG.isDebugEnabled()) {
        LOG.debug("createOnceExt: Node already exists on path " + path);
      }
    }
    return new PathStat(createdPath, setStat);
  }

  /**
   * Delete a path recursively.  When the deletion is recursive, it is a
   * non-atomic operation, hence, not part of ZooKeeper.
   * @param path path to remove (i.e. /tmp will remove /tmp/1 and /tmp/2)
   * @param version expected version (-1 for all)
   * @param recursive if true, remove all children, otherwise behave like
   *        remove()
   * @throws InterruptedException
   * @throws KeeperException
   */
  public void deleteExt(final String path, int version, boolean recursive)
    throws InterruptedException, KeeperException {
    int attempt = 0;
    while (attempt < maxRetryAttempts) {
      try {
        if (!recursive) {
          zooKeeper.delete(path, version);
          return;
        }

        try {
          zooKeeper.delete(path, version);
          return;
        } catch (KeeperException.NotEmptyException e) {
          if (LOG.isDebugEnabled()) {
            LOG.debug("deleteExt: Cannot directly remove node " + path);
          }
        }

        List<String> childList = zooKeeper.getChildren(path, false);
        for (String child : childList) {
          if (progressable != null) {
            progressable.progress();
          }
          deleteExt(path + "/" + child, -1, true);
        }

        zooKeeper.delete(path, version);
        return;
      } catch (KeeperException.ConnectionLossException e) {
        LOG.warn("deleteExt: Connection loss on attempt " +
            attempt + ", waiting " + retryWaitMsecs +
            " msecs before retrying.", e);
      }
      ++attempt;
      Thread.sleep(retryWaitMsecs);
    }
    throw new IllegalStateException("deleteExt: Failed to delete " + path  +
        " after " + attempt + " tries!");
  }

  /**
   * Return the stat of the node of the given path. Return null if no such a
   * node exists.
   * <p>
   * If the watch is true and the call is successful (no exception is thrown),
   * a watch will be left on the node with the given path. The watch will be
   * triggered by a successful operation that creates/delete the node or sets
   * the data on the node.
   *
   * @param path
   *                the node path
   * @param watch
   *                whether need to watch this node
   * @return the stat of the node of the given path; return null if no such a
   *         node exists.
   * @throws KeeperException If the server signals an error
   * @throws InterruptedException If the server transaction is interrupted.
   */
  public Stat exists(String path, boolean watch) throws KeeperException,
      InterruptedException {
    int attempt = 0;
    while (attempt < maxRetryAttempts) {
      try {
        return zooKeeper.exists(path, watch);
      } catch (KeeperException.ConnectionLossException e) {
        LOG.warn("exists: Connection loss on attempt " +
            attempt + ", waiting " + retryWaitMsecs +
            " msecs before retrying.", e);
      }
      ++attempt;
      Thread.sleep(retryWaitMsecs);
    }
    throw new IllegalStateException("exists: Failed to check " + path  +
        " after " + attempt + " tries!");
  }

  /**
   * Return the stat of the node of the given path. Return null if no such a
   * node exists.
   * <p>
   * If the watch is non-null and the call is successful (no exception is
   * thrown), a watch will be left on the node with the given path. The
   * watch will be triggered by a successful operation that
   * creates/delete the node or sets the data on the node.
   *
   * @param path the node path
   * @param watcher explicit watcher
   * @return the stat of the node of the given path; return null if no such a
   *         node exists.
   * @throws KeeperException If the server signals an error
   * @throws InterruptedException If the server transaction is interrupted.
   * @throws IllegalArgumentException if an invalid path is specified
   */
  public Stat exists(final String path, Watcher watcher)
    throws KeeperException, InterruptedException {
    int attempt = 0;
    while (attempt < maxRetryAttempts) {
      try {
        return zooKeeper.exists(path, watcher);
      } catch (KeeperException.ConnectionLossException e) {
        LOG.warn("exists: Connection loss on attempt " +
            attempt + ", waiting " + retryWaitMsecs +
            " msecs before retrying.", e);
      }
      ++attempt;
      Thread.sleep(retryWaitMsecs);
    }
    throw new IllegalStateException("exists: Failed to check " + path  +
        " after " + attempt + " tries!");
  }

  /**
   * Return the data and the stat of the node of the given path.
   * <p>
   * If the watch is non-null and the call is successful (no exception is
   * thrown), a watch will be left on the node with the given path. The watch
   * will be triggered by a successful operation that sets data on the node, or
   * deletes the node.
   * <p>
   * A KeeperException with error code KeeperException.NoNode will be thrown
   * if no node with the given path exists.
   *
   * @param path the given path
   * @param watcher explicit watcher
   * @param stat the stat of the node
   * @return the data of the node
   * @throws KeeperException If the server signals an error with a non-zero
   *         error code
   * @throws InterruptedException If the server transaction is interrupted.
   * @throws IllegalArgumentException if an invalid path is specified
   */
  public byte[] getData(final String path, Watcher watcher, Stat stat)
    throws KeeperException, InterruptedException {
    int attempt = 0;
    while (attempt < maxRetryAttempts) {
      try {
        return zooKeeper.getData(path, watcher, stat);
      } catch (KeeperException.ConnectionLossException e) {
        LOG.warn("getData: Connection loss on attempt " +
            attempt + ", waiting " + retryWaitMsecs +
            " msecs before retrying.", e);
      }
      ++attempt;
      Thread.sleep(retryWaitMsecs);
    }
    throw new IllegalStateException("getData: Failed to get " + path  +
        " after " + attempt + " tries!");
  }

  /**
   * Return the data and the stat of the node of the given path.
   * <p>
   * If the watch is true and the call is successful (no exception is
   * thrown), a watch will be left on the node with the given path. The watch
   * will be triggered by a successful operation that sets data on the node, or
   * deletes the node.
   * <p>
   * A KeeperException with error code KeeperException.NoNode will be thrown
   * if no node with the given path exists.
   *
   * @param path the given path
   * @param watch whether need to watch this node
   * @param stat the stat of the node
   * @return the data of the node
   * @throws KeeperException If the server signals an error with a non-zero
   *         error code
   * @throws InterruptedException If the server transaction is interrupted.
   */
  public byte[] getData(String path, boolean watch, Stat stat)
    throws KeeperException, InterruptedException {
    int attempt = 0;
    while (attempt < maxRetryAttempts) {
      try {
        return zooKeeper.getData(path, watch, stat);
      } catch (KeeperException.ConnectionLossException e) {
        LOG.warn("getData: Connection loss on attempt " +
            attempt + ", waiting " + retryWaitMsecs +
            " msecs before retrying.", e);
      }
      ++attempt;
      Thread.sleep(retryWaitMsecs);
    }
    throw new IllegalStateException("getData: Failed to get " + path  +
        " after " + attempt + " tries!");
  }

  /**
   * Get the children of the path with extensions.
   * Extension 1: Sort the children based on sequence number
   * Extension 2: Get the full path instead of relative path
   *
   * @param path path to znode
   * @param watch set the watch?
   * @param sequenceSorted sort by the sequence number
   * @param fullPath if true, get the fully znode path back
   * @return list of children
   * @throws InterruptedException
   * @throws KeeperException
   */
  public List<String> getChildrenExt(
      final String path,
      boolean watch,
      boolean sequenceSorted,
      boolean fullPath) throws KeeperException, InterruptedException {
    int attempt = 0;
    while (attempt < maxRetryAttempts) {
      try {
        List<String> childList = zooKeeper.getChildren(path, watch);
        /* Sort children according to the sequence number, if desired */
        if (sequenceSorted) {
          Collections.sort(childList, new Comparator<String>() {
            public int compare(String s1, String s2) {
              if ((s1.length() <= SEQUENCE_NUMBER_LENGTH) ||
                  (s2.length() <= SEQUENCE_NUMBER_LENGTH)) {
                throw new RuntimeException(
                    "getChildrenExt: Invalid length for sequence " +
                        " sorting > " +
                        SEQUENCE_NUMBER_LENGTH +
                        " for s1 (" +
                        s1.length() + ") or s2 (" + s2.length() + ")");
              }
              int s1sequenceNumber = Integer.parseInt(
                  s1.substring(s1.length() -
                      SEQUENCE_NUMBER_LENGTH));
              int s2sequenceNumber = Integer.parseInt(
                  s2.substring(s2.length() -
                      SEQUENCE_NUMBER_LENGTH));
              return s1sequenceNumber - s2sequenceNumber;
            }
          });
        }
        if (fullPath) {
          List<String> fullChildList = new ArrayList<String>();
          for (String child : childList) {
            fullChildList.add(path + "/" + child);
          }
          return fullChildList;
        }
        return childList;
      } catch (KeeperException.ConnectionLossException e) {
        LOG.warn("getChildrenExt: Connection loss on attempt " +
            attempt + ", waiting " + retryWaitMsecs +
            " msecs before retrying.", e);
      }
      ++attempt;
      Thread.sleep(retryWaitMsecs);
    }
    throw new IllegalStateException("createExt: Failed to create " + path  +
        " after " + attempt + " tries!");
  }

  /**
   * Close this client object. Once the client is closed, its session becomes
   * invalid. All the ephemeral nodes in the ZooKeeper server associated with
   * the session will be removed. The watches left on those nodes (and on
   * their parents) will be triggered.
   *
   * @throws InterruptedException
   */
  public void close() throws InterruptedException {
    zooKeeper.close();
  }
}