/*
 * 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.oss.driver.api.testinfra.session;

import com.datastax.dse.driver.api.core.graph.ScriptGraphStatement;
import com.datastax.oss.driver.api.core.CqlIdentifier;
import com.datastax.oss.driver.api.core.Version;
import com.datastax.oss.driver.api.core.config.DriverConfigLoader;
import com.datastax.oss.driver.api.core.config.DriverExecutionProfile;
import com.datastax.oss.driver.api.core.cql.SimpleStatement;
import com.datastax.oss.driver.api.core.cql.Statement;
import com.datastax.oss.driver.api.core.metadata.NodeStateListener;
import com.datastax.oss.driver.api.core.metadata.schema.SchemaChangeListener;
import com.datastax.oss.driver.api.core.session.Session;
import com.datastax.oss.driver.api.testinfra.CassandraResourceRule;
import com.datastax.oss.driver.api.testinfra.ccm.BaseCcmRule;
import com.datastax.oss.driver.api.testinfra.ccm.CcmBridge;
import com.datastax.oss.driver.api.testinfra.ccm.SchemaChangeSynchronizer;
import com.datastax.oss.driver.api.testinfra.requirement.BackendType;
import com.datastax.oss.driver.api.testinfra.simulacron.SimulacronRule;
import java.util.Objects;
import org.junit.rules.ExternalResource;

/**
 * Creates and manages a {@link Session} instance for a test.
 *
 * <p>Use it in conjunction with a {@link CassandraResourceRule} that creates the server resource to
 * connect to:
 *
 * <pre>{@code
 * public static @ClassRule CcmRule server = CcmRule.getInstance();
 *
 * // Or: public static @ClassRule SimulacronRule server =
 * //    new SimulacronRule(ClusterSpec.builder().withNodes(3));
 *
 * public static @ClassRule SessionRule sessionRule = new SessionRule(server);
 *
 * public void @Test should_do_something() {
 *   sessionRule.session().execute("some query");
 * }
 * }</pre>
 *
 * Optionally, it can also create a dedicated keyspace (useful to isolate tests that share a common
 * server).
 *
 * <p>If you would rather create a new keyspace manually in each test, see the utility methods in
 * {@link SessionUtils}.
 */
public class SessionRule<SessionT extends Session> extends ExternalResource {

  private static final Version V6_8_0 = Objects.requireNonNull(Version.parse("6.8.0"));

  // the CCM or Simulacron rule to depend on
  private final CassandraResourceRule cassandraResource;
  private final NodeStateListener nodeStateListener;
  private final SchemaChangeListener schemaChangeListener;
  private final CqlIdentifier keyspace;
  private final DriverConfigLoader configLoader;
  private final String graphName;
  private final boolean isCoreGraph;

  // the session that is auto created for this rule and is tied to the given keyspace.
  private SessionT session;

  private DriverExecutionProfile slowProfile;

  /**
   * Returns a builder to construct an instance with a fluent API.
   *
   * @param cassandraResource resource to create clusters for.
   */
  public static CqlSessionRuleBuilder builder(CassandraResourceRule cassandraResource) {
    return new CqlSessionRuleBuilder(cassandraResource);
  }

  /** @see #builder(CassandraResourceRule) */
  public SessionRule(
      CassandraResourceRule cassandraResource,
      boolean createKeyspace,
      NodeStateListener nodeStateListener,
      SchemaChangeListener schemaChangeListener,
      DriverConfigLoader configLoader,
      String graphName,
      boolean isCoreGraph) {
    this.cassandraResource = cassandraResource;
    this.nodeStateListener = nodeStateListener;
    this.schemaChangeListener = schemaChangeListener;
    this.keyspace =
        (cassandraResource instanceof SimulacronRule || !createKeyspace)
            ? null
            : SessionUtils.uniqueKeyspaceId();
    this.configLoader = configLoader;
    this.graphName = graphName;
    this.isCoreGraph = isCoreGraph;
  }

  public SessionRule(
      CassandraResourceRule cassandraResource,
      boolean createKeyspace,
      NodeStateListener nodeStateListener,
      SchemaChangeListener schemaChangeListener,
      DriverConfigLoader configLoader,
      String graphName) {
    this(
        cassandraResource,
        createKeyspace,
        nodeStateListener,
        schemaChangeListener,
        configLoader,
        graphName,
        false);
  }

  public SessionRule(
      CassandraResourceRule cassandraResource,
      boolean createKeyspace,
      NodeStateListener nodeStateListener,
      SchemaChangeListener schemaChangeListener,
      DriverConfigLoader configLoader) {
    this(
        cassandraResource,
        createKeyspace,
        nodeStateListener,
        schemaChangeListener,
        configLoader,
        null,
        false);
  }

  @Override
  protected void before() {
    session =
        SessionUtils.newSession(
            cassandraResource, null, nodeStateListener, schemaChangeListener, null, configLoader);
    slowProfile = SessionUtils.slowProfile(session);
    if (keyspace != null) {
      SessionUtils.createKeyspace(session, keyspace, slowProfile);
      session.execute(
          SimpleStatement.newInstance(String.format("USE %s", keyspace.asCql(false))),
          Statement.SYNC);
    }
    if (graphName != null) {
      BaseCcmRule rule =
          (cassandraResource instanceof BaseCcmRule) ? ((BaseCcmRule) cassandraResource) : null;
      if (rule == null || !CcmBridge.isDistributionOf(BackendType.DSE)) {
        throw new IllegalArgumentException("DseSessionRule should work with DSE.");
      }
      if (rule.getDistributionVersion().compareTo(V6_8_0) >= 0) {
        session()
            .execute(
                ScriptGraphStatement.newInstance(
                        String.format(
                            "system.graph('%s').ifNotExists()%s.create()",
                            this.graphName, isCoreGraph ? ".coreEngine()" : ".classicEngine()"))
                    .setSystemQuery(true),
                ScriptGraphStatement.SYNC);
      } else {
        if (isCoreGraph) {
          throw new IllegalArgumentException(
              "Core graph is not supported for DSE version < " + V6_8_0);
        }
        session()
            .execute(
                ScriptGraphStatement.newInstance(
                        String.format("system.graph('%s').ifNotExists().create()", this.graphName))
                    .setSystemQuery(true),
                ScriptGraphStatement.SYNC);
      }
    }
  }

  @Override
  protected void after() {
    if (graphName != null) {
      session()
          .execute(
              ScriptGraphStatement.newInstance(
                      String.format("system.graph('%s').drop()", this.graphName))
                  .setSystemQuery(true),
              ScriptGraphStatement.SYNC);
    }
    if (keyspace != null) {
      SchemaChangeSynchronizer.withLock(
          () -> {
            SessionUtils.dropKeyspace(session, keyspace, slowProfile);
          });
    }
    session.close();
  }

  /** @return the session created with this rule. */
  public SessionT session() {
    return session;
  }

  /**
   * @return the identifier of the keyspace associated with this rule, or {@code null} if no
   *     keyspace was created (this is always the case if the server resource is a {@link
   *     SimulacronRule}).
   */
  public CqlIdentifier keyspace() {
    return keyspace;
  }

  public String getGraphName() {
    return graphName;
  }

  /** @return a config profile where the request timeout is 30 seconds. * */
  public DriverExecutionProfile slowProfile() {
    return slowProfile;
  }
}
