package io.vertx.servicediscovery.backend.zookeeper;

import io.vertx.core.*;
import io.vertx.core.json.JsonObject;
import io.vertx.servicediscovery.Record;
import io.vertx.servicediscovery.spi.ServiceDiscoveryBackend;
import org.apache.curator.framework.CuratorFramework;
import org.apache.curator.framework.CuratorFrameworkFactory;
import org.apache.curator.framework.api.CuratorEvent;
import org.apache.curator.framework.api.DeleteBuilder;
import org.apache.curator.framework.state.ConnectionState;
import org.apache.curator.framework.state.ConnectionStateListener;
import org.apache.curator.retry.ExponentialBackoffRetry;
import org.apache.zookeeper.CreateMode;
import org.apache.zookeeper.KeeperException;

import java.nio.charset.Charset;
import java.util.ArrayList;
import java.util.List;
import java.util.Objects;
import java.util.UUID;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException;

Author:
Clement Escoffier/**
 * @author <a href="http://escoffier.me">Clement Escoffier</a>
 */
public class ZookeeperBackendService implements ServiceDiscoveryBackend, ConnectionStateListener {

  private static final Charset CHARSET = Charset.forName("UTF-8");
  private CuratorFramework client;
  private String basePath;
  private boolean ephemeral;
  private boolean guaranteed;
  private Vertx vertx;
  private ConnectionState connectionState = ConnectionState.LOST;
  private boolean canBeReadOnly;
  private int connectionTimeoutMs;

  @Override
  public void init(Vertx vertx, JsonObject configuration) {
    this.vertx = vertx;
    String connection = Objects.requireNonNull(configuration.getString("connection"));
    int maxRetries = configuration.getInteger("maxRetries", 3);
    int baseGraceBetweenRetries = configuration
        .getInteger("baseSleepTimeBetweenRetries", 1000);
    canBeReadOnly = configuration.getBoolean("canBeReadOnly", false);
    connectionTimeoutMs = configuration.getInteger("connectionTimeoutMs", 1000);
    basePath = configuration.getString("basePath", "/services");
    ephemeral = configuration.getBoolean("ephemeral", false);
    guaranteed = configuration.getBoolean("guaranteed", false);

    client = CuratorFrameworkFactory.builder()
        .canBeReadOnly(canBeReadOnly)
        .connectString(connection)
        .connectionTimeoutMs(connectionTimeoutMs)
        .retryPolicy(new ExponentialBackoffRetry(baseGraceBetweenRetries, maxRetries))
        .build();
    client.getConnectionStateListenable().addListener(this);
    client.start();
  }

  @Override
  public void store(Record record, Handler<AsyncResult<Record>> resultHandler) {
    if (record.getRegistration() != null) {
      resultHandler.handle(Future.failedFuture("The record has already been registered"));
      return;
    }
    String uuid = UUID.randomUUID().toString();
    record.setRegistration(uuid);

    String content = record.toJson().encode();
    Context context = Vertx.currentContext();

    ensureConnected(x -> {
      if (x.failed()) {
        resultHandler.handle(Future.failedFuture(x.cause()));
      } else {
        try {
          client.create()
              .creatingParentsIfNeeded()
              .withMode(ephemeral ? CreateMode.EPHEMERAL : CreateMode.PERSISTENT)
              .inBackground((curatorFramework, curatorEvent)
                  -> callback(context, record, resultHandler, curatorEvent))
              .withUnhandledErrorListener((s, throwable)
                  -> resultHandler.handle(Future.failedFuture(throwable)))
              .forPath(getPath(record.getRegistration()), content.getBytes(CHARSET));
        } catch (Exception e) {
          resultHandler.handle(Future.failedFuture(e));
        }
      }
    });
  }

  private String getPath(String registration) {
    return basePath + "/" + registration;
  }

  private void callback(Context context, Record record, Handler<AsyncResult<Record>> resultHandler, CuratorEvent curatorEvent) {
    runOnContextIfPossible(context, () -> {
      if (curatorEvent.getResultCode() == KeeperException.Code.OK.intValue()) {
        resultHandler.handle(Future.succeededFuture(record));
      } else {
        KeeperException.Code code =
            KeeperException.Code.get(curatorEvent.getResultCode());
        resultHandler.handle(Future.failedFuture(KeeperException.create(code)));
      }
    });
  }

  @Override
  public void remove(Record record, Handler<AsyncResult<Record>> resultHandler) {
    Objects.requireNonNull(record.getRegistration(), "No registration id in the record");
    remove(record.getRegistration(), resultHandler);
  }

  @Override
  public void remove(String uuid, Handler<AsyncResult<Record>> resultHandler) {
    Objects.requireNonNull(uuid, "No registration id in the record");
    Context context = Vertx.currentContext();

    ensureConnected(x -> {
      if (x.failed()) {
        resultHandler.handle(Future.failedFuture(x.cause()));
      } else {
        getRecordById(context, uuid, record -> {
          if (record == null) {
            resultHandler.handle(Future.failedFuture("Unknown registration " + uuid));
          } else {
            try {
              DeleteBuilder delete = client.delete();
              if (guaranteed) {
                delete.guaranteed();
              }
              delete
                  .deletingChildrenIfNeeded()
                  .inBackground((curatorFramework, curatorEvent)
                      -> callback(context, record, resultHandler, curatorEvent))

                  .withUnhandledErrorListener((s, throwable)
                      -> resultHandler.handle(Future.failedFuture(throwable)))

                  .forPath(getPath(uuid));
            } catch (Exception e) {
              resultHandler.handle(Future.failedFuture(e));
            }
          }
        });
      }
    });
  }

