/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * license agreements; and to You under the Apache License, version 2.0:
 *
 *   https://www.apache.org/licenses/LICENSE-2.0
 *
 * This file is part of the Apache Pekko project, which was derived from Akka.
 */

/*
 * Copyright (C) since 2016 Lightbend Inc. <https://www.lightbend.com>
 */

package docs.javadsl;

import org.apache.pekko.Done;
import org.apache.pekko.NotUsed;
import org.apache.pekko.actor.ActorSystem;
import org.apache.pekko.stream.connectors.orientdb.OrientDbWriteMessage;
import org.apache.pekko.stream.connectors.orientdb.OrientDbSourceSettings;
import org.apache.pekko.stream.connectors.orientdb.OrientDbWriteSettings;
import org.apache.pekko.stream.connectors.orientdb.javadsl.OrientDbFlow;
import org.apache.pekko.stream.connectors.orientdb.javadsl.OrientDbSink;
import org.apache.pekko.stream.connectors.orientdb.javadsl.OrientDbSource;
import org.apache.pekko.stream.connectors.testkit.javadsl.LogCapturingJunit4;
import org.apache.pekko.stream.javadsl.Sink;
import org.apache.pekko.stream.javadsl.Source;
import org.apache.pekko.testkit.javadsl.TestKit;
import com.orientechnologies.orient.core.annotation.OVersion;
import com.orientechnologies.orient.core.db.ODatabaseDocumentInternal;
import com.orientechnologies.orient.core.db.ODatabaseRecordThreadLocal;
import com.orientechnologies.orient.core.db.ODatabaseType;
import com.orientechnologies.orient.core.db.OrientDB;
import com.orientechnologies.orient.core.db.OrientDBConfig;
// #init-settings
import com.orientechnologies.orient.core.db.ODatabasePool;
// #init-settings
import com.orientechnologies.orient.core.db.ODatabaseSession;
import com.orientechnologies.orient.core.record.impl.ODocument;
import com.orientechnologies.orient.object.db.OObjectDatabaseTx;
import org.junit.AfterClass;
import org.junit.BeforeClass;
import org.junit.Rule;
import org.junit.Test;

import java.time.Duration;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.List;
import java.util.concurrent.CompletionStage;
import java.util.concurrent.TimeUnit;
import java.util.function.Consumer;
import java.util.stream.Collectors;

import static org.junit.Assert.assertEquals;

public class OrientDbTest {
  @Rule public final LogCapturingJunit4 logCapturing = new LogCapturingJunit4();

  private static OrientDB orientDB;
  private static ODatabasePool oDatabase;
  private static ODatabaseSession client;
  private static ActorSystem system;

  // #init-settings

  private static String url = "remote:127.0.0.1:2424/";
  private static String dbName = "GratefulDeadConcertsJava";
  private static String username = "root";
  private static String password = "root";
  // #init-settings

  private static String sourceClass = "source1";
  private static String sinkClass1 = "sink1";
  private static String sinkClass2 = "sink2";
  private static String sink3 = "sink3";
  private static String sink6 = "sink6";

  // #define-class
  public static class source1 {

    private String book_title;
    @OVersion private Integer version;

    public void setBook_title(String book_title) {
      this.book_title = book_title;
    }

    public String getBook_title() {
      return book_title;
    }
  }

  public static class sink2 {

    private String book_title;
    @OVersion private Integer version;

    public void setBook_title(String book_title) {
      this.book_title = book_title;
    }

    public String getBook_title() {
      return book_title;
    }
  }
  // #define-class

  public static class KafkaOffset {

    private int offset;

    public KafkaOffset(int offset) {
      this.offset = offset;
    }

    public void setOffset(int offset) {
      this.offset = offset;
    }

    public int getOffset() {
      return offset;
    }
  }

  public static class messagesFromKafka {

    private String book_title;

    private KafkaOffset kafkaOffset;

    public messagesFromKafka(String book_title, KafkaOffset kafkaOffset) {
      this.book_title = book_title;
      this.kafkaOffset = kafkaOffset;
    }

    public void setBook_title(String book_title) {
      this.book_title = book_title;
    }

    public String getBook_title() {
      return book_title;
    }

    public void setKafkaOffset(KafkaOffset kafkaOffset) {
      this.kafkaOffset = kafkaOffset;
    }

    public KafkaOffset getKafkaOffset() {
      return kafkaOffset;
    }
  }

