package io.vertx.ext.jdbc;

import io.vertx.core.AbstractVerticle;
import io.vertx.core.DeploymentOptions;
import io.vertx.core.json.JsonObject;
import io.vertx.ext.jdbc.spi.DataSourceProvider;
import io.vertx.ext.jdbc.spi.impl.AgroalCPDataSourceProvider;
import io.vertx.ext.sql.SQLClient;
import io.vertx.ext.sql.SQLConnection;
import org.junit.BeforeClass;
import org.junit.Test;

import javax.sql.DataSource;
import java.util.List;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.TimeUnit;
import java.util.function.Predicate;
import java.util.stream.Collectors;
import java.util.stream.Stream;

import static java.util.concurrent.TimeUnit.*;

/**
 * @author <a href="mailto:julien@julienviet.com">Julien Viet</a>
 */
public class CloseTest extends JDBCClientTestBase {

  @BeforeClass
  public static void createDb() throws Exception {
    resetDb();
  }

  private static final JsonObject theConfig = DBConfigs.hsqldb();

  public static class NonSharedClientVerticle extends AbstractVerticle {
    @Override
    public void start(io.vertx.core.Promise<Void> f) throws Exception {
      SQLClient client = JDBCClient.create(vertx, theConfig);
      String sql = "SELECT ID, FNAME, LNAME FROM select_table ORDER BY ID";
      client.getConnection(ar1 -> {
        if (ar1.succeeded()) {
          SQLConnection conn = ar1.result();
          conn.query(sql, ar2 -> {
            if (ar2.succeeded()) {
              f.complete();
            } else {
              f.fail(ar2.cause());
            }
          });
        } else {
          f.fail(ar1.cause());
        }
      });
    }
  }

  @Test
  public void testUsingNonSharedInVerticle() throws Exception {
    CompletableFuture<String> id = new CompletableFuture<>();
    vertx.deployVerticle(NonSharedClientVerticle.class.getName(), onSuccess(id::complete));
    close(id.get(10, TimeUnit.SECONDS), false);
  }

  @Test
  public void testUsingNonSharedInVerticle2() throws Exception {
    CompletableFuture<String> id = new CompletableFuture<>();
    vertx.deployVerticle(NonSharedClientVerticle.class.getName(), new DeploymentOptions().setInstances(2), onSuccess(id::complete));
    close(id.get(10, TimeUnit.SECONDS), false);
  }

  public static class SharedClientVerticle extends AbstractVerticle {
    @Override
    public void start(io.vertx.core.Promise<Void> f) throws Exception {
      SQLClient client = JDBCClient.createShared(vertx, theConfig);
      String sql = "SELECT ID, FNAME, LNAME FROM select_table ORDER BY ID";
      client.getConnection(ar1 -> {
        if (ar1.succeeded()) {
          SQLConnection conn = ar1.result();
          conn.query(sql, ar2 -> {
            if (ar2.succeeded()) {
              f.complete();
            } else {
              f.fail(ar2.cause());
            }
          });
        } else {
          f.fail(ar1.cause());
        }
      });
    }
  }

  @Test
  public void testUsingSharedInVerticle() throws Exception {
    CompletableFuture<String> id = new CompletableFuture<>();
    vertx.deployVerticle(SharedClientVerticle.class.getName(), new DeploymentOptions().setInstances(1), onSuccess(id::complete));
    close(id.get(10, TimeUnit.SECONDS), false);
  }

  @Test
  public void testUsingSharedInVerticle2() throws Exception {
    CompletableFuture<String> id = new CompletableFuture<>();
    vertx.deployVerticle(SharedClientVerticle.class.getName(), new DeploymentOptions().setInstances(2), onSuccess(id::complete));
    close(id.get(10, TimeUnit.SECONDS), false);
  }

  private static DataSource ds;

  public static class ProvidedDataSourceVerticle extends AbstractVerticle {
    @Override
    public void start(io.vertx.core.Promise<Void> f) throws Exception {
      SQLClient client = JDBCClient.create(vertx, ds);
      String sql = "SELECT ID, FNAME, LNAME FROM select_table ORDER BY ID";
      client.getConnection(ar1 -> {
        if (ar1.succeeded()) {
          SQLConnection conn = ar1.result();
          conn.query(sql, ar2 -> {
            if (ar2.succeeded()) {
              f.complete();
            } else {
              f.fail(ar2.cause());
            }
          });
        } else {
          f.fail(ar1.cause());
        }
      });
    }
  }

  @Test
  public void testUsingProvidedDataSourceVerticle() throws Exception {
    DataSourceProvider provider = new AgroalCPDataSourceProvider();
    ds = provider.getDataSource(theConfig);
    CompletableFuture<String> id = new CompletableFuture<>();
    vertx.deployVerticle(ProvidedDataSourceVerticle.class.getName(), new DeploymentOptions().setInstances(1), onSuccess(id::complete));
    try {
      close(id.get(10, TimeUnit.SECONDS), true);
    } finally {
      provider.close(ds);
    }
  }

  private void close(String deploymentId, boolean expectedDsThreadStatus) throws Exception {
    List<Thread> getConnThread = findThreads(t -> t.getName().equals("vertx-jdbc-service-get-connection-thread"));
    assertTrue(getConnThread.size() > 0);
    List<Thread> poolThreads = findThreads(t -> t.getName().startsWith("AgroalPooledConnectionPoolManager"));
    assertTrue(poolThreads.size() > 0);
    CountDownLatch closeLatch = new CountDownLatch(1);
    vertx.undeploy(deploymentId, onSuccess(v -> {
      closeLatch.countDown();
    }));
    awaitLatch(closeLatch);
    long start = System.currentTimeMillis();
    while (System.currentTimeMillis() - start < 10000) {
      if (!getConnThread.get(0).isAlive() && poolThreads.stream().allMatch(t -> t.isAlive() == expectedDsThreadStatus)) {
        return;
      }
      MILLISECONDS.sleep(10);
    }
    String msg = Stream
      .concat(Stream.of(getConnThread.get(0)), poolThreads.stream())
      .map(t -> t.getName() + ": state=" + t.getState().name() + "/alive=" + t.isAlive())
      .collect(Collectors.joining(", ",
        "Timeout waiting for end condition:", "."));
    fail(msg);
  }

  private List<Thread> findThreads(Predicate<Thread> predicate) {
    return Thread.getAllStackTraces()
        .keySet()
        .stream()
        .filter(predicate)
        .collect(Collectors.toList());
  }
}
