/*
 * 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) 2019-2022 Lightbend Inc. <https://www.lightbend.com>
 */

package org.apache.pekko.stream.javadsl;

import org.apache.pekko.NotUsed;
import org.apache.pekko.japi.Pair;
import org.apache.pekko.stream.StreamTest;
import org.apache.pekko.stream.testkit.TestPublisher;
import org.apache.pekko.stream.testkit.TestSubscriber;
import org.apache.pekko.stream.testkit.javadsl.TestSink;
import org.apache.pekko.stream.testkit.javadsl.TestSource;
import org.apache.pekko.testkit.PekkoJUnitActorSystemResource;
import org.apache.pekko.testkit.PekkoSpec;
import org.junit.ClassRule;
import org.junit.Test;
import scala.util.Failure;
import scala.util.Success;
import scala.util.Try;

import java.time.Duration;
import java.util.Optional;

import static org.junit.Assert.assertEquals;

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

  @ClassRule
  public static PekkoJUnitActorSystemResource actorSystemResource =
      new PekkoJUnitActorSystemResource("RetryFlowTest", PekkoSpec.testConf());

  public static
  // #withBackoff-signature
  <In, Out, Mat> Flow<In, Out, Mat> withBackoff(
      Duration minBackoff,
      Duration maxBackoff,
      double randomFactor,
      int maxRetries,
      Flow<In, Out, Mat> flow,
      org.apache.pekko.japi.function.Function2<In, Out, Optional<In>> decideRetry)
        // #withBackoff-signature
      {
    return RetryFlow.<In, Out, Mat>withBackoff(
        minBackoff, maxBackoff, randomFactor, maxRetries, flow, decideRetry);
  }

  public static
  // #signature
  <In, InCtx, Out, OutCtx, Mat> FlowWithContext<In, InCtx, Out, OutCtx, Mat> withBackoffAndContext(
      Duration minBackoff,
      Duration maxBackoff,
      double randomFactor,
      int maxRetries,
      FlowWithContext<In, InCtx, Out, OutCtx, Mat> flow,
      org.apache.pekko.japi.function.Function2<
              Pair<In, InCtx>, Pair<Out, OutCtx>, Optional<Pair<In, InCtx>>>
          decideRetry)
        // #signature
      {
    return RetryFlow.<In, InCtx, Out, OutCtx, Mat>withBackoffAndContext(
        minBackoff, maxBackoff, randomFactor, maxRetries, flow, decideRetry);
  }

  @Test
  public void withBackoffShouldRetry() {
    final Duration minBackoff = Duration.ofMillis(10);
    final Duration maxBackoff = Duration.ofSeconds(5);
    final double randomFactor = 0d;
    final int maxRetries = 3;

    // #withBackoff-demo
    Flow<Integer, Integer, NotUsed> flow = // ...
        // the wrapped flow
        // #withBackoff-demo
        Flow.fromFunction(in -> in / 2);

    // #withBackoff-demo

    Flow<Integer, Integer, NotUsed> retryFlow =
        RetryFlow.withBackoff(
            minBackoff,
            maxBackoff,
            randomFactor,
            maxRetries,
            flow,
            (in, out) -> {
              if (out > 0) return Optional.of(out);
              else return Optional.empty();
            });
    // #withBackoff-demo

    final Pair<TestPublisher.Probe<Integer>, TestSubscriber.Probe<Integer>> probes =
        TestSource.<Integer>probe(system)
            .via(retryFlow)
            .toMat(TestSink.probe(system), Keep.both())
            .run(system);

    final TestPublisher.Probe<Integer> source = probes.first();
    final TestSubscriber.Probe<Integer> sink = probes.second();

    sink.request(4);

    source.sendNext(5);
    assertEquals(0, sink.expectNext().intValue());

    source.sendComplete();
    sink.expectComplete();
  }

  @Test
  public void withBackoffAndContextShouldRetry() {
    final Duration minBackoff = Duration.ofMillis(10);
    final Duration maxBackoff = Duration.ofSeconds(5);
    final double randomFactor = 0d;
    final int maxRetries = 3;

    class SomeContext {}

    // #retry-success
    FlowWithContext<Integer, SomeContext, Integer, SomeContext, NotUsed> flow = // ...
        // the wrapped flow
        // #retry-success
        FlowWithContext.fromPairs(
            Flow.fromFunction(
                in -> {
                  final Integer request = in.first();
                  return Pair.create(request / 2, in.second());
                }));

    // #retry-success

    FlowWithContext<Integer, SomeContext, Integer, SomeContext, NotUsed> retryFlow =
        RetryFlow.withBackoffAndContext(
            minBackoff,
            maxBackoff,
            randomFactor,
            maxRetries,
            flow,
            (in, out) -> {
              Integer value = out.first();
              SomeContext context = out.second();
              if (value > 0) {
                return Optional.of(Pair.create(value, context));
              } else {
                return Optional.empty();
              }
            });
    // #retry-success

    final Pair<TestPublisher.Probe<Integer>, TestSubscriber.Probe<Pair<Integer, SomeContext>>>
        probes =
            TestSource.<Integer>probe(system)
                .map(i -> Pair.create(i, new SomeContext()))
                .via(retryFlow)
                .toMat(TestSink.probe(system), Keep.both())
                .run(system);

    final TestPublisher.Probe<Integer> source = probes.first();
    final TestSubscriber.Probe<Pair<Integer, SomeContext>> sink = probes.second();

    sink.request(4);

    source.sendNext(5);
    assertEquals(0, sink.expectNext().first().intValue());

    source.sendComplete();
    sink.expectComplete();
  }

  @Test
  public void retryFailedResponses() {
    final Duration minBackoff = Duration.ofMillis(10);
    final Duration maxBackoff = Duration.ofSeconds(5);
    final double randomFactor = 0d;
    final int maxRetries = 3;
    final FlowWithContext<Integer, Integer, Try<Integer>, Integer, NotUsed> failEvenValuesFlow =
        FlowWithContext.fromPairs(
            Flow.fromFunction(
                in -> {
                  final Integer request = in.first();
                  if (request % 2 == 0)
                    return Pair.create(Failure.apply(new Error("Failed response")), in.second());
                  else return Pair.create(Success.apply(request), in.second());
                }));

    final FlowWithContext<Integer, Integer, Try<Integer>, Integer, NotUsed> retryFlow =
        RetryFlow.withBackoffAndContext(
            minBackoff,
            maxBackoff,
            randomFactor,
            maxRetries,
            failEvenValuesFlow,
            (in, out) -> {
              if (out.first().isFailure()) {
                return Optional.of(Pair.create(in.first() + 1, out.second()));
              } else {
                return Optional.empty();
              }
            });

    final Pair<TestPublisher.Probe<Integer>, TestSubscriber.Probe<Pair<Try<Integer>, Integer>>>
        probes =
            TestSource.<Integer>probe(system)
                .map(i -> Pair.create(i, i))
                .via(retryFlow)
                .toMat(TestSink.probe(system), Keep.both())
                .run(system);

    final TestPublisher.Probe<Integer> source = probes.first();
    final TestSubscriber.Probe<Pair<Try<Integer>, Integer>> sink = probes.second();

    sink.request(1);
    source.sendNext(8);

    Pair<Try<Integer>, Integer> response = sink.expectNext();
    assertEquals(9, response.first().get().intValue());
    assertEquals(8, response.second().intValue());

    source.sendComplete();
    sink.expectComplete();
  }

  @Test
  public void supportFlowWithContext() {
    final Duration minBackoff = Duration.ofMillis(10);
    final Duration maxBackoff = Duration.ofSeconds(5);
    final double randomFactor = 0d;
    final int maxRetries = 5;
    final FlowWithContext<Integer, Integer, Try<Integer>, Integer, NotUsed> flow =
        Flow.<Integer>create()
            .<Integer, Integer, Integer>asFlowWithContext((el, ctx) -> el, ctx -> ctx)
            .map(
                i -> {
                  if (i > 0) return Failure.apply(new RuntimeException("i is larger than 0"));
                  else return Success.apply(i);
                });

    final Pair<TestPublisher.Probe<Integer>, TestSubscriber.Probe<Pair<Try<Integer>, Integer>>>
        probes =
            TestSource.<Integer>probe(system)
                .asSourceWithContext(ctx -> ctx)
                .via(
                    RetryFlow.withBackoffAndContext(
                        minBackoff,
                        maxBackoff,
                        randomFactor,
                        maxRetries,
                        flow,
                        (in, out) -> {
                          if (out.first().isFailure()) {
                            if (out.second() > 0) {
                              return Optional.of(Pair.create(out.second() / 2, out.second() / 2));
                            }
                          }
                          return Optional.empty();
                        }))
                .toMat(TestSink.probe(system), Keep.both())
                .run(system);

    final TestPublisher.Probe<Integer> source = probes.first();
    final TestSubscriber.Probe<Pair<Try<Integer>, Integer>> sink = probes.second();

    sink.request(1);
    source.sendNext(8);

    Pair<Try<Integer>, Integer> response = sink.expectNext();
    assertEquals(0, response.first().get().intValue());
    assertEquals(0, response.second().intValue());

    source.sendComplete();
    sink.expectComplete();
  }
}
