/*
 * 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.extras.codecs.enums;

import static com.datastax.driver.core.DataType.cint;
import static com.datastax.driver.core.DataType.text;
import static com.datastax.driver.extras.codecs.enums.EnumCodecsTest.Bar.BAR_1;
import static com.datastax.driver.extras.codecs.enums.EnumCodecsTest.Bar.BAR_2;
import static com.datastax.driver.extras.codecs.enums.EnumCodecsTest.Foo.FOO_1;
import static com.datastax.driver.extras.codecs.enums.EnumCodecsTest.Foo.FOO_2;
import static com.google.common.collect.Lists.newArrayList;
import static com.google.common.collect.Sets.newHashSet;
import static org.assertj.core.api.Assertions.assertThat;

import com.datastax.driver.core.CCMTestsSupport;
import com.datastax.driver.core.Cluster;
import com.datastax.driver.core.CodecRegistry;
import com.datastax.driver.core.PreparedStatement;
import com.datastax.driver.core.ResultSet;
import com.datastax.driver.core.Row;
import com.datastax.driver.core.TupleType;
import com.datastax.driver.core.TupleValue;
import com.datastax.driver.core.TypeTokens;
import com.datastax.driver.core.UDTValue;
import com.datastax.driver.core.UserType;
import com.datastax.driver.core.utils.CassandraVersion;
import com.datastax.driver.mapping.Mapper;
import com.datastax.driver.mapping.MappingManager;
import com.datastax.driver.mapping.annotations.Accessor;
import com.datastax.driver.mapping.annotations.ClusteringColumn;
import com.datastax.driver.mapping.annotations.PartitionKey;
import com.datastax.driver.mapping.annotations.Query;
import com.datastax.driver.mapping.annotations.Table;
import com.google.common.collect.ImmutableMap;
import java.util.List;
import java.util.Map;
import java.util.Set;
import org.testng.annotations.BeforeMethod;
import org.testng.annotations.Test;

/**
 * A test that validates that Enums are correctly mapped to varchars (with EnumNameCodec) and ints
 * (with EnumOrdinalCodec). It also validates that both codecs may coexist in the same
 * CodecRegistry.
 */
@CassandraVersion("2.1.0")
public class EnumCodecsTest extends CCMTestsSupport {

  private final String insertQuery =
      "INSERT INTO t1 (pk, foo, foos, bar, bars, foobars, tup, udt) VALUES (?, ?, ?, ?, ?, ?, ?, ?)";
  private final String selectQuery =
      "SELECT pk, foo, foos, bar, bars, foobars, tup, udt FROM t1 WHERE pk = ?";

  private final int pk = 42;

  private final List<Foo> foos = newArrayList(FOO_1, FOO_2);
  private final Set<Bar> bars = newHashSet(BAR_1, BAR_2);
  private final Map<Foo, Bar> foobars = ImmutableMap.of(FOO_1, BAR_1, FOO_2, BAR_2);

  private TupleValue tupleValue;
  private UDTValue udtValue;

  @Override
  public void onTestContextInitialized() {
    execute(
        "CREATE TYPE IF NOT EXISTS udt1 (" + "foo int," + "bar text)",
        "CREATE TABLE IF NOT EXISTS t1 ("
            + "pk int, "
            + "foo int, "
            + "foos list<int>, "
            + "bar text, "
            + "bars set<text>, "
            + "foobars map<int,text>, "
            + "tup frozen<tuple<int,varchar>>, "
            + "udt frozen<udt1>,"
            + "primary key (pk, foo))");
  }

  @Override
  public Cluster.Builder createClusterBuilder() {
    return Cluster.builder()
        .withCodecRegistry(
            new CodecRegistry()
                .register(new EnumOrdinalCodec<Foo>(Foo.class))
                .register(new EnumNameCodec<Bar>(Bar.class)));
  }

  @BeforeMethod(groups = "short")
  public void before() {
    TupleType tup = cluster().getMetadata().newTupleType(cint(), text());
    tupleValue = tup.newValue().set(0, FOO_1, Foo.class).set(1, BAR_1, Bar.class);
    UserType udt = cluster().getMetadata().getKeyspace(keyspace).getUserType("udt1");
    udtValue = udt.newValue().set("foo", FOO_1, Foo.class).set("bar", BAR_1, Bar.class);
  }

  @Test(groups = "short")
  public void should_use_enum_codecs_with_simple_statements() {
    session().execute(insertQuery, pk, FOO_1, foos, BAR_1, bars, foobars, tupleValue, udtValue);
    ResultSet rows = session().execute(selectQuery, pk);
    Row row = rows.one();
    assertRow(row);
  }

  @Test(groups = "short")
  public void should_use_enum_codecs_with_prepared_statements_1() {
    session()
        .execute(
            session()
                .prepare(insertQuery)
                .bind(pk, FOO_1, foos, BAR_1, bars, foobars, tupleValue, udtValue));
    PreparedStatement ps = session().prepare(selectQuery);
    ResultSet rows = session().execute(ps.bind(pk));
    Row row = rows.one();
    assertRow(row);
  }

