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

import static com.datastax.driver.core.querybuilder.QueryBuilder.bindMarker;
import static com.datastax.driver.core.querybuilder.QueryBuilder.eq;
import static com.datastax.driver.core.querybuilder.QueryBuilder.fromJson;
import static com.datastax.driver.core.querybuilder.QueryBuilder.insertInto;
import static com.datastax.driver.core.querybuilder.QueryBuilder.quote;
import static com.datastax.driver.core.querybuilder.QueryBuilder.select;
import static org.assertj.core.api.Assertions.assertThat;
import static org.assertj.core.api.Assertions.fail;

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.ProtocolVersion;
import com.datastax.driver.core.ResultSet;
import com.datastax.driver.core.Row;
import com.datastax.driver.core.querybuilder.BuiltStatement;
import com.datastax.driver.core.querybuilder.QueryBuilder;
import com.datastax.driver.core.utils.CassandraVersion;
import java.io.IOException;
import java.io.StringWriter;
import javax.json.Json;
import javax.json.JsonArray;
import javax.json.JsonObject;
import javax.json.JsonStructure;
import javax.json.JsonWriter;
import org.testng.annotations.DataProvider;
import org.testng.annotations.Test;

public class Jsr353JsonCodecTest extends CCMTestsSupport {

  private static final Jsr353JsonCodec jsonCodec = new Jsr353JsonCodec();

  private static final JsonObject alice =
      Json.createObjectBuilder().add("id", 1).add("name", "Alice").build();
  private static final JsonObject bob =
      Json.createObjectBuilder().add("id", 2).add("name", "Bob").build();
  private static final JsonObject charlie =
      Json.createObjectBuilder().add("id", 3).add("name", "Charlie").build();

  private static final JsonArray bobAndCharlie =
      Json.createArrayBuilder().add(bob).add(charlie).build();

  private static final String insertQuery = "INSERT INTO t1 (c1, c2) VALUES (?, ?)";
  private static final String selectQuery = "SELECT c1, c2 FROM t1 WHERE c1 = ? and c2 = ?";

  private static final String notAJsonString = "this text is not json";

  private static final String TABLE1 = "test1";
  private static final String TABLE2 = "test2";

  @Override
  public void onTestContextInitialized() {
    execute(
        "CREATE TABLE t1 (c1 text, c2 text, PRIMARY KEY (c1, c2))",
        String.format(
            "CREATE TABLE %s (k text, \"miXeD\" text, i int, f float, PRIMARY KEY (k, \"miXeD\"))",
            TABLE1),
        insertInto(TABLE1)
            .value("k", "key0")
            .value(quote("miXeD"), "a")
            .value("i", 1)
            .value("f", 1.1)
            .toString(),
        insertInto(TABLE1)
            .value("k", "key0")
            .value(quote("miXeD"), "b")
            .value("i", 2)
            .value("f", 2.5)
            .toString(),
        insertInto(TABLE1)
            .value("k", "key0")
            .value(quote("miXeD"), "c")
            .value("i", 3)
            .value("f", 3.7)
            .toString(),
        insertInto(TABLE1)
            .value("k", "key0")
            .value(quote("miXeD"), "d")
            .value("i", 4)
            .value("f", 5.0)
            .toString());

    if (ccm().getProtocolVersion().compareTo(ProtocolVersion.V3) >= 0) {
      execute(
          "CREATE TYPE address (street text, zipcode int, phones list<text>)",
          String.format("CREATE TABLE %s (k text PRIMARY KEY, v frozen<address>)", TABLE2));
    }
  }

  @Override
  public Cluster.Builder createClusterBuilder() {
    return Cluster.builder()
        .withCodecRegistry(
            new CodecRegistry().register(jsonCodec) // global JsonStructure <-> varchar codec
            );
  }

  @DataProvider(name = "Jsr353JsonCodecTest")
  public static Object[][] parameters() {
    return new Object[][] {{alice}, {bobAndCharlie}};
  }

  @Test(groups = "unit")
  public void test_cql_text_to_json() {
    Jsr353JsonCodec codec = new Jsr353JsonCodec();
    String json = "'{\"id\":1,\"name\":\"John Doe\"}'";
    JsonObject object = Json.createObjectBuilder().add("id", 1).add("name", "John Doe").build();
    assertThat(codec.format(object)).isEqualTo(json);
    assertThat(codec.parse(json)).isEqualTo(object);
  }

