/*
 * Copyright DataStax, Inc.
 *
 * Licensed 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.type.codec.registry;

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

import com.datastax.oss.driver.api.core.CqlSession;
import com.datastax.oss.driver.api.core.ProtocolVersion;
import com.datastax.oss.driver.api.core.cql.BoundStatement;
import com.datastax.oss.driver.api.core.cql.PreparedStatement;
import com.datastax.oss.driver.api.core.cql.ResultSet;
import com.datastax.oss.driver.api.core.cql.Row;
import com.datastax.oss.driver.api.core.cql.SimpleStatement;
import com.datastax.oss.driver.api.core.data.UdtValue;
import com.datastax.oss.driver.api.core.type.DataType;
import com.datastax.oss.driver.api.core.type.DataTypes;
import com.datastax.oss.driver.api.core.type.UserDefinedType;
import com.datastax.oss.driver.api.core.type.codec.CodecNotFoundException;
import com.datastax.oss.driver.api.core.type.codec.MappingCodec;
import com.datastax.oss.driver.api.core.type.codec.TypeCodec;
import com.datastax.oss.driver.api.core.type.codec.TypeCodecs;
import com.datastax.oss.driver.api.core.type.codec.registry.MutableCodecRegistry;
import com.datastax.oss.driver.api.core.type.reflect.GenericType;
import com.datastax.oss.driver.api.testinfra.ccm.CcmRule;
import com.datastax.oss.driver.api.testinfra.session.SessionRule;
import com.datastax.oss.driver.api.testinfra.session.SessionUtils;
import com.datastax.oss.driver.categories.ParallelizableTests;
import com.datastax.oss.driver.internal.core.type.codec.IntCodec;
import com.datastax.oss.driver.internal.core.type.codec.UdtCodec;
import com.datastax.oss.driver.internal.core.type.codec.extras.OptionalCodec;
import com.google.common.collect.ImmutableMap;
import edu.umd.cs.findbugs.annotations.NonNull;
import edu.umd.cs.findbugs.annotations.Nullable;
import java.nio.ByteBuffer;
import java.util.HashMap;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.Optional;
import org.assertj.core.util.Maps;
import org.junit.BeforeClass;
import org.junit.ClassRule;
import org.junit.Rule;
import org.junit.Test;
import org.junit.experimental.categories.Category;
import org.junit.rules.RuleChain;
import org.junit.rules.TestName;
import org.junit.rules.TestRule;

@Category(ParallelizableTests.class)
public class CodecRegistryIT {

  private static final CcmRule CCM_RULE = CcmRule.getInstance();

  private static final SessionRule<CqlSession> SESSION_RULE = SessionRule.builder(CCM_RULE).build();

  @ClassRule
  public static final TestRule CHAIN = RuleChain.outerRule(CCM_RULE).around(SESSION_RULE);

  @Rule public TestName name = new TestName();

  @BeforeClass
  public static void createSchema() {
    // table with simple primary key, single cell.
    SESSION_RULE
        .session()
        .execute(
            SimpleStatement.builder("CREATE TABLE IF NOT EXISTS test (k text primary key, v int)")
                .setExecutionProfile(SESSION_RULE.slowProfile())
                .build());
    // table with map value
    SESSION_RULE
        .session()
        .execute(
            SimpleStatement.builder(
                    "CREATE TABLE IF NOT EXISTS test2 (k0 text, k1 int, v map<int,text>, primary key (k0, k1))")
                .setExecutionProfile(SESSION_RULE.slowProfile())
                .build());
    // table with UDT
    SESSION_RULE
        .session()
        .execute(
            SimpleStatement.builder("CREATE TYPE IF NOT EXISTS coordinates (x int, y int)")
                .setExecutionProfile(SESSION_RULE.slowProfile())
                .build());
    SESSION_RULE
        .session()
        .execute(
            SimpleStatement.builder(
                    "CREATE TABLE IF NOT EXISTS test3 (k0 text, k1 int, v map<text,frozen<coordinates>>, primary key (k0, k1))")
                .setExecutionProfile(SESSION_RULE.slowProfile())
                .build());
  }

  // A simple codec that allows float values to be used for cassandra int column type.
  private static class FloatCIntCodec implements TypeCodec<Float> {

    private static final IntCodec intCodec = new IntCodec();

    @NonNull
    @Override
    public GenericType<Float> getJavaType() {
      return GenericType.of(Float.class);
    }

    @NonNull
    @Override
    public DataType getCqlType() {
      return DataTypes.INT;
    }

    @Override
    public ByteBuffer encode(Float value, @NonNull ProtocolVersion protocolVersion) {
      return intCodec.encode(value.intValue(), protocolVersion);
    }

    @Override
    public Float decode(ByteBuffer bytes, @NonNull ProtocolVersion protocolVersion) {
      return intCodec.decode(bytes, protocolVersion).floatValue();
    }

    @NonNull
    @Override
    public String format(Float value) {
      return intCodec.format(value.intValue());
    }

    @Override
    public Float parse(String value) {
      return intCodec.parse(value).floatValue();
    }
  }

  @Test
  public void should_throw_exception_if_no_codec_registered_for_type_set() {
    PreparedStatement prepared =
        SESSION_RULE.session().prepare("INSERT INTO test (k, v) values (?, ?)");

    // float value for int column should not work since no applicable codec.
    Throwable t =
        catchThrowable(
            () ->
                prepared
                    .boundStatementBuilder()
                    .setString(0, name.getMethodName())
                    .setFloat(1, 3.14f)
                    .build());

    assertThat(t).isInstanceOf(CodecNotFoundException.class);
  }

  @Test
  public void should_throw_exception_if_no_codec_registered_for_type_get() {
    PreparedStatement prepared =
        SESSION_RULE.session().prepare("INSERT INTO test (k, v) values (?, ?)");

    BoundStatement insert =
        prepared.boundStatementBuilder().setString(0, name.getMethodName()).setInt(1, 2).build();
    SESSION_RULE.session().execute(insert);

    ResultSet result =
        SESSION_RULE
            .session()
            .execute(
                SimpleStatement.builder("SELECT v from test where k = ?")
                    .addPositionalValue(name.getMethodName())
                    .build());

    List<Row> rows = result.all();
    assertThat(rows).hasSize(1);

    // should not be able to access int column as float as no codec is registered to handle that.
    Row row = rows.iterator().next();

    Throwable t = catchThrowable(() -> assertThat(row.getFloat("v")).isEqualTo(3.0f));

    assertThat(t).isInstanceOf(CodecNotFoundException.class);
  }

  @Test
  public void should_be_able_to_register_and_use_custom_codec() {
    // create a cluster with a registered codec from Float <-> cql int.
    try (CqlSession session =
        (CqlSession)
            SessionUtils.baseBuilder()
                .addTypeCodecs(new FloatCIntCodec())
                .addContactEndPoints(CCM_RULE.getContactPoints())
                .withKeyspace(SESSION_RULE.keyspace())
                .build()) {
      PreparedStatement prepared = session.prepare("INSERT INTO test (k, v) values (?, ?)");

      // float value for int column should work.
      BoundStatement insert =
          prepared
              .boundStatementBuilder()
              .setString(0, name.getMethodName())
              .setFloat(1, 3.14f)
              .build();
      session.execute(insert);

      ResultSet result =
          session.execute(
              SimpleStatement.builder("SELECT v from test where k = ?")
                  .addPositionalValue(name.getMethodName())
                  .build());

      List<Row> rows = result.all();
      assertThat(rows).hasSize(1);

      // should be able to retrieve value back as float, some precision is lost due to going from
      // int -> float.
      Row row = rows.iterator().next();
      assertThat(row.getFloat("v")).isEqualTo(3.0f);
      assertThat(row.getFloat(0)).isEqualTo(3.0f);
    }
  }

  @Test
  public void should_register_custom_codec_at_runtime() {
    // Still create a separate session because we don't want to interfere with other tests
    try (CqlSession session = SessionUtils.newSession(CCM_RULE, SESSION_RULE.keyspace())) {

      MutableCodecRegistry registry =
          (MutableCodecRegistry) session.getContext().getCodecRegistry();
      registry.register(new FloatCIntCodec());

      PreparedStatement prepared = session.prepare("INSERT INTO test (k, v) values (?, ?)");

      // float value for int column should work.
      BoundStatement insert =
          prepared
              .boundStatementBuilder()
              .setString(0, name.getMethodName())
              .setFloat(1, 3.14f)
              .build();
      session.execute(insert);

      ResultSet result =
          session.execute(
              SimpleStatement.builder("SELECT v from test where k = ?")
                  .addPositionalValue(name.getMethodName())
                  .build());

      List<Row> rows = result.all();
      assertThat(rows).hasSize(1);

      // should be able to retrieve value back as float, some precision is lost due to going from
      // int -> float.
      Row row = rows.iterator().next();
      assertThat(row.getFloat("v")).isEqualTo(3.0f);
      assertThat(row.getFloat(0)).isEqualTo(3.0f);
    }
  }

  @Test
  public void should_be_able_to_register_and_use_custom_codec_with_generic_type() {
    // create a cluster with registered codecs using OptionalCodec
    OptionalCodec<Map<Integer, String>> optionalMapCodec =
        new OptionalCodec<>(TypeCodecs.mapOf(TypeCodecs.INT, TypeCodecs.TEXT));
    TypeCodec<Map<Integer, Optional<String>>> mapWithOptionalValueCodec =
        TypeCodecs.mapOf(TypeCodecs.INT, new OptionalCodec<>(TypeCodecs.TEXT));

    try (CqlSession session =
        (CqlSession)
            SessionUtils.baseBuilder()
                .addTypeCodecs(optionalMapCodec, mapWithOptionalValueCodec)
                .addContactEndPoints(CCM_RULE.getContactPoints())
                .withKeyspace(SESSION_RULE.keyspace())
                .build()) {
      PreparedStatement prepared =
          session.prepare("INSERT INTO test2 (k0, k1, v) values (?, ?, ?)");

      // optional map should work.
      Map<Integer, String> v0 = Maps.newHashMap(0, "value");
      Optional<Map<Integer, String>> v0Opt = Optional.of(v0);
      BoundStatement insert =
          prepared
              .boundStatementBuilder()
              .setString(0, name.getMethodName())
              .setInt(1, 0)
              .set(
                  2,
                  v0Opt,
                  optionalMapCodec
                      .getJavaType()) // use java type so has to be looked up in registry.
              .build();
      session.execute(insert);

      // optional absent map should work.
      Optional<Map<Integer, String>> absentMap = Optional.empty();
      insert =
          prepared
              .boundStatementBuilder()
              .setString(0, name.getMethodName())
              .setInt(1, 1)
              .set(2, absentMap, optionalMapCodec.getJavaType())
              .build();
      session.execute(insert);

      // map with optional value should work - note that you can't have null values in collections,
      // so this is not technically practical but want to validate that custom codec resolution
      // works
      // when it's composed in a collection codec.
      Map<Integer, Optional<String>> v2Map = Maps.newHashMap(1, Optional.of("hello"));
      insert =
          prepared
              .boundStatementBuilder()
              .setString(0, name.getMethodName())
              .setInt(1, 2)
              .set(2, v2Map, mapWithOptionalValueCodec.getJavaType())
              .build();
      session.execute(insert);

      ResultSet result =
          session.execute(
              SimpleStatement.builder("SELECT v from test2 where k0 = ?")
                  .addPositionalValues(name.getMethodName())
                  .build());

      List<Row> rows = result.all();
      assertThat(rows).hasSize(3);

      Iterator<Row> iterator = rows.iterator();
      // row (at key 0) should have v0
      Row row = iterator.next();
      // should be able to retrieve value back as an optional map.
      assertThat(row.get(0, optionalMapCodec.getJavaType())).isEqualTo(v0Opt);
      // should be able to retrieve value back as map.
      assertThat(row.getMap(0, Integer.class, String.class)).isEqualTo(v0);

      // next row (at key 1) should be absent (null value).
      row = iterator.next();
      // value should be null.
      assertThat(row.isNull(0)).isTrue();
      // getting with codec should return Optional.empty()
      assertThat(row.get(0, optionalMapCodec.getJavaType())).isEqualTo(absentMap);
      // getting with map should return an empty map.
      assertThat(row.getMap(0, Integer.class, String.class)).isEmpty();

      // next row (at key 2) should have v2
      row = iterator.next();
      // getting with codec should return with the correct type.
      assertThat(row.get(0, mapWithOptionalValueCodec.getJavaType())).isEqualTo(v2Map);
      // getting with map should return a map without optional value.
      assertThat(row.getMap(0, Integer.class, String.class)).isEqualTo(Maps.newHashMap(1, "hello"));
    }
  }

  @Test
  public void should_be_able_to_handle_empty_collections() {
    try (CqlSession session =
        (CqlSession)
            SessionUtils.<CqlSession>baseBuilder()
                .addContactEndPoints(CCM_RULE.getContactPoints())
                .withKeyspace(SESSION_RULE.keyspace())
                .build()) {

      // Using prepared statements (CQL type is known)
      PreparedStatement prepared =
          session.prepare("INSERT INTO test2 (k0, k1, v) values (?, ?, ?)");

      BoundStatement insert =
          prepared
              .boundStatementBuilder()
              .setString(0, name.getMethodName())
              .setInt(1, 0)
              .setMap(2, new HashMap<>(), Integer.class, String.class)
              .build();
      session.execute(insert);

      // Using simple statements (CQL type is unknown)
      session.execute(
          SimpleStatement.newInstance(
              "INSERT INTO test2 (k0, k1, v) values (?, ?, ?)",
              name.getMethodName(),
              1,
              new HashMap<>()));

      ResultSet result =
          session.execute(
              SimpleStatement.builder("SELECT v from test2 where k0 = ?")
                  .addPositionalValues(name.getMethodName())
                  .build());

      List<Row> rows = result.all();
      assertThat(rows).hasSize(2);

      Row row1 = rows.get(0);
      assertThat(row1.isNull(0)).isTrue();
      assertThat(row1.getMap(0, Integer.class, String.class)).isEmpty();

      Row row2 = rows.get(1);
      assertThat(row2.isNull(0)).isTrue();
      assertThat(row2.getMap(0, Integer.class, String.class)).isEmpty();
    }
  }

  private static final class Coordinates {

    public final int x;
    public final int y;

    public Coordinates(int x, int y) {
      this.x = x;
      this.y = y;
    }

    @Override
    public boolean equals(Object o) {
      if (this == o) {
        return true;
      }
      if (o == null || getClass() != o.getClass()) {
        return false;
      }
      Coordinates that = (Coordinates) o;
      return this.x == that.x && this.y == that.y;
    }

    @Override
    public int hashCode() {
      return Objects.hash(x, y);
    }
  }

  private static class CoordinatesCodec extends MappingCodec<UdtValue, Coordinates> {

    public CoordinatesCodec(@NonNull TypeCodec<UdtValue> innerCodec) {
      super(innerCodec, GenericType.of(Coordinates.class));
    }

    @NonNull
    @Override
    public UserDefinedType getCqlType() {
      return (UserDefinedType) super.getCqlType();
    }

    @Nullable
    @Override
    protected Coordinates innerToOuter(@Nullable UdtValue value) {
      return value == null ? null : new Coordinates(value.getInt("x"), value.getInt("y"));
    }

    @Nullable
    @Override
    protected UdtValue outerToInner(@Nullable Coordinates value) {
      return value == null
          ? null
          : getCqlType().newValue().setInt("x", value.x).setInt("y", value.y);
    }
  }

  @Test
  public void should_register_and_use_custom_codec_for_user_defined_type() {

    Map<String, Coordinates> coordinatesMap = ImmutableMap.of("home", new Coordinates(12, 34));
    GenericType<Map<String, Coordinates>> coordinatesMapType =
        GenericType.mapOf(String.class, Coordinates.class);

    // Still create a separate session because we don't want to interfere with other tests
    try (CqlSession session = SessionUtils.newSession(CCM_RULE, SESSION_RULE.keyspace())) {

      // register the mapping codec for UDT coordinates
      UserDefinedType coordinatesUdt =
          session
              .getMetadata()
              .getKeyspace(SESSION_RULE.keyspace())
              .flatMap(ks -> ks.getUserDefinedType("coordinates"))
              .orElseThrow(IllegalStateException::new);
      MutableCodecRegistry codecRegistry =
          (MutableCodecRegistry) session.getContext().getCodecRegistry();

      // Retrieve the inner codec
      TypeCodec<UdtValue> innerCodec = codecRegistry.codecFor(coordinatesUdt);
      assertThat(innerCodec).isInstanceOf(UdtCodec.class);

      // Create the "outer" codec and register it
      CoordinatesCodec coordinatesCodec = new CoordinatesCodec(innerCodec);
      codecRegistry.register(coordinatesCodec);

      // Test that the codec will be used to create on-the-fly codecs
      assertThat(codecRegistry.codecFor(Coordinates.class)).isSameAs(coordinatesCodec);
      assertThat(codecRegistry.codecFor(coordinatesMapType).accepts(coordinatesMap)).isTrue();

      // test insertion
      PreparedStatement prepared =
          session.prepare("INSERT INTO test3 (k0, k1, v) values (?, ?, ?)");
      BoundStatement insert =
          prepared
              .boundStatementBuilder()
              .setString(0, name.getMethodName())
              .setInt(1, 0)
              .set(
                  2,
                  coordinatesMap,
                  coordinatesMapType) // use java type so has to be looked up in registry.
              .build();
      session.execute(insert);

      // test retrieval
      ResultSet result =
          session.execute(
              SimpleStatement.builder("SELECT v from test3 where k0 = ? AND k1 = ?")
                  .addPositionalValues(name.getMethodName(), 0)
                  .build());
      List<Row> rows = result.all();
      assertThat(rows).hasSize(1);
      Row row = rows.get(0);
      assertThat(row.get(0, coordinatesMapType)).isEqualTo(coordinatesMap);
      assertThat(row.getMap(0, String.class, Coordinates.class)).isEqualTo(coordinatesMap);
    }
  }
}
