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

package akka.stream.javadsl;

import akka.NotUsed;
import akka.japi.Pair;
import akka.stream.*;
import akka.testkit.AkkaJUnitActorSystemResource;
import akka.testkit.AkkaSpec;
import org.junit.ClassRule;
import org.junit.Test;
import scala.collection.Seq;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashSet;
import java.util.List;
import java.util.concurrent.CompletionStage;
import java.util.concurrent.TimeUnit;

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

public class GraphDslTest extends StreamTest {
  public GraphDslTest() {
    super(actorSystemResource);
  }

  @ClassRule
  public static AkkaJUnitActorSystemResource actorSystemResource = new AkkaJUnitActorSystemResource("GraphDslTest",
          AkkaSpec.testConf());

  @Test
  public void demonstrateBuildSimpleGraph() throws Exception {
    //#simple-graph-dsl
    final Source<Integer, NotUsed> in = Source.from(Arrays.asList(1, 2, 3, 4, 5));
    final Sink<List<String>, CompletionStage<List<String>>> sink = Sink.head();
    final Flow<Integer, Integer, NotUsed> f1 = Flow.of(Integer.class).map(elem -> elem + 10);
    final Flow<Integer, Integer, NotUsed> f2 = Flow.of(Integer.class).map(elem -> elem + 20);
    final Flow<Integer, String, NotUsed> f3 = Flow.of(Integer.class).map(elem -> elem.toString());
    final Flow<Integer, Integer, NotUsed> f4 = Flow.of(Integer.class).map(elem -> elem + 30);

    final RunnableGraph<CompletionStage<List<String>>> result =
            RunnableGraph.fromGraph(
                    GraphDSL     // create() function binds sink, out which is sink's out port and builder DSL
                            .create(   // we need to reference out's shape in the builder DSL below (in to() function)
                                    sink,                // previously created sink (Sink)
                                    (builder, out) -> {  // variables: builder (GraphDSL.Builder) and out (SinkShape)
                                      final UniformFanOutShape<Integer, Integer> bcast = builder.add(Broadcast.create(2));
                                      final UniformFanInShape<Integer, Integer> merge = builder.add(Merge.create(2));

                                      final Outlet<Integer> source = builder.add(in).out();
                                      builder.from(source).via(builder.add(f1))
                                              .viaFanOut(bcast).via(builder.add(f2)).viaFanIn(merge)
                                              .via(builder.add(f3.grouped(1000))).to(out);  // to() expects a SinkShape
                                      builder.from(bcast).via(builder.add(f4)).toFanIn(merge);
                                      return ClosedShape.getInstance();
                                    }));
    //#simple-graph-dsl
    final List<String> list = result.run(materializer).toCompletableFuture().get(3, TimeUnit.SECONDS);
    final String[] res = list.toArray(new String[]{});
    Arrays.sort(res, null);
    assertArrayEquals(new String[]{"31", "32", "33", "34", "35", "41", "42", "43", "44", "45"}, res);
  }

  @Test
  @SuppressWarnings("unused")
  public void demonstrateConnectErrors() {
    try {
      //#simple-graph
      final RunnableGraph<NotUsed> g =
              RunnableGraph.<NotUsed>fromGraph(
                      GraphDSL
                              .create((b) -> {
                                        final SourceShape<Integer> source1 = b.add(Source.from(Arrays.asList(1, 2, 3, 4, 5)));
                                        final SourceShape<Integer> source2 = b.add(Source.from(Arrays.asList(1, 2, 3, 4, 5)));
                                        final FanInShape2<Integer, Integer, Pair<Integer, Integer>> zip = b.add(Zip.create());
                                        b.from(source1).toInlet(zip.in0());
                                        b.from(source2).toInlet(zip.in1());
                                        return ClosedShape.getInstance();
                                      }
                              )
              );
      // unconnected zip.out (!) => "The inlets [] and outlets [] must correspond to the inlets [] and outlets [ZipWith2.out]"
      //#simple-graph
      org.junit.Assert.fail("expected IllegalArgumentException");
    } catch (IllegalStateException e) {
      assertTrue(e != null && e.getMessage() != null && e.getMessage().contains("ZipWith2.out"));
    }
  }

