/*
 * 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.core.cql;

import static org.assertj.core.api.Assertions.assertThat;

import com.codahale.metrics.Gauge;
import com.datastax.oss.driver.api.core.CqlSession;
import com.datastax.oss.driver.api.core.config.DefaultDriverOption;
import com.datastax.oss.driver.api.core.config.DriverConfigLoader;
import com.datastax.oss.driver.api.core.context.DriverContext;
import com.datastax.oss.driver.api.core.cql.PreparedStatement;
import com.datastax.oss.driver.api.core.metrics.DefaultSessionMetric;
import com.datastax.oss.driver.api.core.session.ProgrammaticArguments;
import com.datastax.oss.driver.api.core.session.SessionBuilder;
import com.datastax.oss.driver.api.testinfra.ccm.CustomCcmRule;
import com.datastax.oss.driver.api.testinfra.ccm.SchemaChangeSynchronizer;
import com.datastax.oss.driver.api.testinfra.session.SessionRule;
import com.datastax.oss.driver.api.testinfra.session.SessionUtils;
import com.datastax.oss.driver.categories.IsolatedTests;
import com.datastax.oss.driver.internal.core.context.DefaultDriverContext;
import com.datastax.oss.driver.internal.core.cql.CqlPrepareAsyncProcessor;
import com.datastax.oss.driver.internal.core.cql.CqlPrepareSyncProcessor;
import com.datastax.oss.driver.internal.core.metadata.schema.events.TypeChangeEvent;
import com.datastax.oss.driver.internal.core.session.BuiltInRequestProcessors;
import com.datastax.oss.driver.internal.core.session.RequestProcessor;
import com.datastax.oss.driver.internal.core.session.RequestProcessorRegistry;
import com.datastax.oss.driver.shaded.guava.common.cache.RemovalListener;
import com.datastax.oss.driver.shaded.guava.common.util.concurrent.Uninterruptibles;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableSet;
import edu.umd.cs.findbugs.annotations.NonNull;
import java.nio.ByteBuffer;
import java.time.Duration;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.Optional;
import java.util.Set;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicReference;
import java.util.function.Consumer;
import org.junit.AfterClass;
import org.junit.BeforeClass;
import org.junit.Rule;
import org.junit.Test;
import org.junit.experimental.categories.Category;
import org.junit.rules.RuleChain;
import org.junit.rules.TestRule;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

// These tests must be isolated because setup modifies SessionUtils.SESSION_BUILDER_CLASS_PROPERTY
@Category(IsolatedTests.class)
public class PreparedStatementCachingIT {

  private CustomCcmRule ccmRule = CustomCcmRule.builder().build();

  private SessionRule<CqlSession> sessionRule =
      SessionRule.builder(ccmRule)
          .withConfigLoader(
              SessionUtils.configLoaderBuilder()
                  .withInt(DefaultDriverOption.REQUEST_PAGE_SIZE, 2)
                  .withDuration(DefaultDriverOption.REQUEST_TIMEOUT, Duration.ofSeconds(30))
                  .build())
          .build();

  @Rule public TestRule chain = RuleChain.outerRule(ccmRule).around(sessionRule);

  private static class PreparedStatementRemovalEvent {

    private final ByteBuffer queryId;

    public PreparedStatementRemovalEvent(ByteBuffer queryId) {
      this.queryId = queryId;
    }

    @Override
    public boolean equals(Object o) {
      if (this == o) return true;
      if (o == null || !(o instanceof PreparedStatementRemovalEvent)) return false;
      PreparedStatementRemovalEvent that = (PreparedStatementRemovalEvent) o;
      return Objects.equals(queryId, that.queryId);
    }

    @Override
    public int hashCode() {
      return Objects.hash(queryId);
    }

    @Override
    public String toString() {
      return "PreparedStatementRemovalEvent{" + "queryId=" + queryId + '}';
    }
  }

  private static class TestCqlPrepareAsyncProcessor extends CqlPrepareAsyncProcessor {

    private static final Logger LOG =
        LoggerFactory.getLogger(PreparedStatementCachingIT.TestCqlPrepareAsyncProcessor.class);

    private static RemovalListener<Object, Object> buildCacheRemoveCallback(
        @NonNull Optional<DefaultDriverContext> context) {
      return (evt) -> {
        try {
          CompletableFuture<PreparedStatement> future =
              (CompletableFuture<PreparedStatement>) evt.getValue();
          ByteBuffer queryId = Uninterruptibles.getUninterruptibly(future).getId();
          context.ifPresent(
              ctx -> ctx.getEventBus().fire(new PreparedStatementRemovalEvent(queryId)));
        } catch (Exception e) {
          LOG.error("Unable to register removal handler", e);
        }
      };
    }

    public TestCqlPrepareAsyncProcessor(@NonNull Optional<DefaultDriverContext> context) {
      // Default CqlPrepareAsyncProcessor uses weak values here as well.  We avoid doing so
      // to prevent cache entries from unexpectedly disappearing mid-test.
      super(context, builder -> builder.removalListener(buildCacheRemoveCallback(context)));
    }
  }

  private static class TestDefaultDriverContext extends DefaultDriverContext {
    public TestDefaultDriverContext(
        DriverConfigLoader configLoader, ProgrammaticArguments programmaticArguments) {
      super(configLoader, programmaticArguments);
    }

    @Override
    protected RequestProcessorRegistry buildRequestProcessorRegistry() {
      // Re-create the processor cache to insert the TestCqlPrepareAsyncProcessor with it's strong
      // prepared statement cache, see JAVA-3062
      List<RequestProcessor<?, ?>> processors =
          BuiltInRequestProcessors.createDefaultProcessors(this);
      processors.removeIf((processor) -> processor instanceof CqlPrepareAsyncProcessor);
      processors.removeIf((processor) -> processor instanceof CqlPrepareSyncProcessor);
      CqlPrepareAsyncProcessor asyncProcessor = new TestCqlPrepareAsyncProcessor(Optional.of(this));
      processors.add(2, asyncProcessor);
      processors.add(3, new CqlPrepareSyncProcessor(asyncProcessor));
      return new RequestProcessorRegistry(
          getSessionName(), processors.toArray(new RequestProcessor[0]));
    }
  }

  private static class TestSessionBuilder extends SessionBuilder {

    @Override
    protected Object wrap(@NonNull CqlSession defaultSession) {
      return defaultSession;
    }

    @Override
    protected DriverContext buildContext(
        DriverConfigLoader configLoader, ProgrammaticArguments programmaticArguments) {
      return new TestDefaultDriverContext(configLoader, programmaticArguments);
    }
  }

  @BeforeClass
  public static void setup() {
    System.setProperty(
        SessionUtils.SESSION_BUILDER_CLASS_PROPERTY, PreparedStatementCachingIT.class.getName());
  }

  @AfterClass
  public static void teardown() {
    System.clearProperty(SessionUtils.SESSION_BUILDER_CLASS_PROPERTY);
  }

  public static SessionBuilder builder() {
    return new TestSessionBuilder();
  }

  private void invalidationResultSetTest(
      Consumer<CqlSession> setupTestSchema, Set<String> expectedChangedTypes) {
    invalidationTestInner(
        setupTestSchema,
        "select f from test_table_caching_1 where e = ?",
        "select h from test_table_caching_2 where g = ?",
        expectedChangedTypes);
  }

  private void invalidationVariableDefsTest(
      Consumer<CqlSession> setupTestSchema,
      boolean isCollection,
      Set<String> expectedChangedTypes) {
    String condition = isCollection ? "contains ?" : "= ?";
    invalidationTestInner(
        setupTestSchema,
        String.format("select e from test_table_caching_1 where f %s allow filtering", condition),
        String.format("select g from test_table_caching_2 where h %s allow filtering", condition),
        expectedChangedTypes);
  }

  private void invalidationTestInner(
      Consumer<CqlSession> setupTestSchema,
      String preparedStmtQueryType1,
      String preparedStmtQueryType2,
      Set<String> expectedChangedTypes) {

    try (CqlSession session = sessionWithCacheSizeMetric()) {

      assertThat(getPreparedCacheSize(session)).isEqualTo(0);
      setupTestSchema.accept(session);

      session.prepare(preparedStmtQueryType1);
      ByteBuffer queryId2 = session.prepare(preparedStmtQueryType2).getId();
      assertThat(getPreparedCacheSize(session)).isEqualTo(2);

      CountDownLatch preparedStmtCacheRemoveLatch = new CountDownLatch(1);
      CountDownLatch typeChangeEventLatch = new CountDownLatch(expectedChangedTypes.size());

      DefaultDriverContext ctx = (DefaultDriverContext) session.getContext();
      Map<String, Boolean> changedTypes = new ConcurrentHashMap<>();
      AtomicReference<Optional<ByteBuffer>> removedQueryIds =
          new AtomicReference<>(Optional.empty());
      AtomicReference<Optional<String>> typeChangeEventError =
          new AtomicReference<>(Optional.empty());
      AtomicReference<Optional<String>> removedQueryEventError =
          new AtomicReference<>(Optional.empty());
      ctx.getEventBus()
          .register(
              TypeChangeEvent.class,
              (e) -> {
                // expect one event per type changed and for every parent type that nests it
                if (Boolean.TRUE.equals(
                    changedTypes.putIfAbsent(e.oldType.getName().toString(), true))) {
                  // store an error if we see duplicate change event
                  // any non-empty error will fail the test so it's OK to do this multiple times
                  typeChangeEventError.set(Optional.of("Duplicate type change event " + e));
                }
                typeChangeEventLatch.countDown();
              });
      ctx.getEventBus()
          .register(
              PreparedStatementRemovalEvent.class,
              (e) -> {
                if (!removedQueryIds.compareAndSet(Optional.empty(), Optional.of(e.queryId))) {
                  // store an error if we see multiple cache invalidation events
                  // any non-empty error will fail the test so it's OK to do this multiple times
                  removedQueryEventError.set(
                      Optional.of("Unable to set reference for PS removal event"));
                }
                preparedStmtCacheRemoveLatch.countDown();
              });

      // alter test_type_caching_2 to trigger cache invalidation and above events
      session.execute("ALTER TYPE test_type_caching_2 add i blob");

      session.checkSchemaAgreement();

      // wait for latches and fail if they don't reach zero before timeout
      assertThat(
              Uninterruptibles.awaitUninterruptibly(
                  preparedStmtCacheRemoveLatch, 120, TimeUnit.SECONDS))
          .withFailMessage("preparedStmtCacheRemoveLatch did not trigger before timeout")
          .isTrue();
      assertThat(Uninterruptibles.awaitUninterruptibly(typeChangeEventLatch, 20, TimeUnit.SECONDS))
          .withFailMessage("typeChangeEventLatch did not trigger before timeout")
          .isTrue();

      /* Okay, the latch triggered so cache processing should now be done.  Let's validate :allthethings: */
      assertThat(changedTypes.keySet()).isEqualTo(expectedChangedTypes);
      assertThat(removedQueryIds.get()).isNotEmpty().get().isEqualTo(queryId2);
      assertThat(getPreparedCacheSize(session)).isEqualTo(1);

      // check no errors were seen in callback (and report those as fail msgs)
      // if something is broken these may still succeed due to timing
      // but shouldn't intermittently fail if the code is working properly
      assertThat(typeChangeEventError.get())
          .withFailMessage(() -> typeChangeEventError.get().get())
          .isEmpty();
      assertThat(removedQueryEventError.get())
          .withFailMessage(() -> removedQueryEventError.get().get())
          .isEmpty();
    }
  }

  Consumer<CqlSession> setupCacheEntryTestBasic =
      (session) -> {
        session.execute("CREATE TYPE test_type_caching_1 (a text, b int)");
        session.execute("CREATE TYPE test_type_caching_2 (c int, d text)");
        session.execute(
            "CREATE TABLE test_table_caching_1 (e int primary key, f frozen<test_type_caching_1>)");
        session.execute(
            "CREATE TABLE test_table_caching_2 (g int primary key, h frozen<test_type_caching_2>)");
      };

  @Test
  public void should_invalidate_cache_entry_on_basic_udt_change_result_set() {
    SchemaChangeSynchronizer.withLock(
        () -> {
          invalidationResultSetTest(
              setupCacheEntryTestBasic, ImmutableSet.of("test_type_caching_2"));
        });
  }

  @Test
  public void should_invalidate_cache_entry_on_basic_udt_change_variable_defs() {
    SchemaChangeSynchronizer.withLock(
        () -> {
          invalidationVariableDefsTest(
              setupCacheEntryTestBasic, false, ImmutableSet.of("test_type_caching_2"));
        });
  }

  Consumer<CqlSession> setupCacheEntryTestCollection =
      (session) -> {
        session.execute("CREATE TYPE test_type_caching_1 (a text, b int)");
        session.execute("CREATE TYPE test_type_caching_2 (c int, d text)");
        session.execute(
            "CREATE TABLE test_table_caching_1 (e int primary key, f list<frozen<test_type_caching_1>>)");
        session.execute(
            "CREATE TABLE test_table_caching_2 (g int primary key, h list<frozen<test_type_caching_2>>)");
      };

  @Test
  public void should_invalidate_cache_entry_on_collection_udt_change_result_set() {
    SchemaChangeSynchronizer.withLock(
        () -> {
          invalidationResultSetTest(
              setupCacheEntryTestCollection, ImmutableSet.of("test_type_caching_2"));
        });
  }

  @Test
  public void should_invalidate_cache_entry_on_collection_udt_change_variable_defs() {
    SchemaChangeSynchronizer.withLock(
        () -> {
          invalidationVariableDefsTest(
              setupCacheEntryTestCollection, true, ImmutableSet.of("test_type_caching_2"));
        });
  }

  Consumer<CqlSession> setupCacheEntryTestTuple =
      (session) -> {
        session.execute("CREATE TYPE test_type_caching_1 (a text, b int)");
        session.execute("CREATE TYPE test_type_caching_2 (c int, d text)");
        session.execute(
            "CREATE TABLE test_table_caching_1 (e int primary key, f tuple<int, test_type_caching_1, text>)");
        session.execute(
            "CREATE TABLE test_table_caching_2 (g int primary key, h tuple<text, test_type_caching_2, int>)");
      };

  @Test
  public void should_invalidate_cache_entry_on_tuple_udt_change_result_set() {
    SchemaChangeSynchronizer.withLock(
        () -> {
          invalidationResultSetTest(
              setupCacheEntryTestTuple, ImmutableSet.of("test_type_caching_2"));
        });
  }

  @Test
  public void should_invalidate_cache_entry_on_tuple_udt_change_variable_defs() {
    SchemaChangeSynchronizer.withLock(
        () -> {
          invalidationVariableDefsTest(
              setupCacheEntryTestTuple, false, ImmutableSet.of("test_type_caching_2"));
        });
  }

  Consumer<CqlSession> setupCacheEntryTestNested =
      (session) -> {
        session.execute("CREATE TYPE test_type_caching_1 (a text, b int)");
        session.execute("CREATE TYPE test_type_caching_2 (c int, d text)");
        session.execute("CREATE TYPE test_type_caching_3 (e frozen<test_type_caching_1>, f int)");
        session.execute("CREATE TYPE test_type_caching_4 (g int, h frozen<test_type_caching_2>)");
        session.execute(
            "CREATE TABLE test_table_caching_1 (e int primary key, f frozen<test_type_caching_3>)");
        session.execute(
            "CREATE TABLE test_table_caching_2 (g int primary key, h frozen<test_type_caching_4>)");
      };

  @Test
  public void should_invalidate_cache_entry_on_nested_udt_change_result_set() {
    SchemaChangeSynchronizer.withLock(
        () -> {
          invalidationResultSetTest(
              setupCacheEntryTestNested,
              ImmutableSet.of("test_type_caching_2", "test_type_caching_4"));
        });
  }

  @Test
  public void should_invalidate_cache_entry_on_nested_udt_change_variable_defs() {
    SchemaChangeSynchronizer.withLock(
        () -> {
          invalidationVariableDefsTest(
              setupCacheEntryTestNested,
              false,
              ImmutableSet.of("test_type_caching_2", "test_type_caching_4"));
        });
  }

  /* ========================= Infrastructure copied from PreparedStatementIT ========================= */
  private CqlSession sessionWithCacheSizeMetric() {
    return SessionUtils.newSession(
        ccmRule,
        sessionRule.keyspace(),
        SessionUtils.configLoaderBuilder()
            .withInt(DefaultDriverOption.REQUEST_PAGE_SIZE, 2)
            .withDuration(DefaultDriverOption.REQUEST_TIMEOUT, Duration.ofSeconds(30))
            .withStringList(
                DefaultDriverOption.METRICS_SESSION_ENABLED,
                ImmutableList.of(DefaultSessionMetric.CQL_PREPARED_CACHE_SIZE.getPath()))
            .build());
  }

  @SuppressWarnings("unchecked")
  private static long getPreparedCacheSize(CqlSession session) {
    return session
        .getMetrics()
        .flatMap(metrics -> metrics.getSessionMetric(DefaultSessionMetric.CQL_PREPARED_CACHE_SIZE))
        .map(metric -> ((Gauge<Long>) metric).getValue())
        .orElseThrow(
            () ->
                new AssertionError(
                    "Could not access metric "
                        + DefaultSessionMetric.CQL_PREPARED_CACHE_SIZE.getPath()));
  }
}