  @BeforeClass
  public static void setup() throws Exception {
    system = ActorSystem.create();

    orientDB = new OrientDB(url, username, password, OrientDBConfig.defaultConfig());
    orientDB.createIfNotExists(dbName, ODatabaseType.PLOCAL);

    // #init-settings

    oDatabase = new ODatabasePool(orientDB, dbName, username, password);

    system.registerOnTermination(() -> oDatabase.close());
    // #init-settings
    client = oDatabase.acquire();

    register(sourceClass);

    flush(sourceClass, "book_title", "Akka in Action");
    flush(sourceClass, "book_title", "Programming in Scala");
    flush(sourceClass, "book_title", "Learning Scala");
    flush(sourceClass, "book_title", "Scala for Spark in Production");
    flush(sourceClass, "book_title", "Scala Puzzlers");
    flush(sourceClass, "book_title", "Effective Akka");
    flush(sourceClass, "book_title", "Akka Concurrency");
  }

  @AfterClass
  public static void tearDown() throws Exception {
    unregister(sourceClass);
    unregister(sinkClass1);
    unregister(sinkClass2);
    unregister(sink3);
    unregister(sink6);

    client.close();
    oDatabase.close();
    if (orientDB.exists(dbName)) {
      orientDB.drop(dbName);
    }
    orientDB.close();
    TestKit.shutdownActorSystem(system);
  }

  private static void register(String className) {
    if (!client.getMetadata().getSchema().existsClass(className)) {
      client.getMetadata().getSchema().createClass(className);
    }
  }

  private static void flush(String className, String fieldName, String fieldValue) {
    ODocument oDocument = new ODocument().field(fieldName, fieldValue);
    oDocument.setClassNameIfExists(className);
    client.save(oDocument);
  }

  private static void unregister(String className) {
    if (client.getMetadata().getSchema().existsClass(className)) {
      client.getMetadata().getSchema().dropClass(className);
    }
  }

  @Test
  public void sourceSettings() {
    // #source-settings
    // re-iterating default values
    OrientDbSourceSettings sourceSettings =
        OrientDbSourceSettings.create(oDatabase).withSkip(0).withLimit(10);
    // #source-settings
  }

  @Test
  public void writeSettings() {
    // #write-settings
    // re-iterating default values
    OrientDbWriteSettings updateSettings = OrientDbWriteSettings.create(oDatabase);
    // #write-settings
  }

  @Test
  public void oDocObjectStream() throws Exception {
    // Copy sourceClass to sinkClass1 through ODocument stream
    CompletionStage<Done> f1 =
        OrientDbSource.create(sourceClass, OrientDbSourceSettings.create(oDatabase))
            .map(m -> OrientDbWriteMessage.create(m.oDocument()))
            .groupedWithin(10, Duration.ofMillis(10))
            .runWith(
                OrientDbSink.create(sinkClass1, OrientDbWriteSettings.create(oDatabase)), system);

    f1.toCompletableFuture().get(10, TimeUnit.SECONDS);

    // #run-odocument
    CompletionStage<List<String>> result =
        OrientDbSource.create(sinkClass1, OrientDbSourceSettings.create(oDatabase))
            .map(m -> m.oDocument().<String>field("book_title"))
            .runWith(Sink.seq(), system);
    // #run-odocument

    List<String> res = new ArrayList<>(result.toCompletableFuture().get(10, TimeUnit.SECONDS));

    List<String> expect =
        Arrays.asList(
            "Akka Concurrency",
            "Akka in Action",
            "Effective Akka",
            "Learning Scala",
            "Programming in Scala",
            "Scala Puzzlers",
            "Scala for Spark in Production");

    Collections.sort(res);
    assertEquals(expect, res);
  }

  @Test
  public void typedStream() throws Exception {
    // Copy sourceClass/book to sinkClass2/book through Typed stream
    // #run-typed
    CompletionStage<Done> f1 =
        OrientDbSource.typed(
                sourceClass, OrientDbSourceSettings.create(oDatabase), source1.class, null)
            .map(
                readResult -> {
                  ODatabaseDocumentInternal db =
                      (ODatabaseDocumentInternal) oDatabase.acquire();
                  new OObjectDatabaseTx(db);
                  ODatabaseRecordThreadLocal.instance().set(db);
                  sink2 sink = new sink2();
                  sink.setBook_title(readResult.oDocument().getBook_title());
                  return OrientDbWriteMessage.create(sink);
                })
            .groupedWithin(10, Duration.ofMillis(10))
            .runWith(
                OrientDbSink.typed(
                    sinkClass2, OrientDbWriteSettings.create(oDatabase), sink2.class),
                system);
    // #run-typed

    f1.toCompletableFuture().get(10, TimeUnit.SECONDS);

    CompletionStage<List<String>> f2 =
        OrientDbSource.typed(
                sinkClass2, OrientDbSourceSettings.create(oDatabase), sink2.class, null)
            .map(
                m -> {
                  ODatabaseDocumentInternal db =
                      (ODatabaseDocumentInternal) oDatabase.acquire();
                  new OObjectDatabaseTx(db);
                  ODatabaseRecordThreadLocal.instance().set(db);
                  return m.oDocument().getBook_title();
                })
            .runWith(Sink.seq(), system);

    List<String> result = new ArrayList<>(f2.toCompletableFuture().get(10, TimeUnit.SECONDS));

    List<String> expect =
        Arrays.asList(
            "Akka Concurrency",
            "Akka in Action",
            "Effective Akka",
            "Learning Scala",
            "Programming in Scala",
            "Scala Puzzlers",
            "Scala for Spark in Production");

    Collections.sort(result);
    assertEquals(expect, result);
  }