  @Test(groups = "short")
  public void should_use_enum_codecs_with_prepared_statements_2() {
    session()
        .execute(
            session()
                .prepare(insertQuery)
                .bind()
                .setInt(0, pk)
                .set(1, FOO_1, Foo.class)
                .setList(2, foos, Foo.class)
                .set(3, BAR_1, Bar.class)
                .set(4, bars, TypeTokens.setOf(Bar.class))
                .setMap(5, foobars, Foo.class, Bar.class)
                .setTupleValue(6, tupleValue)
                .setUDTValue(7, udtValue));
    PreparedStatement ps = session().prepare(selectQuery);
    ResultSet rows = session().execute(ps.bind().setInt(0, pk));
    Row row = rows.one();
    assertRow(row);
  }

  @Test(groups = "short")
  public void should_use_mapper_to_store_and_retrieve_nulls_with_enum_codecs() {
    // given
    MappingManager manager = new MappingManager(session());
    Mapper<Mapped> mapper = manager.mapper(Mapped.class);
    Mapped pojo = new Mapped();
    pojo.pk = 42;
    pojo.foo = FOO_1;
    // when
    mapper.save(pojo);
    Mapped actual = mapper.get(42, FOO_1);
    // then
    assertThat(actual).isEqualToComparingFieldByField(pojo);
  }

  @Test(groups = "short")
  public void should_use_mapper_to_store_and_retrieve_values_with_enum_codecs() {
    // given
    MappingManager manager = new MappingManager(session());
    Mapper<Mapped> mapper = manager.mapper(Mapped.class);
    Mapped pojo = new Mapped();
    pojo.pk = 42;
    pojo.foo = FOO_1;
    pojo.bar = BAR_1;
    pojo.foos = newArrayList(FOO_2, FOO_1);
    pojo.bars = newHashSet(BAR_1, BAR_2);
    pojo.foobars = ImmutableMap.of(FOO_1, BAR_2);
    // when
    mapper.save(pojo);
    Mapped actual = mapper.get(42, FOO_1);
    // then
    assertThat(actual).isEqualToComparingFieldByField(pojo);
  }

  @Test(groups = "short")
  public void should_use_accessor_to_store_and_retrieve_values_with_enum_codecs() {
    // given
    MappedAccessor accessor = new MappingManager(session()).createAccessor(MappedAccessor.class);
    Mapped expected = new Mapped();
    expected.pk = 42;
    expected.foo = FOO_1;
    expected.bar = BAR_2;
    // when
    accessor.insert(42, FOO_1, BAR_2);
    Mapped result = accessor.getByFoo(42, FOO_1);
    // then
    assertThat(result).isEqualToComparingFieldByField(expected);
  }

