/*
 * 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.hadoop.hive.llap.registry.impl;

import com.google.common.annotations.VisibleForTesting;
import com.google.common.base.Preconditions;
import org.apache.curator.framework.CuratorFramework;
import org.apache.curator.framework.api.BackgroundCallback;
import org.apache.curator.framework.api.CuratorEvent;
import org.apache.curator.framework.state.ConnectionState;
import org.apache.curator.framework.state.ConnectionStateListener;
import org.apache.zookeeper.CreateMode;
import org.apache.zookeeper.KeeperException;
import org.apache.zookeeper.WatchedEvent;
import org.apache.zookeeper.Watcher;
import org.apache.zookeeper.Watcher.Event.EventType;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;


import java.io.Closeable;
import java.io.IOException;
import java.nio.charset.Charset;
import java.nio.charset.StandardCharsets;
import java.util.Arrays;
import java.util.Iterator;
import java.util.List;
import java.util.Random;
import java.util.TreeSet;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicReference;

/**
 * We would have used the curator ephemeral node with some extra logic but it doesn't handle
 * the EXISTS condition, which is crucial here; so we c/p parts of Curator and add our logic.
 */
public class SlotZnode implements Closeable {
  static final Charset CHARSET = StandardCharsets.UTF_8;
  private final static Logger LOG = LoggerFactory.getLogger(SlotZnode.class);

  private final AtomicReference<CountDownLatch> initialCreateLatch =
      new AtomicReference<CountDownLatch>(new CountDownLatch(1));
  private final AtomicReference<String> nodePath = new AtomicReference<String>(null);
  private final Random rdm = new Random();
  private final CuratorFramework client;
  private final String basePath, prefix, workerPrefix, dataStr;
  private final byte[] data;
  private int currentSlot;
  private final AtomicReference<State> state = new AtomicReference<State>(State.LATENT);
  private int fallbackCount = 0; // Test-only counter.
  private final BackgroundCallback backgroundCallback = new BackgroundCallback() {
    @Override
    public void processResult(CuratorFramework client, CuratorEvent event) throws Exception {
      processCreateResult(client, event);
    }
  };
  private final Watcher watcher = new Watcher() {
    @Override
    public void process(WatchedEvent event) {
      processWatchedEvent(event);
    }
  };
  private final BackgroundCallback checkExistsCallback = new BackgroundCallback() {
    @Override
    public void processResult(CuratorFramework client, CuratorEvent event) throws Exception {
      processWatchResult(event);
    }
  };
  private final ConnectionStateListener connectionStateListener = new ConnectionStateListener() {
    @Override
    public void stateChanged(CuratorFramework client, ConnectionState newState) {
      processConnectionState(newState);
    }
  };

  private enum State {
    LATENT,
    INITIAL_SELECTION,
    AFTER_SELECTION,
    CLOSED
  }

  public SlotZnode(
      CuratorFramework client, String basePath, String prefix, String workerPrefix, String data) {
    this.client = Preconditions.checkNotNull(client, "client cannot be null");
    this.basePath = Preconditions.checkNotNull(basePath, "basePath cannot be null");
    this.prefix = Preconditions.checkNotNull(prefix, "prefix cannot be null");
    this.workerPrefix = workerPrefix;
    Preconditions.checkNotNull(data, "data cannot be null");
    this.dataStr = data;
    this.data = data.getBytes(CHARSET);
  }

  @VisibleForTesting
  public int getFallbackCount() {
    return fallbackCount;
  }

  private void chooseSlotToTake() throws Exception {
    int slotToTake = -1;
    while (true) {
      List<String> allChildNodes;
      try {
        allChildNodes = client.getChildren().forPath(basePath);
      } catch (Exception e) {
        LOG.error("Cannot list nodes to get slots; failing", e);
        throw e;
      }
      TreeSet<Integer> slots = new TreeSet<>();
      int approxWorkerCount = 0;
      for (String child : allChildNodes) {
        if (!child.startsWith(prefix)) {
          if (child.startsWith(workerPrefix)) {
            ++approxWorkerCount;
          }
        } else {
          slots.add(Integer.parseInt(child.substring(prefix.length())));
        }
      }
      Iterator<Integer> slotIter = slots.iterator();
      slotToTake = 0;
      while (slotIter.hasNext()) {
        int nextTaken = slotIter.next();
        if (slotToTake < nextTaken) break;
        slotToTake = nextTaken + 1;
      }
      // There may be a race for this slot so re-query after a delay with some probability.
      if (slotToTake != currentSlot || !shouldFallBackOnCollision(approxWorkerCount)) break;
      ++fallbackCount;
      Thread.sleep(rdm.nextInt(200)); // arbitrary
    }

    currentSlot = slotToTake;
    LOG.info("Will attempt to take slot " + currentSlot);
  }


  private boolean shouldFallBackOnCollision(int approxWorkerCount) {
    // Ideally, we'd want 1 worker to try for every slot; e.g. if there are 4 workers we want 3
    // to re-read, i.e. probability of falling back = 0.75, or 1/4 < random([0,1)). However, we
    // make it slightly more probable (2.0x) to avoid too much re-reading. This is hand-wavy.
    if (approxWorkerCount == 0) return false;
    return (2.0f / approxWorkerCount) <= rdm.nextDouble();
  }

  private String getSlotPath(int slot) {
    return String.format("%s/%s%010d", basePath, prefix, slot);
  }