  @Test(groups = "unit")
  public void test_nulls() {
    Jsr353JsonCodec codec = new Jsr353JsonCodec();
    assertThat(codec.format(null)).isEqualTo("NULL");
    assertThat(codec.parse(null)).isNull();
  }

  @Test(groups = "short", dataProvider = "Jsr353JsonCodecTest")
  @CassandraVersion("2.0.0")
  public void should_use_custom_codec_with_simple_statements(JsonStructure object)
      throws IOException {
    session().execute(insertQuery, notAJsonString, object);
    ResultSet rows = session().execute(selectQuery, notAJsonString, object);
    Row row = rows.one();
    assertRow(row, object);
  }

  @Test(groups = "short", dataProvider = "Jsr353JsonCodecTest")
  public void should_use_custom_codec_with_built_statements_1(JsonStructure object)
      throws IOException {
    BuiltStatement insertStmt =
        insertInto("t1").value("c1", bindMarker()).value("c2", bindMarker());
    BuiltStatement selectStmt =
        select("c1", "c2").from("t1").where(eq("c1", bindMarker())).and(eq("c2", bindMarker()));
    session().execute(session().prepare(insertStmt).bind(notAJsonString, object));
    ResultSet rows = session().execute(session().prepare(selectStmt).bind(notAJsonString, object));
    Row row = rows.one();
    assertRow(row, object);
  }

  @Test(groups = "short", dataProvider = "Jsr353JsonCodecTest")
  public void should_use_custom_codec_with_built_statements_2(JsonStructure object)
      throws IOException {
    BuiltStatement insertStmt = insertInto("t1").value("c1", notAJsonString).value("c2", object);
    BuiltStatement selectStmt =
        select("c1", "c2").from("t1").where(eq("c1", notAJsonString)).and(eq("c2", object));
    session().execute(insertStmt);
    ResultSet rows = session().execute(selectStmt);
    Row row = rows.one();
    assertRow(row, object);
  }

  @Test(groups = "short", dataProvider = "Jsr353JsonCodecTest")
  public void should_use_custom_codec_with_prepared_statements_1(JsonStructure object)
      throws IOException {
    session().execute(session().prepare(insertQuery).bind(notAJsonString, object));
    PreparedStatement ps = session().prepare(selectQuery);
    ResultSet rows = session().execute(ps.bind(notAJsonString, object));
    Row row = rows.one();
    assertRow(row, object);
  }

  @Test(groups = "short", dataProvider = "Jsr353JsonCodecTest")
  public void should_use_custom_codec_with_prepared_statements_2(JsonStructure object)
      throws IOException {
    session()
        .execute(
            session()
                .prepare(insertQuery)
                .bind()
                .setString(0, notAJsonString)
                .set(1, object, JsonStructure.class));
    PreparedStatement ps = session().prepare(selectQuery);
    ResultSet rows =
        session()
            .execute(ps.bind().setString(0, notAJsonString).set(1, object, JsonStructure.class));
    Row row = rows.one();
    assertRow(row, object);
  }

  private void assertRow(Row row, JsonStructure object) throws IOException {
    String json;
    StringWriter sw = new StringWriter();
    JsonWriter writer = Json.createWriter(sw);
    writer.write(object);
    json = sw.toString();
    assertThat(row.getString(0)).isEqualTo(notAJsonString);
    assertThat(row.getObject(1)).isEqualTo(json);
    assertThat(row.getString(1)).isEqualTo(json);
    assertThat(row.get(1, String.class)).isEqualTo(json);
    assertThat(row.get(1, JsonStructure.class)).isEqualTo(object);
  }

