/*
 * 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.NotUsed;
import org.apache.pekko.actor.ActorSystem;
import org.apache.pekko.actor.Cancellable;
import org.apache.pekko.japi.Pair;
import org.apache.pekko.stream.connectors.hdfs.*;
import org.apache.pekko.stream.connectors.hdfs.javadsl.HdfsFlow;
import org.apache.pekko.stream.connectors.hdfs.util.JavaTestUtils;
import org.apache.pekko.stream.connectors.testkit.javadsl.LogCapturingJunit4;
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.apache.pekko.testkit.javadsl.TestKit;
import org.apache.pekko.util.ByteString;
import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.fs.FileStatus;
import org.apache.hadoop.fs.FileSystem;
import org.apache.hadoop.fs.Path;
import org.apache.hadoop.hdfs.MiniDFSCluster;
import org.apache.hadoop.io.SequenceFile;
import org.apache.hadoop.io.Text;
import org.apache.hadoop.io.compress.DefaultCodec;
import org.junit.*;

import scala.concurrent.duration.Duration;

import java.io.IOException;
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.BiFunction;
import java.util.stream.Collectors;

import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertTrue;

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

  private static MiniDFSCluster hdfsCluster = null;
  private static ActorSystem system;
  private static String destination = JavaTestUtils.destination();
  private static List<ByteString> books = JavaTestUtils.books();
  private static FileSystem fs = null;
  private static HdfsWritingSettings settings = HdfsWritingSettings.create();

  // #define-kafka-classes
  public static class Book {
    final String title;

    Book(String title) {
      this.title = title;
    }
  }

  static class KafkaCommitter {
    List<Integer> committedOffsets = new ArrayList<>();

    void commit(KafkaOffset offset) {
      committedOffsets.add(offset.offset);
    }
  }

  static class KafkaOffset {
    final int offset;

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

  static class KafkaMessage {
    final Book book;
    final KafkaOffset offset;

    KafkaMessage(Book book, KafkaOffset offset) {
      this.book = book;
      this.offset = offset;
    }
  }
  // #define-kafka-classes

  @Test
  public void testDataWriterFileSizeRotationWithFiveFile() throws Exception {
    Flow<HdfsWriteMessage<ByteString, NotUsed>, RotationMessage, NotUsed> flow =
        HdfsFlow.data(
            fs, SyncStrategy.count(50), RotationStrategy.size(0.01, FileUnit.KB()), settings);

    CompletionStage<List<RotationMessage>> resF =
        Source.from(books).map(HdfsWriteMessage::create).via(flow).runWith(Sink.seq(), system);

    List<RotationMessage> result = new ArrayList<>(resF.toCompletableFuture().get());
    List<RotationMessage> expect =
        Arrays.asList(
            new RotationMessage(output("0"), 0),
            new RotationMessage(output("1"), 1),
            new RotationMessage(output("2"), 2),
            new RotationMessage(output("3"), 3),
            new RotationMessage(output("4"), 4));

    assertEquals(expect, result);
  }

  @Test
  public void testDataWriterFileSizeRotationWithTwoFile() throws Exception {
    List<ByteString> data = JavaTestUtils.generateFakeContent(1.0, FileUnit.KB().byteCount());

    Flow<HdfsWriteMessage<ByteString, NotUsed>, RotationMessage, NotUsed> flow =
        HdfsFlow.data(
            fs, SyncStrategy.count(500), RotationStrategy.size(0.5, FileUnit.KB()), settings);

    CompletionStage<List<RotationMessage>> resF =
        Source.from(data).map(HdfsWriteMessage::create).via(flow).runWith(Sink.seq(), system);

    List<RotationMessage> logs = new ArrayList<>(resF.toCompletableFuture().get());

    assertEquals(logs.size(), 2);

    List<FileStatus> files = JavaTestUtils.getFiles(fs);
    assertEquals(
        files.stream().map(FileStatus::getLen).reduce((a, b) -> a + b).get().longValue(), 1024L);
    files.forEach(file -> assertTrue(file.getLen() <= 512L));
    JavaTestUtils.verifyOutputFileSize(fs, logs);
    JavaTestUtils.verifyFlattenContent(fs, logs, data);
  }

  @Test
  public void testDataWriterDetectUpstreamFinish() throws Exception {
    List<ByteString> data = JavaTestUtils.generateFakeContent(1.0, FileUnit.KB().byteCount());

    // #define-data
    Flow<HdfsWriteMessage<ByteString, NotUsed>, RotationMessage, NotUsed> flow =
        HdfsFlow.data(
            fs, SyncStrategy.count(500), RotationStrategy.size(1, FileUnit.GB()), settings);
    // #define-data
    CompletionStage<List<RotationMessage>> resF =
        Source.from(data).map(HdfsWriteMessage::create).via(flow).runWith(Sink.seq(), system);

    List<RotationMessage> logs = new ArrayList<>(resF.toCompletableFuture().get());
    assertEquals(logs.size(), 1);

    assertEquals(JavaTestUtils.getFiles(fs).get(0).getLen(), 1024L);
    JavaTestUtils.verifyOutputFileSize(fs, logs);
    JavaTestUtils.verifyFlattenContent(fs, logs, data);
  }

  @Test
  public void testDataWriterWithBufferRotation() throws Exception {
    Flow<HdfsWriteMessage<ByteString, NotUsed>, RotationMessage, NotUsed> flow =
        HdfsFlow.data(fs, SyncStrategy.count(1), RotationStrategy.count(2), settings);

    CompletionStage<List<RotationMessage>> resF =
        Source.from(books).map(HdfsWriteMessage::create).via(flow).runWith(Sink.seq(), system);

    List<RotationMessage> logs = new ArrayList<>(resF.toCompletableFuture().get());

    assertEquals(logs.size(), 3);
    JavaTestUtils.verifyFlattenContent(fs, logs, books);
  }

  @Test
  public void testDataWriterWithTimeRotation() throws Exception {
    Pair<Cancellable, CompletionStage<List<RotationMessage>>> resF =
        Source.tick(
                java.time.Duration.ofMillis(0),
                java.time.Duration.ofMillis(50),
                ByteString.fromString("I love Pekko Connectors!"))
            .map(HdfsWriteMessage::create)
            .via(
                HdfsFlow.data(
                    fs,
                    SyncStrategy.none(),
                    RotationStrategy.time(Duration.create(500, TimeUnit.MILLISECONDS)),
                    settings))
            .toMat(Sink.seq(), Keep.both())
            .run(system);

    system
        .scheduler()
        .scheduleOnce(
            java.time.Duration.ofMillis(1500),
            () -> resF.first().cancel(),
            system.dispatcher()); // cancel within 1500 milliseconds

    List<RotationMessage> logs = new ArrayList<>(resF.second().toCompletableFuture().get());
    JavaTestUtils.verifyOutputFileSize(fs, logs);
    assertTrue(logs.size() == 3 || logs.size() == 4);
  }

  @Test
  public void testDataWriterWithNoRotation() throws Exception {
    Flow<HdfsWriteMessage<ByteString, NotUsed>, RotationMessage, NotUsed> flow =
        HdfsFlow.data(fs, SyncStrategy.none(), RotationStrategy.none(), settings);

    CompletionStage<List<RotationMessage>> resF =
        Source.from(books).map(HdfsWriteMessage::create).via(flow).runWith(Sink.seq(), system);

    List<RotationMessage> logs = new ArrayList<>(resF.toCompletableFuture().get());

    JavaTestUtils.verifyOutputFileSize(fs, logs);
    List<Integer> list = new ArrayList<>();
    books.iterator().forEachRemaining(s -> list.add(s.toArray().length));
    assertEquals(
        JavaTestUtils.getFiles(fs).get(0).getLen(),
        list.stream().mapToInt(Integer::intValue).sum());
  }

  @Test
  public void testDataWriterKafkaExample() throws Exception {
    // #kafka-example
    // We're going to pretend we got messages from kafka.
    // After we've written them to HDFS, we want
    // to commit the offset to Kafka
    List<KafkaMessage> messagesFromKafka =
        Arrays.asList(
            new KafkaMessage(new Book("Pekko Concurrency"), new KafkaOffset(0)),
            new KafkaMessage(new Book("Pekko in Action"), new KafkaOffset(1)),
            new KafkaMessage(new Book("Effective Pekko"), new KafkaOffset(2)),
            new KafkaMessage(new Book("Learning Scala"), new KafkaOffset(3)),
            new KafkaMessage(new Book("Scala Puzzlers"), new KafkaOffset(4)),
            new KafkaMessage(new Book("Scala for Spark in Production"), new KafkaOffset(5)));

    final KafkaCommitter kafkaCommitter = new KafkaCommitter();

    Flow<HdfsWriteMessage<ByteString, KafkaOffset>, OutgoingMessage<KafkaOffset>, NotUsed> flow =
        HdfsFlow.dataWithPassThrough(
            fs,
            SyncStrategy.count(50),
            RotationStrategy.count(4),
            HdfsWritingSettings.create().withNewLine(true));

    CompletionStage<List<RotationMessage>> resF =
        Source.from(messagesFromKafka)
            .map(
                kafkaMessage -> {
                  Book book = kafkaMessage.book;
                  // Transform message so that we can write to hdfs\
                  return HdfsWriteMessage.create(
                      ByteString.fromString(book.title), kafkaMessage.offset);
                })
            .via(flow)
            .map(
                message -> {
                  if (message instanceof WrittenMessage<KafkaOffset> writtenMessage) {
                    kafkaCommitter.commit(writtenMessage.passThrough());
                    return message;
                  } else {
                    return message;
                  }
                })
            .collectType(RotationMessage.class) // Collect only rotation messages
            .runWith(Sink.seq(), system);
    // #kafka-example

    ArrayList<RotationMessage> logs = new ArrayList<>(resF.toCompletableFuture().get());
    List<RotationMessage> expect =
        Arrays.asList(new RotationMessage(output("0"), 0), new RotationMessage(output("1"), 1));

    // Make sure all messages was committed to kafka
    assertEquals(Arrays.asList(0, 1, 2, 3, 4, 5), kafkaCommitter.committedOffsets);

    assertEquals(logs, expect);
    JavaTestUtils.verifyOutputFileSize(fs, logs);
    assertEquals(
        JavaTestUtils.readLogs(fs, logs).stream()
            .map(string -> string.split("\n"))
            .flatMap(Arrays::stream)
            .collect(Collectors.toList()),
        messagesFromKafka.stream().map(message -> message.book.title).collect(Collectors.toList()));
  }

  @Test
  public void testCompressedDataWriterWithSizeRotation() throws Exception {
    // #define-codec
    DefaultCodec codec = new DefaultCodec();
    codec.setConf(fs.getConf());
    // #define-codec

    // #define-compress
    Flow<HdfsWriteMessage<ByteString, NotUsed>, RotationMessage, NotUsed> flow =
        HdfsFlow.compressed(
            fs, SyncStrategy.count(50), RotationStrategy.size(0.1, FileUnit.MB()), codec, settings);
    // #define-compress

    List<ByteString> content =
        JavaTestUtils.generateFakeContentWithPartitions(1, FileUnit.MB().byteCount(), 30);

    CompletionStage<List<RotationMessage>> resF =
        Source.fromIterator(content::iterator)
            .map(HdfsWriteMessage::create)
            .via(flow)
            .runWith(Sink.seq(), system);

    List<RotationMessage> logs = new ArrayList<>(resF.toCompletableFuture().get());
    List<RotationMessage> expect =
        Arrays.asList(
            new RotationMessage(output("0.deflate"), 0),
            new RotationMessage(output("1.deflate"), 1),
            new RotationMessage(output("2.deflate"), 2),
            new RotationMessage(output("3.deflate"), 3),
            new RotationMessage(output("4.deflate"), 4),
            new RotationMessage(output("5.deflate"), 5));

    assertEquals(logs, expect);
    JavaTestUtils.verifyOutputFileSize(fs, logs);
    JavaTestUtils.verifyLogsWithCodec(fs, content, logs, codec);
  }

  @Test
  public void testCompressedDataWriterWithBufferRotation() throws Exception {
    DefaultCodec codec = new DefaultCodec();
    codec.setConf(fs.getConf());

    Flow<HdfsWriteMessage<ByteString, NotUsed>, RotationMessage, NotUsed> flow =
        HdfsFlow.compressed(fs, SyncStrategy.count(1), RotationStrategy.count(1), codec, settings);

    CompletionStage<List<RotationMessage>> resF =
        Source.from(books).map(HdfsWriteMessage::create).via(flow).runWith(Sink.seq(), system);

    List<RotationMessage> logs = new ArrayList<>(resF.toCompletableFuture().get());
    List<RotationMessage> expect =
        Arrays.asList(
            new RotationMessage(output("0.deflate"), 0),
            new RotationMessage(output("1.deflate"), 1),
            new RotationMessage(output("2.deflate"), 2),
            new RotationMessage(output("3.deflate"), 3),
            new RotationMessage(output("4.deflate"), 4));

    assertEquals(logs, expect);
    JavaTestUtils.verifyOutputFileSize(fs, logs);
    JavaTestUtils.verifyLogsWithCodec(fs, books, logs, codec);
  }

  @Test
  public void testCompressedDataWriterWithNoRotation() throws Exception {
    DefaultCodec codec = new DefaultCodec();
    codec.setConf(fs.getConf());

    Flow<HdfsWriteMessage<ByteString, NotUsed>, RotationMessage, NotUsed> flow =
        HdfsFlow.compressed(fs, SyncStrategy.none(), RotationStrategy.none(), codec, settings);

    List<ByteString> content =
        JavaTestUtils.generateFakeContentWithPartitions(1, FileUnit.MB().byteCount(), 30);

    CompletionStage<List<RotationMessage>> resF =
        Source.fromIterator(content::iterator)
            .map(HdfsWriteMessage::create)
            .via(flow)
            .runWith(Sink.seq(), system);

    List<RotationMessage> logs = new ArrayList<>(resF.toCompletableFuture().get());
    List<RotationMessage> expect =
        Collections.singletonList(new RotationMessage(output("0.deflate"), 0));

    assertEquals(logs, expect);
    JavaTestUtils.verifyOutputFileSize(fs, logs);
    JavaTestUtils.verifyLogsWithCodec(fs, content, logs, codec);
  }

  @Test
  public void testSequenceWriterWithSizeRotationWithoutCompression() throws Exception {
    // #define-sequence-compressed
    Flow<HdfsWriteMessage<Pair<Text, Text>, NotUsed>, RotationMessage, NotUsed> flow =
        HdfsFlow.sequence(
            fs,
            SyncStrategy.none(),
            RotationStrategy.size(1, FileUnit.MB()),
            settings,
            Text.class,
            Text.class);
    // #define-sequence-compressed

    List<Pair<Text, Text>> content =
        JavaTestUtils.generateFakeContentForSequence(0.5, FileUnit.MB().byteCount());

    CompletionStage<List<RotationMessage>> resF =
        Source.fromIterator(content::iterator)
            .map(HdfsWriteMessage::create)
            .via(flow)
            .runWith(Sink.seq(), system);

    List<RotationMessage> logs = new ArrayList<>(resF.toCompletableFuture().get());

    JavaTestUtils.verifyOutputFileSize(fs, logs);
    JavaTestUtils.verifySequenceFile(fs, content, logs);
  }

  @Test
  public void testSequenceWriterWithSizeRotationWithCompression() throws Exception {
    DefaultCodec codec = new DefaultCodec();
    codec.setConf(fs.getConf());

    // #define-sequence
    Flow<HdfsWriteMessage<Pair<Text, Text>, NotUsed>, RotationMessage, NotUsed> flow =
        HdfsFlow.sequence(
            fs,
            SyncStrategy.none(),
            RotationStrategy.size(1, FileUnit.MB()),
            SequenceFile.CompressionType.BLOCK,
            codec,
            settings,
            Text.class,
            Text.class);
    // #define-sequence

    List<Pair<Text, Text>> content =
        JavaTestUtils.generateFakeContentForSequence(0.5, FileUnit.MB().byteCount());

    CompletionStage<List<RotationMessage>> resF =
        Source.fromIterator(content::iterator)
            .map(HdfsWriteMessage::create)
            .via(flow)
            .runWith(Sink.seq(), system);

    List<RotationMessage> logs = new ArrayList<>(resF.toCompletableFuture().get());

    JavaTestUtils.verifyOutputFileSize(fs, logs);
    JavaTestUtils.verifySequenceFile(fs, content, logs);
  }

  @Test
  public void testSequenceWriterWithBufferRotation() throws Exception {
    Flow<HdfsWriteMessage<Pair<Text, Text>, NotUsed>, RotationMessage, NotUsed> flow =
        HdfsFlow.sequence(
            fs, SyncStrategy.none(), RotationStrategy.count(1), settings, Text.class, Text.class);

    List<Pair<Text, Text>> content = JavaTestUtils.booksForSequenceWriter();

    CompletionStage<List<RotationMessage>> resF =
        Source.fromIterator(content::iterator)
            .map(HdfsWriteMessage::create)
            .via(flow)
            .runWith(Sink.seq(), system);

    List<RotationMessage> logs = new ArrayList<>(resF.toCompletableFuture().get());

    assertEquals(logs.size(), 5);
    JavaTestUtils.verifyOutputFileSize(fs, logs);
    JavaTestUtils.verifySequenceFile(fs, content, logs);
  }

  @Test
  public void testSequenceWriterWithNoRotation() throws Exception {
    Flow<HdfsWriteMessage<Pair<Text, Text>, NotUsed>, RotationMessage, NotUsed> flow =
        HdfsFlow.sequence(
            fs, SyncStrategy.none(), RotationStrategy.none(), settings, Text.class, Text.class);

    List<Pair<Text, Text>> content =
        JavaTestUtils.generateFakeContentForSequence(0.5, FileUnit.MB().byteCount());

    CompletionStage<List<RotationMessage>> resF =
        Source.fromIterator(content::iterator)
            .map(HdfsWriteMessage::create)
            .via(flow)
            .runWith(Sink.seq(), system);

    List<RotationMessage> logs = new ArrayList<>(resF.toCompletableFuture().get());

    assertEquals(logs.size(), 1);
    JavaTestUtils.verifyOutputFileSize(fs, logs);
    JavaTestUtils.verifySequenceFile(fs, content, logs);
  }

  @BeforeClass
  public static void setup() throws Exception {
    hdfsCluster = JavaTestUtils.setupCluster();

    // #init-client
    Configuration conf = new Configuration();
    conf.set("fs.default.name", "hdfs://localhost:54310");

    fs = FileSystem.get(conf);
    // #init-client

    system = ActorSystem.create();
  }

  @AfterClass
  public static void teardown() throws Exception {
    fs.close();
    hdfsCluster.shutdown();
    TestKit.shutdownActorSystem(system);
  }

  @After
  public void afterEach() throws IOException {
    fs.delete(new Path(destination), true);
    fs.delete(settings.pathGenerator().apply(0L, 0L).getParent(), true);
  }

  private static String output(String s) {
    return destination + s;
  }

  private static void documentation() {
    // #define-generator
    BiFunction<Long, Long, String> func =
        (rotationCount, timestamp) -> "/tmp/pekko-connectors/" + rotationCount + "-" + timestamp;
    FilePathGenerator pathGenerator = FilePathGenerator.create(func);
    // #define-generator
    // #define-settings
    HdfsWritingSettings.create()
        .withOverwrite(true)
        .withNewLine(false)
        .withLineSeparator(System.getProperty("line.separator"))
        .withPathGenerator(pathGenerator);
    // #define-settings
  }
}