  @Test
  public void typedStreamWithPassThrough() throws Exception {
    // #kafka-example
    // We're going to pretend we got messages from kafka.
    // After we've written them to OrientDB, we want
    // to commit the offset to Kafka

    List<Integer> committedOffsets = new ArrayList<>();
    List<messagesFromKafka> messagesFromKafkas =
        Arrays.asList(
            new messagesFromKafka("Akka Concurrency", new KafkaOffset(0)),
            new messagesFromKafka("Akka in Action", new KafkaOffset(1)),
            new messagesFromKafka("Effective Akka", new KafkaOffset(2)));

    Consumer<KafkaOffset> commitToKafka =
        new Consumer<KafkaOffset>() {
          @Override
          public void accept(KafkaOffset kafkaOffset) {
            committedOffsets.add(kafkaOffset.getOffset());
          }
        };

    Source.from(messagesFromKafkas)
        .map(
            kafkaMessage -> {
              String book_title = kafkaMessage.getBook_title();
              return OrientDbWriteMessage.create(
                  new ODocument().field("book_title", book_title), kafkaMessage.kafkaOffset);
            })
        .groupedWithin(10, Duration.ofMillis(10))
        .via(OrientDbFlow.createWithPassThrough(sink6, OrientDbWriteSettings.create(oDatabase)))
        .map(
            messages -> {
              ODatabaseDocumentInternal db =
                  (ODatabaseDocumentInternal) oDatabase.acquire();
              new OObjectDatabaseTx(db);
              ODatabaseRecordThreadLocal.instance().set(db);
              messages.stream().forEach(message -> commitToKafka.accept(message.passThrough()));
              return NotUsed.getInstance();
            })
        .runWith(Sink.seq(), system)
        .toCompletableFuture()
        .get(10, TimeUnit.SECONDS);
    // #kafka-example

    assertEquals(Arrays.asList(0, 1, 2), committedOffsets);

    List<Object> result2 =
        OrientDbSource.create(sink6, OrientDbSourceSettings.create(oDatabase), null)
            .map(m -> m.oDocument().field("book_title"))
            .runWith(Sink.seq(), system)
            .toCompletableFuture()
            .get(10, TimeUnit.SECONDS);

    assertEquals(
        messagesFromKafkas.stream()
            .map(m -> m.getBook_title())
            .sorted()
            .collect(Collectors.toList()),
        result2.stream().sorted().collect(Collectors.toList()));
  }

  @Test
  public void flow() throws Exception {
    // Copy sourceClass to sink3 through ODocument stream
    // #run-flow
    CompletionStage<List<List<OrientDbWriteMessage<ODocument, NotUsed>>>> f1 =
        OrientDbSource.create(sourceClass, OrientDbSourceSettings.create(oDatabase), null)
            .map(m -> OrientDbWriteMessage.create(m.oDocument()))
            .groupedWithin(10, Duration.ofMillis(10))
            .via(OrientDbFlow.create(sink3, OrientDbWriteSettings.create(oDatabase)))
            .runWith(Sink.seq(), system);
    // #run-flow

    f1.toCompletableFuture().get(10, TimeUnit.SECONDS);

    // Assert docs in sink3
    CompletionStage<List<String>> f2 =
        OrientDbSource.create(sink3, OrientDbSourceSettings.create(oDatabase), null)
            .map(m -> m.oDocument().<String>field("book_title"))
            .runWith(Sink.seq(), system);

    List<String> result2 = new ArrayList<>(f2.toCompletableFuture().get(10, TimeUnit.SECONDS));

    List<String> expect =
        Arrays.asList(
            "Akka Concurrency",
            "Akka in Action",
            "Effective Akka",
            "Learning Scala",
            "Programming in Scala",
            "Scala Puzzlers",
            "Scala for Spark in Production");

    Collections.sort(result2);
    assertEquals(expect, result2);
  }
}