  private void assertRow(Row row) {

    assertThat(row.getInt(0)).isEqualTo(pk);

    assertThat(row.getObject(1))
        .isEqualTo(FOO_1.ordinal()); // uses the built-in IntCodec because CQL type is int
    assertThat(row.getInt("foo")).isEqualTo(FOO_1.ordinal()); // uses the built-in IntCodec
    assertThat(row.get(1, Integer.class)).isEqualTo(FOO_1.ordinal()); // forces IntCodec
    assertThat(row.get("foo", Foo.class)).isEqualTo(FOO_1); // forces EnumOrdinalCodec

    assertThat(row.getObject(2))
        .isEqualTo(
            newArrayList(
                FOO_1.ordinal(),
                FOO_2.ordinal())); // uses the built-in ListCodec(IntCodec) because CQL type is
    // list<int>
    assertThat(row.getList(2, Integer.class))
        .isEqualTo(newArrayList(FOO_1.ordinal(), FOO_2.ordinal()));
    assertThat(row.getList("foos", Foo.class)).isEqualTo(newArrayList(FOO_1, FOO_2));
    assertThat(row.get(2, TypeTokens.listOf(Integer.class)))
        .isEqualTo(newArrayList(FOO_1.ordinal(), FOO_2.ordinal()));
    assertThat(row.get("foos", TypeTokens.listOf(Foo.class))).isEqualTo(newArrayList(FOO_1, FOO_2));

    assertThat(row.getObject(3))
        .isEqualTo(BAR_1.name()); // uses the built-in VarcharCodec because CQL type is varchar
    assertThat(row.getString("bar")).isEqualTo(BAR_1.name()); // forces VarcharCodec
    assertThat(row.get(3, String.class)).isEqualTo(BAR_1.name()); // forces VarcharCodec
    assertThat(row.get("bar", Bar.class)).isEqualTo(BAR_1); // forces EnumNameCodec

    assertThat(row.getObject(4))
        .isEqualTo(
            newHashSet(
                BAR_1.name(),
                BAR_2.name())); // uses the built-in SetCodec(VarcharCodec) because CQL type is
    // set<varchar>
    assertThat(row.getSet(4, String.class)).isEqualTo(newHashSet(BAR_1.name(), BAR_2.name()));
    assertThat(row.getSet("bars", Bar.class)).isEqualTo(newHashSet(BAR_1, BAR_2));
    assertThat(row.get(4, TypeTokens.setOf(String.class)))
        .isEqualTo(newHashSet(BAR_1.name(), BAR_2.name()));
    assertThat(row.get("bars", TypeTokens.setOf(Bar.class))).isEqualTo(newHashSet(BAR_1, BAR_2));

    // uses default built-in codec MapCodec(IntCodec, VarcharCodec) because CQL type is map<int,
    // varchar>
    assertThat(row.getObject(5))
        .isEqualTo(ImmutableMap.of(FOO_1.ordinal(), BAR_1.name(), FOO_2.ordinal(), BAR_2.name()));

    // getMap with combinations of built-in or EnumX codecs
    assertThat(row.getMap(5, Integer.class, String.class))
        .isEqualTo(ImmutableMap.of(FOO_1.ordinal(), BAR_1.name(), FOO_2.ordinal(), BAR_2.name()));
    assertThat(row.getMap(5, Foo.class, String.class))
        .isEqualTo(ImmutableMap.of(FOO_1, BAR_1.name(), FOO_2, BAR_2.name()));
    assertThat(row.getMap(5, Integer.class, Bar.class))
        .isEqualTo(ImmutableMap.of(FOO_1.ordinal(), BAR_1, FOO_2.ordinal(), BAR_2));
    assertThat(row.getMap(5, Foo.class, Bar.class))
        .isEqualTo(ImmutableMap.of(FOO_1, BAR_1, FOO_2, BAR_2));

    // get + TypeToken with combinations of built-in or EnumX codecs
    assertThat(row.get("foobars", TypeTokens.mapOf(Integer.class, String.class)))
        .isEqualTo(ImmutableMap.of(FOO_1.ordinal(), BAR_1.name(), FOO_2.ordinal(), BAR_2.name()));
    assertThat(row.get("foobars", TypeTokens.mapOf(Foo.class, String.class)))
        .isEqualTo(ImmutableMap.of(FOO_1, BAR_1.name(), FOO_2, BAR_2.name()));
    assertThat(row.get("foobars", TypeTokens.mapOf(Integer.class, Bar.class)))
        .isEqualTo(ImmutableMap.of(FOO_1.ordinal(), BAR_1, FOO_2.ordinal(), BAR_2));
    assertThat(row.get("foobars", TypeTokens.mapOf(Foo.class, Bar.class)))
        .isEqualTo(ImmutableMap.of(FOO_1, BAR_1, FOO_2, BAR_2));

    assertThat(row.getTupleValue(6).getInt(0)).isEqualTo(FOO_1.ordinal());
    assertThat(row.get(6, TupleValue.class).get(0, Foo.class)).isEqualTo(FOO_1);
    assertThat(row.getTupleValue("tup").getString(1)).isEqualTo(BAR_1.name());
    assertThat(row.get("tup", TupleValue.class).get(1, Bar.class)).isEqualTo(BAR_1);

    assertThat(row.getUDTValue(7).getInt("foo")).isEqualTo(FOO_1.ordinal());
    assertThat(row.get(7, UDTValue.class).get("foo", Foo.class)).isEqualTo(FOO_1);
    assertThat(row.getUDTValue("udt").getString("bar")).isEqualTo(BAR_1.name());
    assertThat(row.get("udt", UDTValue.class).get("bar", Bar.class)).isEqualTo(BAR_1);
  }

  @SuppressWarnings("unused")
  @Table(name = "t1")
  public static class Mapped {

    @PartitionKey private int pk;

    @ClusteringColumn private Foo foo;

    private List<Foo> foos;

    private Bar bar;

    private Set<Bar> bars;

    private Map<Foo, Bar> foobars;

    public int getPk() {
      return pk;
    }

    public void setPk(int pk) {
      this.pk = pk;
    }

    public Foo getFoo() {
      return foo;
    }

    public void setFoo(Foo foo) {
      this.foo = foo;
    }

    public List<Foo> getFoos() {
      return foos;
    }

    public void setFoos(List<Foo> foos) {
      this.foos = foos;
    }

    public Bar getBar() {
      return bar;
    }

    public void setBar(Bar bar) {
      this.bar = bar;
    }

    public Set<Bar> getBars() {
      return bars;
    }

    public void setBars(Set<Bar> bars) {
      this.bars = bars;
    }

    public Map<Foo, Bar> getFoobars() {
      return foobars;
    }

    public void setFoobars(Map<Foo, Bar> foobars) {
      this.foobars = foobars;
    }
  }

  @Accessor
  public interface MappedAccessor {

    @Query("select * from t1 where pk=? and foo=?")
    Mapped getByFoo(int pk, Foo foo);

    @Query("insert into t1 (pk, foo, bar) values (?,?,?)")
    ResultSet insert(int pk, Foo foo, Bar bar);
  }

  enum Foo {
    FOO_1,
    FOO_2
  }

  enum Bar {
    BAR_1,
    BAR_2
  }
}