  @Test
  public void demonstrateReusingFlowInGraph() throws Exception {
    //#graph-dsl-reusing-a-flow
    final Sink<Integer, CompletionStage<Integer>> topHeadSink = Sink.head();
    final Sink<Integer, CompletionStage<Integer>> bottomHeadSink = Sink.head();
    final Flow<Integer, Integer, NotUsed> sharedDoubler = Flow.of(Integer.class).map(elem -> elem * 2);

    final RunnableGraph<Pair<CompletionStage<Integer>, CompletionStage<Integer>>> g =
            RunnableGraph.<Pair<CompletionStage<Integer>, CompletionStage<Integer>>>fromGraph(
                    GraphDSL.create(
                            topHeadSink, // import this sink into the graph
                            bottomHeadSink, // and this as well
                            Keep.both(),
                            (b, top, bottom) -> {
                              final UniformFanOutShape<Integer, Integer> bcast =
                                      b.add(Broadcast.create(2));

                              b.from(b.add(Source.single(1))).viaFanOut(bcast)
                                      .via(b.add(sharedDoubler)).to(top);
                              b.from(bcast).via(b.add(sharedDoubler)).to(bottom);
                              return ClosedShape.getInstance();
                            }
                    )
            );
    //#graph-dsl-reusing-a-flow
    final Pair<CompletionStage<Integer>, CompletionStage<Integer>> pair = g.run(materializer);
    assertEquals(Integer.valueOf(2), pair.first().toCompletableFuture().get(3, TimeUnit.SECONDS));
    assertEquals(Integer.valueOf(2), pair.second().toCompletableFuture().get(3, TimeUnit.SECONDS));
  }

  @Test
  public void demonstrateMatValue() throws Exception {
    //#graph-dsl-matvalue
    final Sink<Integer, CompletionStage<Integer>> foldSink = Sink.<Integer, Integer>fold(0, (a, b) -> {
      return a + b;
    });

    final Flow<CompletionStage<Integer>, Integer, NotUsed> flatten =
            Flow.<CompletionStage<Integer>>create().mapAsync(4, x -> x);

    final Flow<Integer, Integer, CompletionStage<Integer>> foldingFlow = Flow.fromGraph(
            GraphDSL.create(foldSink,
                    (b, fold) -> {
                      return FlowShape.of(
                              fold.in(),
                              b.from(b.materializedValue()).via(b.add(flatten)).out());
                    }));
    //#graph-dsl-matvalue

    //#graph-dsl-matvalue-cycle
    // This cannot produce any value:
    final Source<Integer, CompletionStage<Integer>> cyclicSource = Source.fromGraph(
            GraphDSL.create(foldSink,
                    (b, fold) -> {
                      // - Fold cannot complete until its upstream mapAsync completes
                      // - mapAsync cannot complete until the materialized Future produced by
                      //   fold completes
                      // As a result this Source will never emit anything, and its materialited
                      // Future will never complete
                      b.from(b.materializedValue()).via(b.add(flatten)).to(fold);
                      return SourceShape.of(b.from(b.materializedValue()).via(b.add(flatten)).out());
                    }));

    //#graph-dsl-matvalue-cycle
  }

  @Test
  public void beAbleToConstructClosedGraphFromList() throws Exception {
    //#graph-from-list
    //create the source
    final Source<String, NotUsed> in = Source.from(Arrays.asList("ax", "bx", "cx"));
    //generate the sinks from code
    List<String> prefixes = Arrays.asList("a", "b", "c");
    final List<Sink<String, CompletionStage<String>>> list = new ArrayList<>();
    for (String prefix : prefixes) {
      final Sink<String, CompletionStage<String>> sink =
              Flow.of(String.class).filter(str -> str.startsWith(prefix)).toMat(Sink.head(), Keep.right());
      list.add(sink);
    }

    final RunnableGraph<List<CompletionStage<String>>> g = RunnableGraph.fromGraph(
            GraphDSL.create(
                    list,
                    (GraphDSL.Builder<List<CompletionStage<String>>> builder, List<SinkShape<String>> outs) -> {
                      final UniformFanOutShape<String, String> bcast = builder.add(Broadcast.create(outs.size()));

                      final Outlet<String> source = builder.add(in).out();
                      builder.from(source).viaFanOut(bcast);

                      for (SinkShape<String> sink : outs) {
                        builder.from(bcast).to(sink);
                      }

                      return ClosedShape.getInstance();
                    }));
    List<CompletionStage<String>> result = g.run(materializer);
    //#graph-from-list

    assertEquals(3, result.size());
    assertEquals("ax", result.get(0).toCompletableFuture().get(1, TimeUnit.SECONDS));
    assertEquals("bx", result.get(1).toCompletableFuture().get(1, TimeUnit.SECONDS));
    assertEquals("cx", result.get(2).toCompletableFuture().get(1, TimeUnit.SECONDS));


  }

}
