/*
 * 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 org.apache.pekko.stream.connectors.pravega;

import docs.javadsl.PravegaBaseTestCase;
import io.pravega.client.ClientConfig;
import io.pravega.client.admin.KeyValueTableManager;
import io.pravega.client.stream.Serializer;
import io.pravega.client.stream.impl.UTF8StringSerializer;
import io.pravega.client.tables.KeyValueTableConfiguration;
import io.pravega.client.tables.TableKey;
import java.nio.ByteBuffer;
import java.util.ArrayList;
import java.util.List;
import java.util.Optional;
import java.util.concurrent.*;
import org.apache.pekko.Done;
import org.apache.pekko.NotUsed;
import org.apache.pekko.japi.Pair;
import org.apache.pekko.stream.connectors.pravega.javadsl.PravegaTable;
import org.apache.pekko.stream.javadsl.Flow;
import org.apache.pekko.stream.javadsl.Keep;
import org.apache.pekko.stream.javadsl.Sink;
import org.apache.pekko.stream.javadsl.Source;
import org.junit.jupiter.api.Assertions;
import org.junit.jupiter.api.BeforeAll;
import org.junit.jupiter.api.Test;

public class PravegaKVTableTestCase extends PravegaBaseTestCase {

  private long timeoutSeconds = 10;

  static String scope = newScope();
  static String tableName = newTableName();

  UTF8StringSerializer serializer = new UTF8StringSerializer();

  Serializer<Integer> intSerializer =
      new Serializer<Integer>() {
        public ByteBuffer serialize(Integer value) {
          ByteBuffer buff = ByteBuffer.allocate(4).putInt(value);
          buff.position(0);
          return buff;
        }

        public Integer deserialize(ByteBuffer serializedValue) {

          return serializedValue.getInt();
        }
      };

  TableWriterSettings<Integer, String> tablewriterSettings =
      TableWriterSettingsBuilder.<Integer, String>create(system, intSerializer, serializer)
          .withKeyExtractor(id -> new TableKey(intSerializer.serialize(id)))
          .build();

  @Test
  public void writeAndReadInKVTable()
      throws ExecutionException, InterruptedException, TimeoutException {

    final List<Pair<Integer, String>> events =
        List.of(
            new Pair<Integer, String>(1, "One"),
            new Pair<Integer, String>(2, "Two"),
            new Pair<Integer, String>(3, "Three"),
            new Pair<Integer, String>(4, "Four"));

    Sink<Pair<Integer, String>, CompletionStage<Done>> sink =
        PravegaTable.sink(scope, tableName, tablewriterSettings);

    CompletionStage<Done> done = Source.from(events).toMat(sink, Keep.right()).run(system);

    done.toCompletableFuture().get(timeoutSeconds, TimeUnit.SECONDS);

    TableReaderSettings<Integer, String> tableReaderSettings =
        TableReaderSettingsBuilder.<Integer, String>create(system, intSerializer, serializer)
            .withKeyExtractor(id -> new TableKey(intSerializer.serialize(id)))
            .build();

    final CompletionStage<String> readingDone =
        PravegaTable.source(scope, tableName, tableReaderSettings)
            .runWith(
                Sink.fold(
                    "",
                    (acc, p) -> {
                      if (acc == "") return p.value();
                      return acc + ", " + p.value();
                    }),
                system);

    String result = readingDone.toCompletableFuture().get(timeoutSeconds, TimeUnit.SECONDS);
    Assertions.assertTrue(
        result.equals("One, Two, Three, Four"), "Read 2 elements [%s]".formatted(result));

    Flow<Integer, Optional<String>, NotUsed> readFlow =
        PravegaTable.readFlow(scope, tableName, tableReaderSettings);

    List<Integer> ids = List.of(1, 2, 3, 4);

    CompletionStage<List<String>> readFlowFut =
        Source.from(ids)
            .via(readFlow)
            .runWith(
                Sink.fold(
                    new ArrayList<String>(),
                    (acc, p) -> {
                      acc.add(p.get());
                      return acc;
                    }),
                system);

    List<String> values = readFlowFut.toCompletableFuture().get(timeoutSeconds, TimeUnit.SECONDS);

    Assertions.assertEquals(values, List.of("One", "Two", "Three", "Four"));
  }

  @BeforeAll
  public static void setUp() {
    createScope(scope);

    ClientConfig clientConfig = ClientConfig.builder().build();

    KeyValueTableConfiguration keyValueTableConfig =
        KeyValueTableConfiguration.builder().partitionCount(2).primaryKeyLength(4).build();
    KeyValueTableManager keyValueTableManager = KeyValueTableManager.create(clientConfig);

    if (keyValueTableManager.createKeyValueTable(scope, tableName, keyValueTableConfig))
      LOGGER.info("Created KeyValue table [{}] in scope [{}]", tableName, scope);
    else LOGGER.info("KeyValue table [{}] already exists in scope [{}]", tableName, scope);

    keyValueTableManager.close();
  }
}
