/*
 * 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 static org.bson.codecs.configuration.CodecRegistries.fromProviders;
import static org.bson.codecs.configuration.CodecRegistries.fromRegistries;
import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertTrue;

import com.mongodb.client.model.Filters;
import com.mongodb.client.model.InsertManyOptions;
import com.mongodb.client.model.Updates;
import com.mongodb.reactivestreams.client.MongoClient;
import com.mongodb.reactivestreams.client.MongoClients;
import com.mongodb.reactivestreams.client.MongoCollection;
import com.mongodb.reactivestreams.client.MongoDatabase;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.CompletionStage;
import java.util.concurrent.TimeUnit;
import java.util.stream.IntStream;
import org.apache.pekko.Done;
import org.apache.pekko.NotUsed;
import org.apache.pekko.actor.ActorSystem;
import org.apache.pekko.japi.Pair;
import org.apache.pekko.stream.Materializer;
import org.apache.pekko.stream.connectors.mongodb.DocumentReplace;
import org.apache.pekko.stream.connectors.mongodb.DocumentUpdate;
import org.apache.pekko.stream.connectors.mongodb.javadsl.MongoSink;
import org.apache.pekko.stream.connectors.testkit.javadsl.LogCapturingExtension;
import org.apache.pekko.stream.javadsl.Sink;
import org.apache.pekko.stream.javadsl.Source;
import org.apache.pekko.stream.testkit.javadsl.StreamTestKit;
import org.bson.Document;
import org.bson.codecs.ValueCodecProvider;
import org.bson.codecs.configuration.CodecRegistry;
import org.bson.codecs.pojo.PojoCodecProvider;
import org.bson.conversions.Bson;
import org.junit.jupiter.api.*;
import org.junit.jupiter.api.extension.ExtendWith;

@ExtendWith(LogCapturingExtension.class)
public class MongoSinkTest {

  private static ActorSystem system;

  private final MongoClient client;
  private final MongoDatabase db;
  private final MongoCollection<Document> numbersDocumentColl;
  private final MongoCollection<Number> numbersColl;
  private final MongoCollection<Document> domainObjectsDocumentColl;
  private final MongoCollection<DomainObject> domainObjectsColl;

  private final List<Integer> testRange = IntStream.range(1, 10).boxed().toList();

  public MongoSinkTest() {
    system = ActorSystem.create();

    PojoCodecProvider codecProvider =
        PojoCodecProvider.builder().register(Number.class, DomainObject.class).build();
    CodecRegistry codecRegistry =
        fromRegistries(fromProviders(codecProvider, new ValueCodecProvider()));

    client = MongoClients.create("mongodb://localhost:27017");
    db = client.getDatabase("MongoSinkTest");
    numbersColl = db.getCollection("numbers", Number.class).withCodecRegistry(codecRegistry);
    numbersDocumentColl = db.getCollection("numbers");
    domainObjectsColl =
        db.getCollection("domainObjects", DomainObject.class).withCodecRegistry(codecRegistry);
    domainObjectsDocumentColl = db.getCollection("domainObjects");
  }

  private void insertTestRange() throws Exception {
    Source.from(testRange)
        .map(i -> Document.parse("{\"value\":" + i + "}"))
        .runWith(MongoSink.insertOne(numbersDocumentColl), system)
        .toCompletableFuture()
        .get(5, TimeUnit.SECONDS);
  }

  private void insertDomainObjects() throws Exception {
    Source.from(testRange)
        .map(
            i -> {
              DomainObject domainObject =
                  new DomainObject(
                      i, "first-property-%s".formatted(i), "second-property-%s".formatted(i));
              System.out.println("%s inserting %s".formatted(i, domainObject));
              return domainObject;
            })
        .runWith(MongoSink.insertOne(domainObjectsColl), system)
        .toCompletableFuture()
        .get(5, TimeUnit.SECONDS);
  }

  private void deleteCollection(MongoCollection<Document> collection) throws Exception {
    Source.fromPublisher(collection.deleteMany(new Document()))
        .runWith(Sink.head(), system)
        .toCompletableFuture()
        .get(5, TimeUnit.SECONDS);
  }

  private void checkCollectionForLeaks(MongoCollection<Document> collection) throws Exception {
    Source.fromPublisher(collection.deleteMany(new Document()))
        .runWith(Sink.head(), system)
        .toCompletableFuture()
        .get(5, TimeUnit.SECONDS);
  }

  @BeforeEach
  public void cleanDb() throws Exception {
    deleteCollection(numbersDocumentColl);
    deleteCollection(domainObjectsDocumentColl);
  }

  @AfterEach
  public void checkForLeaks() throws Exception {
    checkCollectionForLeaks(numbersDocumentColl);
    checkCollectionForLeaks(domainObjectsDocumentColl);
    StreamTestKit.assertAllStagesStopped(Materializer.matFromSystem(system));
  }

  @AfterAll
  public static void terminateActorSystem() {
    system.terminate();
  }

  @Test
  public void saveWithInsertOne() throws Exception {
    final Source<Document, NotUsed> source =
        Source.from(testRange).map(i -> Document.parse("{\"value\":" + i + "}"));
    final CompletionStage<Done> completion =
        source.runWith(MongoSink.insertOne(numbersDocumentColl), system);

    completion.toCompletableFuture().get(5, TimeUnit.SECONDS);
    final CompletionStage<List<Document>> found =
        Source.fromPublisher(numbersDocumentColl.find()).runWith(Sink.seq(), system);

    assertEquals(
        testRange,
        found.toCompletableFuture().get(5, TimeUnit.SECONDS).stream()
            .map(n -> n.getInteger("value"))
            .toList());
  }

  @Test
  public void saveWithInsertOneAndCodecSupport() throws Exception {
    // #insert-one
    List<Number> testRangeObjects = testRange.stream().map(Number::new).toList();
    final CompletionStage<Done> completion =
        Source.from(testRangeObjects).runWith(MongoSink.insertOne(numbersColl), system);
    // #insert-one

    completion.toCompletableFuture().get(5, TimeUnit.SECONDS);
    final CompletionStage<List<Number>> found =
        Source.fromPublisher(numbersColl.find()).runWith(Sink.seq(), system);

    assertEquals(testRangeObjects, found.toCompletableFuture().get(5, TimeUnit.SECONDS));
  }

  @Test
  public void saveWithInsertMany() throws Exception {
    final Source<Document, NotUsed> source =
        Source.from(testRange).map(i -> Document.parse("{\"value\":" + i + "}"));
    final CompletionStage<Done> completion =
        source.grouped(2).runWith(MongoSink.insertMany(numbersDocumentColl), system);

    completion.toCompletableFuture().get(5, TimeUnit.SECONDS);
    final CompletionStage<List<Document>> found =
        Source.fromPublisher(numbersDocumentColl.find()).runWith(Sink.seq(), system);

    assertEquals(
        testRange,
        found.toCompletableFuture().get(5, TimeUnit.SECONDS).stream()
            .map(n -> n.getInteger("value"))
            .toList());
  }

  @Test
  public void saveWithInsertManyAndCodecSupport() throws Exception {
    // #insert-many
    final List<Number> testRangeObjects = testRange.stream().map(Number::new).toList();
    final CompletionStage<Done> completion =
        Source.from(testRangeObjects).grouped(2).runWith(MongoSink.insertMany(numbersColl), system);
    // #insert-many

    completion.toCompletableFuture().get(5, TimeUnit.SECONDS);
    final CompletionStage<List<Number>> found =
        Source.fromPublisher(numbersColl.find()).runWith(Sink.seq(), system);

    assertEquals(testRangeObjects, found.toCompletableFuture().get(5, TimeUnit.SECONDS));
  }

  @Test
  public void saveWithInsertManyWithOptions() throws Exception {
    final Source<Document, NotUsed> source =
        Source.from(testRange).map(i -> Document.parse("{\"value\":" + i + "}"));
    final CompletionStage<Done> completion =
        source
            .grouped(2)
            .runWith(
                MongoSink.insertMany(numbersDocumentColl, new InsertManyOptions().ordered(false)),
                system);

    completion.toCompletableFuture().get(5, TimeUnit.SECONDS);
    final CompletionStage<List<Document>> found =
        Source.fromPublisher(numbersDocumentColl.find()).runWith(Sink.seq(), system);

    assertEquals(
        testRange,
        found.toCompletableFuture().get(5, TimeUnit.SECONDS).stream()
            .map(n -> n.getInteger("value"))
            .toList());
  }

  @Test
  public void saveWithInsertManyWithOptionsAndCodecSupport() throws Exception {
    List<Number> testRangeObjects = testRange.stream().map(Number::new).toList();
    final CompletionStage<Done> completion =
        Source.from(testRangeObjects)
            .grouped(2)
            .runWith(
                MongoSink.insertMany(numbersColl, new InsertManyOptions().ordered(false)), system);

    completion.toCompletableFuture().get(5, TimeUnit.SECONDS);
    final CompletionStage<List<Number>> found =
        Source.fromPublisher(numbersColl.find()).runWith(Sink.seq(), system);

    assertEquals(testRangeObjects, found.toCompletableFuture().get(5, TimeUnit.SECONDS));
  }

  @Test
  public void updateWithUpdateOne() throws Exception {
    insertTestRange();

    // #update-one
    final Source<DocumentUpdate, NotUsed> source =
        Source.from(testRange)
            .map(
                i ->
                    DocumentUpdate.create(
                        Filters.eq("value", i), Updates.set("updateValue", i * -1)));
    final CompletionStage<Done> completion =
        source.runWith(MongoSink.updateOne(numbersDocumentColl), system);
    // #update-one

    completion.toCompletableFuture().get(5, TimeUnit.SECONDS);

    final CompletionStage<List<Document>> found =
        Source.fromPublisher(numbersDocumentColl.find()).runWith(Sink.seq(), system);

    assertEquals(
        testRange.stream().map(i -> Pair.create(i, i * -1)).toList(),
        found.toCompletableFuture().get(5, TimeUnit.SECONDS).stream()
            .map(d -> Pair.create(d.getInteger("value"), d.getInteger("updateValue")))
            .toList());
  }

  @Test
  public void updateWithUpdateMany() throws Exception {
    insertTestRange();

    final Source<DocumentUpdate, NotUsed> source =
        Source.single(0)
            .map(
                i -> DocumentUpdate.create(Filters.gte("value", 0), Updates.set("updateValue", 0)));
    final CompletionStage<Done> completion =
        source.runWith(MongoSink.updateMany(numbersDocumentColl), system);

    completion.toCompletableFuture().get(5, TimeUnit.SECONDS);

    final CompletionStage<List<Document>> found =
        Source.fromPublisher(numbersDocumentColl.find()).runWith(Sink.seq(), system);

    assertEquals(
        testRange.stream().map(i -> Pair.create(i, 0)).toList(),
        found.toCompletableFuture().get(5, TimeUnit.SECONDS).stream()
            .map(d -> Pair.create(d.getInteger("value"), d.getInteger("updateValue")))
            .toList());
  }

  @Test
  public void deleteWithDeleteOne() throws Exception {
    insertTestRange();

    // #delete-one
    final Source<Bson, NotUsed> source = Source.from(testRange).map(i -> Filters.eq("value", i));
    final CompletionStage<Done> completion =
        source.runWith(MongoSink.deleteOne(numbersDocumentColl), system);
    // #delete-one

    completion.toCompletableFuture().get(5, TimeUnit.SECONDS);

    final CompletionStage<List<Document>> found =
        Source.fromPublisher(numbersDocumentColl.find()).runWith(Sink.seq(), system);

    assertTrue(found.toCompletableFuture().get(5, TimeUnit.SECONDS).isEmpty());
  }

  @Test
  public void deleteWithDeleteMany() throws Exception {
    insertTestRange();

    final Source<Bson, NotUsed> source = Source.single(0).map(i -> Filters.gte("value", 0));
    final CompletionStage<Done> completion =
        source.runWith(MongoSink.deleteMany(numbersDocumentColl), system);

    completion.toCompletableFuture().get(5, TimeUnit.SECONDS);

    final CompletionStage<List<Document>> found =
        Source.fromPublisher(numbersDocumentColl.find()).runWith(Sink.seq(), system);

    assertTrue(found.toCompletableFuture().get(5, TimeUnit.SECONDS).isEmpty());
  }

  @Test
  public void replaceWithReplaceOne() throws Exception {
    insertDomainObjects();

    // #replace-one
    final Source<DocumentReplace<DomainObject>, NotUsed> source =
        Source.from(testRange)
            .map(
                i ->
                    DocumentReplace.create(
                        Filters.eq("_id", i),
                        new DomainObject(
                            i,
                            "updated-first-property-%s".formatted(i),
                            "updated-second-property-%s".formatted(i))));
    final CompletionStage<Done> completion =
        source.runWith(MongoSink.replaceOne(domainObjectsColl), system);
    // #replace-one

    completion.toCompletableFuture().get(5, TimeUnit.SECONDS);

    final List<DomainObject> found =
        new ArrayList<>(
            Source.fromPublisher(domainObjectsColl.find())
                .runWith(Sink.seq(), system)
                .toCompletableFuture()
                .get(5, TimeUnit.SECONDS));

    final List<DomainObject> expected =
        testRange.stream()
            .map(
                i ->
                    new DomainObject(
                        i,
                        "updated-first-property-%s".formatted(i),
                        "updated-second-property-%s".formatted(i)))
            .toList();

    assertEquals(expected, found);
  }
}