  /**
   * Validates that the {@code select().json()} capability in {@link QueryBuilder} appropriately
   * generates a 'SELECT JSON' query and that the returned {@link ResultSet} has rows with a column
   * that is parsable using {@link Jsr353JsonCodec}.
   *
   * @jira_ticket JAVA-743
   * @since 3.1.0
   */
  @Test(groups = "short")
  @CassandraVersion("2.2.0")
  public void should_support_select_json_output() throws Exception {
    // when
    ResultSet r = session().execute(select().json().from(TABLE1).where(eq("k", "key0")));
    // then
    assertThat(r.getAvailableWithoutFetching()).isEqualTo(4);
    for (Row row : r) {
      JsonStructure json = row.get(0, jsonCodec);
      assertThat(json).isInstanceOf(JsonObject.class);
      JsonObject obj = (JsonObject) json;
      assertThat(obj.containsKey("k")).isTrue();
      assertThat(obj.containsKey(quote("miXeD"))).isTrue();
      assertThat(obj.containsKey("i")).isTrue();
      assertThat(obj.containsKey("f")).isTrue();

      assertThat(obj.getString("k")).isEqualTo("key0");

      int i = obj.getInt("i");
      switch (i) {
        case 1:
          assertThat(obj.getString(quote("miXeD"))).isEqualTo("a");
          assertThat(obj.getJsonNumber("f").doubleValue()).isEqualTo(1.1);
          break;
        case 2:
          assertThat(obj.getString(quote("miXeD"))).isEqualTo("b");
          assertThat(obj.getJsonNumber("f").doubleValue()).isEqualTo(2.5);
          break;
        case 3:
          assertThat(obj.getString(quote("miXeD"))).isEqualTo("c");
          assertThat(obj.getJsonNumber("f").doubleValue()).isEqualTo(3.7);
          break;
        case 4:
          assertThat(obj.getString(quote("miXeD"))).isEqualTo("d");
          assertThat(obj.getJsonNumber("f").doubleValue()).isEqualTo(5.0);
          break;
        default:
          fail("Unexpected value for i = " + i);
      }
    }
  }

  /**
   * Validates that the {@code insertInto().json()} capability in {@link QueryBuilder} creates a
   * query that interprets the input as a json document and that fetching that data back as json
   * returns a row with a column that matches the contents of the input json document when retrieved
   * using get with {@link JsonStructure}.
   *
   * @jira_ticket JAVA-743
   * @since 3.1.0
   */
  @Test(groups = "short")
  @CassandraVersion("2.2.0")
  public void should_support_insert_json() throws Exception {
    // when
    String key = "should_support_insert_json_with_codec_format";
    JsonObject input =
        Json.createObjectBuilder()
            .add("k", key)
            .add(quote("miXeD"), "miXeDValue")
            .add("f", 8.6)
            .add("i", 4)
            .build();
    session().execute(insertInto(TABLE1).json(input));
    // then
    ResultSet r = session().execute(select().json().from(TABLE1).where(eq("k", key)));
    assertThat(r.getAvailableWithoutFetching()).isEqualTo(1);

    JsonStructure output = r.one().get(0, JsonStructure.class);
    assertThat(output).isEqualTo(input);
  }

  /**
   * Validates that the {@code fromJson()} and {@code toJson} capability in {@link QueryBuilder}
   * allows specifying json input ({@code fromJson}) and retrieval of json column output ({@code
   * toJson}). This test inserts a row with a UDT providing the UDT column value as json using
   * {@code toJson} and then retrieves that column value using {@code fromJson} and ensures the
   * output matches the input {@link JsonStructure} for the UDT.
   *
   * @jira_ticket JAVA-743
   * @since 3.1.0
   */
  @Test(groups = "short")
  @CassandraVersion("2.2.0")
  public void should_support_fromJson_and_toJson() throws Exception {
    String key = "should_support_fromJson_and_toJson";
    JsonObject inputAddr =
        Json.createObjectBuilder()
            .add("street", "1234 Ln.")
            .add("zipcode", 86753)
            .add("phones", Json.createArrayBuilder().add("555-5555").add("867-5309").build())
            .build();

    // insert using fromJson on address column, which is a UDT.
    session().execute(insertInto(TABLE2).value("k", key).value("v", fromJson(inputAddr)));

    // retrieve using toJson on address column.
    ResultSet r =
        session()
            .execute(select().toJson("v").as("v").column("k").from(TABLE2).where(eq("k", key)));

    assertThat(r.getAvailableWithoutFetching()).isEqualTo(1);

    Row row = r.one();

    assertThat(row.getString("k")).isEqualTo(key);
    assertThat(row.get("v", JsonStructure.class)).isEqualTo(inputAddr);
  }
}