  public boolean start(long timeout, TimeUnit unit) throws Exception {
    Preconditions.checkState(state.compareAndSet(State.LATENT, State.INITIAL_SELECTION), "Already started");
    CountDownLatch localLatch = initialCreateLatch.get();
    client.getConnectionStateListenable().addListener(connectionStateListener);
    chooseSlotToTake();
    startCreateCurrentNode();
    return localLatch.await(timeout, unit);
  }

  @Override
  public void close() throws IOException {
    State currentState = state.getAndSet(State.CLOSED);
    if (currentState == State.CLOSED || currentState == State.LATENT) return;
    client.getConnectionStateListenable().removeListener(connectionStateListener);
    String localNodePath = nodePath.getAndSet(null);
    if (localNodePath == null) return;
    try {
      client.delete().guaranteed().forPath(localNodePath);
    } catch (KeeperException.NoNodeException ignore) {
    } catch (Exception e) {
      LOG.error("Deleting node: " + localNodePath, e);
      throw new IOException(e);
    }
  }

  public int getCurrentSlot() {
    assert isActive();
    return currentSlot;
  }

  private void startCreateCurrentNode() {
    if (!isActive()) return;
    String createPath = null;
    try {
      createPath = getSlotPath(currentSlot);
      LOG.info("Attempting to create " + createPath);
      client.create().withMode(CreateMode.EPHEMERAL).inBackground(backgroundCallback)
        .forPath(createPath, data);
    } catch (Exception e) {
      LOG.error("Creating node. Path: " + createPath, e);
      throw new RuntimeException(e);
    }
  }

  private void watchNode() throws Exception {
    if (!isActive()) return;
    String localNodePath = nodePath.get();
    if (localNodePath == null) return;
    try {
      client.checkExists().usingWatcher(watcher).inBackground(
          checkExistsCallback).forPath(localNodePath);
    } catch (Exception e) {
      LOG.error("Watching node: " + localNodePath, e);
      throw e;
    }
  }

  private boolean isActive() {
    State localState = state.get();
    return (localState != State.LATENT && localState != State.CLOSED);
  }

  private void processWatchResult(CuratorEvent event) throws Exception {
    if (event.getResultCode() != KeeperException.Code.NONODE.intValue()) return;
    LOG.info("Trying to reacquire because of the NONODE event");
    startCreateCurrentNode();
  }


  private void processConnectionState(ConnectionState newState) {
    if (newState != ConnectionState.RECONNECTED) return;
    LOG.info("Trying to reacquire because of the RECONNECTED event");
    startCreateCurrentNode();
  }


  private void processWatchedEvent(WatchedEvent event) {
    if (event.getType() != EventType.NodeDeleted) return;
    String localPath = nodePath.get();
    if (localPath == null) return;
    if (!localPath.equals(event.getPath())) {
      LOG.info("Ignoring the NodeDeleted event for " + event.getPath());
      return;
    }
    LOG.info("Trying to reacquire because of the NodeDeleted event");
    startCreateCurrentNode();
  }


  private void processCreateResult(CuratorFramework client, CuratorEvent event) throws Exception {
    boolean doesExist = event.getResultCode() == KeeperException.Code.NODEEXISTS.intValue();
    if (!doesExist && event.getResultCode() != KeeperException.Code.OK.intValue()) {
      if (LOG.isInfoEnabled()) {
        LOG.info("Trying to reacquire due to create error: " + event);
      }
      startCreateCurrentNode(); // TODO: a pattern from Curator. Better error handling?
      return;
    }
    State localState = state.get();
    switch (localState) {
    case CLOSED:
    case LATENT:
      return;
    case INITIAL_SELECTION:
      if (doesExist) {
        LOG.info("Slot " + currentSlot + " was occupied");
        chooseSlotToTake(); // Try another slot.
        startCreateCurrentNode();
      } else {
        handleCreatedNode(event.getName());
      }
      break;
    case AFTER_SELECTION:
      if (doesExist) {
        processExistsFromCreate(client, event.getPath());
      } else {
        handleCreatedNode(event.getName());
      }
      break;
    default:
      throw new AssertionError("Unknown state " + localState);
    }
  }


  private void processExistsFromCreate(CuratorFramework client, String path) throws Exception {
    byte[] actual;
    try {
      actual = client.getData().forPath(path);
    } catch (Exception ex) {
      LOG.error("Error getting data for the node; will retry creating", ex);
      startCreateCurrentNode();
      return;
    }
    if (Arrays.equals(actual, data)) {
      handleCreatedNode(path);
    } else {
      if (LOG.isInfoEnabled()) {
        LOG.info("Data at {} is from a different node: {} (we are {})",
            path, new String(actual, CHARSET), dataStr);
      }
      nodePath.getAndSet(null);
      chooseSlotToTake(); // Try another slot.
      startCreateCurrentNode();
    }
  }

  private void handleCreatedNode(String path) throws Exception {
    while (true) {
      State localState = state.get();
      if (localState == State.CLOSED || localState == State.LATENT) return;
      if (state.compareAndSet(localState, State.AFTER_SELECTION)) break;
    }
    nodePath.set(path);
    watchNode();
    CountDownLatch localLatch = initialCreateLatch.getAndSet(null);
    if (localLatch != null) {
      localLatch.countDown();
    }
    LOG.info("Acquired the slot znode {}{}", path,
        localLatch != null ? "; this is the initial assignment" : "");
  }

  @VisibleForTesting
  public String getActualPath() {
    return nodePath.get();
  }
}