  private void getRecordById(Context context, String uuid, Handler<Record> handler) {
    ensureConnected(x -> {
      if (x.failed()) {
        handler.handle(null);
      } else {
        try {
          client.getData()
              .inBackground((curatorFramework, curatorEvent)
                  -> runOnContextIfPossible(context, () -> {
                if (curatorEvent.getResultCode() == KeeperException.Code.OK.intValue()) {
                  JsonObject json
                      = new JsonObject(new String(curatorEvent.getData(), CHARSET));
                  handler.handle(new Record(json));
                } else {
                  handler.handle(null);
                }
              }))
              .forPath(getPath(uuid));
        } catch (Exception e) {
          handler.handle(null);
        }
      }
    });

  }

  private void runOnContextIfPossible(Context context, Runnable runnable) {
    if (context != null) {
      context.runOnContext(v -> runnable.run());
    } else {
      runnable.run();
    }
  }

  @Override
  public void update(Record record, Handler<AsyncResult<Void>> resultHandler) {
    Objects.requireNonNull(record.getRegistration(), "No registration id in the record");
    Context context = Vertx.currentContext();
    ensureConnected(x -> {
      if (x.failed()) {
        resultHandler.handle(Future.failedFuture(x.cause()));
      } else {
        try {
          client.setData()
              .inBackground((framework, event)
                  -> runOnContextIfPossible(context, () -> {
                if (event.getResultCode() == KeeperException.Code.OK.intValue()) {
                  resultHandler.handle(Future.succeededFuture());
                } else {
                  KeeperException.Code code = KeeperException.Code.get(event.getResultCode());
                  resultHandler.handle(Future.failedFuture(KeeperException.create(code)));
                }
              }))
              .withUnhandledErrorListener((message, e) -> resultHandler.handle(Future.failedFuture(e)))
              .forPath(getPath(record.getRegistration()),
                  record.toJson().encode().getBytes(CHARSET));
        } catch (Exception e) {
          resultHandler.handle(Future.failedFuture(e));
        }
      }
    });
  }

  @Override
  public void getRecords(Handler<AsyncResult<List<Record>>> resultHandler) {
    Context context = Vertx.currentContext();
    ensureConnected(
        x -> {
          if (x.failed()) {
            resultHandler.handle(Future.failedFuture(x.cause()));
          } else {
            try {
              client.getChildren()
                  .inBackground((fmk, event) -> {
                    List<String> children = event.getChildren();
                    List<Future> futures = new ArrayList<>();
                    for (String child : children) {
                      Promise<Record> promise = Promise.promise();
                      getRecord(child, promise);
                      futures.add(promise.future());
                    }

                    CompositeFuture.all(futures)
                        .setHandler(
                            ar -> runOnContextIfPossible(context, () -> {
                              if (ar.failed()) {
                                resultHandler.handle(Future.failedFuture(ar.cause()));
                              } else {
                                List<Record> records = new ArrayList<>();
                                for (Future future : futures) {
                                  records.add((Record) future.result());
                                }
                                resultHandler.handle(Future.succeededFuture(records));
                              }
                            }));
                  })
                  .withUnhandledErrorListener((message, e) -> resultHandler.handle(Future.failedFuture(e)))
                  .forPath(basePath);
            } catch (Exception e) {
              resultHandler.handle(Future.failedFuture(e));
            }
          }
        }
    );

  }

  @Override
  public void getRecord(String uuid, Handler<AsyncResult<Record>> handler) {
    Objects.requireNonNull(uuid);
    Context context = Vertx.currentContext();

    ensureConnected(x -> {
      if (x.failed()) {
        handler.handle(Future.failedFuture(x.cause()));
      } else {
        try {
          client.getData()
              .inBackground((fmk, curatorEvent)
                  -> runOnContextIfPossible(context, () -> {
                if (curatorEvent.getResultCode() == KeeperException.Code.OK.intValue()) {
                  JsonObject json
                      = new JsonObject(new String(curatorEvent.getData(), CHARSET));
                  handler.handle(Future.succeededFuture(new Record(json)));
                } else if (curatorEvent.getResultCode() == KeeperException.Code.NONODE.intValue()) {
                  handler.handle(Future.succeededFuture(null));
                } else {
                  KeeperException.Code code = KeeperException.Code.get(curatorEvent.getResultCode());
                  handler.handle(Future.failedFuture(KeeperException.create(code)));
                }
              }))
              .withUnhandledErrorListener((message, e) -> handler.handle(Future.failedFuture(e)))
              .forPath(getPath(uuid));
        } catch (Exception e) {
          handler.handle(Future.failedFuture(e));
        }
      }
    });
  }

  @Override
  public synchronized void stateChanged(CuratorFramework client, ConnectionState newState) {
    connectionState = newState;
  }

  private synchronized void ensureConnected(Handler<AsyncResult<Void>> handler) {
    switch (connectionState) {
      case CONNECTED:
      case RECONNECTED:
        handler.handle(Future.succeededFuture());
        break;
      case READ_ONLY:
        if (canBeReadOnly) {
          handler.handle(Future.succeededFuture());
          break;
        } // Else attempt to reconnect
      case LOST:
      case SUSPENDED:
        vertx.executeBlocking(
            future -> {
              try {
                if (client.blockUntilConnected(connectionTimeoutMs, TimeUnit.MILLISECONDS)) {
                  future.complete();
                } else {
                  future.fail(new TimeoutException());
                }
              } catch (Exception e) {
                future.fail(e);
              }
            }, ar -> {
              if (ar.failed()) {
                handler.handle(Future.failedFuture(KeeperException.create(KeeperException.Code.CONNECTIONLOSS)));
              } else {
                handler.handle(Future.succeededFuture());
              }
            });
        break;
    }

  }
}