/*
 * 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 com.datastax.driver.core;

import static org.assertj.core.api.Assertions.assertThat;
import static org.assertj.core.api.Assertions.fail;
import static org.testng.Assert.assertEquals;
import static org.testng.Assert.assertTrue;

import com.datastax.driver.core.exceptions.InvalidQueryException;
import com.datastax.driver.core.utils.CassandraVersion;
import com.google.common.base.Function;
import com.google.common.base.Throwables;
import com.google.common.util.concurrent.AsyncFunction;
import com.google.common.util.concurrent.ListenableFuture;
import com.google.common.util.concurrent.Uninterruptibles;
import java.nio.ByteBuffer;
import java.util.Collection;
import java.util.Map;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.Executor;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.TimeUnit;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.testng.annotations.DataProvider;
import org.testng.annotations.Test;

public class AsyncQueryTest extends CCMTestsSupport {

  Logger logger = LoggerFactory.getLogger(AsyncQueryTest.class);

  @DataProvider(name = "keyspace")
  public static Object[][] keyspace() {
    return new Object[][] {{"asyncquerytest"}, {"\"AsyncQueryTest\""}};
  }

  @Override
  public void onTestContextInitialized() {
    for (Object[] objects : keyspace()) {
      String keyspace = (String) objects[0];
      execute(
          String.format(
              "create keyspace %s WITH replication = {'class': 'SimpleStrategy', 'replication_factor': 1}",
              keyspace),
          String.format("create table %s.foo(k int, v int, primary key (k, v))", keyspace));
      for (int v = 1; v <= 100; v++)
        execute(String.format("insert into %s.foo (k, v) values (1, %d)", keyspace, v));
    }
  }

  /** Checks that a cancelled query releases the connection (JAVA-407). */
  @Test(groups = "short")
  public void cancelled_query_should_release_the_connection() throws InterruptedException {
    ResultSetFuture future = session().executeAsync("select release_version from system.local");
    future.cancel(true);
    assertTrue(future.isCancelled());

    TimeUnit.MILLISECONDS.sleep(100);

    HostConnectionPool pool = getPool(session());
    for (Connection connection : pool.connections) {
      assertEquals(connection.inFlight.get(), 0);
    }
  }

  @Test(groups = "short")
  public void should_init_cluster_and_session_if_needed() throws Exception {
    // For this test we need an uninitialized cluster, so we can't reuse the one provided by the
    // parent class. Rebuild a new one with the same (unique) host.
    Cluster cluster2 = register(createClusterBuilder().build());
    try {
      Session session2 = cluster2.newSession();

      // Neither cluster2 nor session2 are initialized at this point
      assertThat(cluster2.manager.metadata).isNull();

      ResultSetFuture future = session2.executeAsync("select release_version from system.local");
      Row row = Uninterruptibles.getUninterruptibly(future).one();

      assertThat(row.getString(0)).isNotEmpty();
    } finally {
      cluster2.close();
    }
  }

  @Test(
      groups = "short",
      dataProvider = "keyspace",
      enabled = false,
      description =
          "disabled because the blocking USE call in the current pool implementation makes it deadlock")
  public void should_chain_query_on_async_session_init_with_same_executor(String keyspace)
      throws Exception {
    ListenableFuture<Integer> resultFuture =
        connectAndQuery(keyspace, GuavaCompatibility.INSTANCE.sameThreadExecutor());

    Integer result = Uninterruptibles.getUninterruptibly(resultFuture);
    assertThat(result).isEqualTo(1);
  }

  @Test(groups = "short", dataProvider = "keyspace")
  public void should_chain_query_on_async_session_init_with_different_executor(String keyspace)
      throws Exception {
    ExecutorService executor = Executors.newFixedThreadPool(1);

    ListenableFuture<Integer> resultFuture = connectAndQuery(keyspace, executor);

    Integer result = Uninterruptibles.getUninterruptibly(resultFuture);
    assertThat(result).isEqualTo(1);

    executor.shutdownNow();
  }

  @Test(groups = "short")
  public void should_propagate_error_to_chained_query_if_session_init_fails() throws Exception {
    ListenableFuture<Integer> resultFuture =
        connectAndQuery("wrong_keyspace", GuavaCompatibility.INSTANCE.sameThreadExecutor());

    try {
      Uninterruptibles.getUninterruptibly(resultFuture);
    } catch (ExecutionException e) {
      assertThat(e.getCause())
          .isInstanceOf(InvalidQueryException.class)
          .hasMessage("Keyspace 'wrong_keyspace' does not exist");
    }
  }

  @Test(groups = "short")
  public void should_fail_when_synchronous_call_on_io_thread() throws Exception {
    for (int i = 0; i < 1000; i++) {
      ResultSetFuture f = session().executeAsync("select release_version from system.local");
      ListenableFuture<Thread> f2 =
          GuavaCompatibility.INSTANCE.transform(
              f,
              new Function<ResultSet, Thread>() {
                @Override
                public Thread apply(ResultSet input) {
                  session().execute("select release_version from system.local");
                  return Thread.currentThread();
                }
              });
      if (isFailed(
          f2, IllegalStateException.class, "Detected a synchronous call on an I/O thread")) {
        return;
      }
    }
    fail("callback was not executed on io thread in 1000 attempts, something may be wrong.");
  }

  @Test(groups = "short")
  public void should_fail_when_synchronous_call_on_io_thread_with_session_wrapper()
      throws Exception {
    final Session session = new SessionWrapper(session());
    for (int i = 0; i < 1000; i++) {
      ResultSetFuture f = session.executeAsync("select release_version from system.local");
      ListenableFuture<Thread> f2 =
          GuavaCompatibility.INSTANCE.transform(
              f,
              new Function<ResultSet, Thread>() {
                @Override
                public Thread apply(ResultSet input) {
                  session.execute("select release_version from system.local");
                  return Thread.currentThread();
                }
              });
      if (isFailed(
          f2, IllegalStateException.class, "Detected a synchronous call on an I/O thread")) {
        return;
      }
    }
    fail("callback was not executed on io thread in 1000 attempts, something may be wrong.");
  }

  @Test(groups = "short")
  @CassandraVersion(value = "2.0.0", description = "Paging is not supported until 2.0")
  public void should_fail_when_auto_paging_on_io_thread() throws Exception {
    for (int i = 0; i < 1000; i++) {
      Statement statement = new SimpleStatement("select v from asyncquerytest.foo where k = 1");
      // Ensure results will be paged (there are 100 rows in the test table)
      statement.setFetchSize(10);
      ResultSetFuture f = session().executeAsync(statement);
      ListenableFuture<Thread> f2 =
          GuavaCompatibility.INSTANCE.transform(
              f,
              new Function<ResultSet, Thread>() {
                @Override
                public Thread apply(ResultSet rs) {
                  rs.all(); // Consume the whole result set
                  return Thread.currentThread();
                }
              });
      if (isFailed(
          f2, IllegalStateException.class, "Detected a synchronous call on an I/O thread")) {
        return;
      }
    }
    fail("callback was not executed on io thread in 1000 attempts, something may be wrong.");
  }

  private boolean isFailed(
      ListenableFuture<Thread> future, Class<?> expectedException, String expectedMessage) {
    try {
      Thread executedThread = future.get();
      if (executedThread != Thread.currentThread()) {
        fail("Expected a failed future, callback was executed on " + executedThread);
      } else {
        // Callback was invoked on the same thread, which indicates that the future completed
        // before the transform callback was registered. Try again to produce case where callback
        // is called on io thread.
        logger.warn("Future completed before transform callback registered, will try again.");
      }
    } catch (Exception e) {
      assertThat(Throwables.getRootCause(e))
          .isInstanceOf(expectedException)
          .hasMessageContaining(expectedMessage);
      return true;
    }
    return false;
  }

  private ListenableFuture<Integer> connectAndQuery(String keyspace, Executor executor) {
    ListenableFuture<Session> sessionFuture = cluster().connectAsync(keyspace);
    ListenableFuture<ResultSet> queryFuture =
        GuavaCompatibility.INSTANCE.transformAsync(
            sessionFuture,
            new AsyncFunction<Session, ResultSet>() {
              @Override
              public ListenableFuture<ResultSet> apply(Session session) throws Exception {
                return session.executeAsync("select v from foo where k = 1");
              }
            },
            executor);
    return GuavaCompatibility.INSTANCE.transform(
        queryFuture,
        new Function<ResultSet, Integer>() {
          @Override
          public Integer apply(ResultSet rs) {
            return rs.one().getInt(0);
          }
        },
        executor);
  }

  private static HostConnectionPool getPool(Session session) {
    Collection<HostConnectionPool> pools = ((SessionManager) session).pools.values();
    assertEquals(pools.size(), 1);
    return pools.iterator().next();
  }

  private static class SessionWrapper extends AbstractSession {

    private final Session session;

    public SessionWrapper(Session session) {
      this.session = session;
    }

    @Override
    public ResultSet execute(Statement statement) {
      // test a custom call to checkNotInEventLoop()
      checkNotInEventLoop();
      return executeAsync(statement).getUninterruptibly();
    }

    @Override
    public String getLoggedKeyspace() {
      return session.getLoggedKeyspace();
    }

    @Override
    public Session init() {
      return session.init();
    }

    @Override
    public ListenableFuture<Session> initAsync() {
      return session.initAsync();
    }

    @Override
    public ResultSetFuture executeAsync(Statement statement) {
      return session.executeAsync(statement);
    }

    @Override
    public CloseFuture closeAsync() {
      return session.closeAsync();
    }

    @Override
    public boolean isClosed() {
      return session.isClosed();
    }

    @Override
    public Cluster getCluster() {
      return session.getCluster();
    }

    @Override
    public State getState() {
      return session.getState();
    }

    @Override
    protected ListenableFuture<PreparedStatement> prepareAsync(
        String query, Map<String, ByteBuffer> customPayload) {
      return ((SessionManager) session).prepareAsync(query, customPayload);
    }
  }
}
