/*
 * 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.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertTrue;
import static org.junit.jupiter.api.Assertions.fail;

import com.typesafe.config.Config;
import jakarta.jms.*;
import java.nio.charset.StandardCharsets;
import java.util.*;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.CompletionStage;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.TimeUnit;
import java.util.function.Function;
import java.util.stream.IntStream;
import java.util.stream.Stream;
import org.apache.activemq.artemis.jms.client.ActiveMQConnectionFactory;
import org.apache.activemq.artemis.jms.client.ActiveMQQueue;
import org.apache.activemq.artemis.jms.client.ActiveMQSession;
import org.apache.activemq.artemis.jms.client.ActiveMQTextMessage;
import org.apache.activemq.artemis.junit.EmbeddedActiveMQResource;
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.connectors.jakartams.*;
import org.apache.pekko.stream.connectors.jakartams.Destination;
import org.apache.pekko.stream.connectors.jakartams.javadsl.JmsConsumer;
import org.apache.pekko.stream.connectors.jakartams.javadsl.JmsConsumerControl;
import org.apache.pekko.stream.connectors.jakartams.javadsl.JmsProducer;
import org.apache.pekko.stream.connectors.jakartams.javadsl.JmsProducerStatus;
import org.apache.pekko.stream.connectors.testkit.javadsl.LogCapturingExtension;
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.junit.jupiter.api.AfterAll;
import org.junit.jupiter.api.BeforeAll;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.extension.ExtendWith;
import scala.util.Failure;
import scala.util.Success;
import scala.util.Try;

final class DummyJavaTests implements java.io.Serializable {

  private static final long serialVersionUID = 1234567L;

  private final String value;

  DummyJavaTests(String value) {
    this.value = value;
  }

  @Override
  public boolean equals(Object o) {
    if (this == o) {
      return true;
    }

    if (o instanceof DummyJavaTests dummy) {
      return dummy.value.equals(this.value);
    }
    return false;
  }

  @Override
  public int hashCode() {
    return value != null ? value.hashCode() : 0;
  }
}

@ExtendWith(LogCapturingExtension.class)
public class JmsConnectorsTest {

  private static ActorSystem system;
  private static Config producerConfig;
  private static Config browseConfig;

  @BeforeAll
  public static void setup() {
    system = ActorSystem.create();
    producerConfig = system.settings().config().getConfig(JmsProducerSettings.configPath());
    browseConfig = system.settings().config().getConfig(JmsBrowseSettings.configPath());
  }

  @AfterAll
  public static void teardown() {
    TestKit.shutdownActorSystem(system);
  }

  private List<JmsTextMessage> createTestMessageList() {
    List<Integer> intsIn = List.of(1, 2, 3, 4, 5, 6, 7, 8, 9, 10);
    List<JmsTextMessage> msgsIn = new ArrayList<>();
    for (Integer n : intsIn) {

      // #create-messages-with-properties
      JmsTextMessage message =
          org.apache.pekko.stream.connectors.jakartams.JmsTextMessage.create(n.toString())
              .withProperty("Number", n)
              .withProperty("IsOdd", n % 2 == 1)
              .withProperty("IsEven", n % 2 == 0);
      // #create-messages-with-properties

      msgsIn.add(message);
    }

    return msgsIn;
  }

  @Test
  public void publishAndConsumeJmsTextMessage() throws Exception {
    withServer(
        server -> {
          // #connection-factory
          // #text-sink
          // #text-source
          ConnectionFactory connectionFactory = new ActiveMQConnectionFactory(server.getVmURL());
          // #text-source
          // #connection-factory
          // #text-sink

          // #text-sink

          Sink<String, CompletionStage<Done>> jmsSink =
              JmsProducer.textSink(
                  JmsProducerSettings.create(system, connectionFactory).withQueue("test"));

          List<String> in = List.of("a", "b", "c", "d", "e", "f", "g", "h", "i", "j", "k");
          CompletionStage<Done> finished = Source.from(in).runWith(jmsSink, system);
          // #text-sink

          // #text-source
          Source<String, JmsConsumerControl> jmsSource =
              JmsConsumer.textSource(
                  JmsConsumerSettings.create(system, connectionFactory).withQueue("test"));

          CompletionStage<List<String>> result =
              jmsSource.take(in.size()).runWith(Sink.seq(), system);
          // #text-source

          assertEquals(Done.getInstance(), finished.toCompletableFuture().get(3, TimeUnit.SECONDS));
          assertEquals(in, result.toCompletableFuture().get(3, TimeUnit.SECONDS));
        });
  }

  @Test
  public void publishAndConsumeJmsObjectMessage() throws Exception {
    withServer(
        server -> {
          // #connection-factory-object
          // #object-sink
          // #object-source
          ActiveMQConnectionFactory connectionFactory =
              new ActiveMQConnectionFactory(server.getVmURL());
          // connectionFactory.(List.of(DummyJavaTests.class.getPackage().getName()));

          // #object-source
          // #connection-factory-object
          // #object-sink

          // #object-sink
          Sink<java.io.Serializable, CompletionStage<Done>> jmsSink =
              JmsProducer.objectSink(
                  JmsProducerSettings.create(system, connectionFactory).withQueue("test"));

          java.io.Serializable in = new DummyJavaTests("javaTest");
          CompletionStage<Done> finished = Source.single(in).runWith(jmsSink, system);
          // #object-sink

          // #object-source
          Source<java.io.Serializable, JmsConsumerControl> jmsSource =
              JmsConsumer.objectSource(
                  JmsConsumerSettings.create(system, connectionFactory).withQueue("test"));

          CompletionStage<java.io.Serializable> result =
              jmsSource.take(1).runWith(Sink.head(), system);
          // #object-source

          assertEquals(Done.getInstance(), finished.toCompletableFuture().get(3, TimeUnit.SECONDS));
          Object resultObject = result.toCompletableFuture().get(3, TimeUnit.SECONDS);
          assertEquals(resultObject, in);
        });
  }

  @Test
  public void publishAndConsumeJmsByteMessage() throws Exception {
    withConnectionFactory(
        connectionFactory -> {
          // #bytearray-sink
          Sink<byte[], CompletionStage<Done>> jmsSink =
              JmsProducer.bytesSink(
                  JmsProducerSettings.create(producerConfig, connectionFactory).withQueue("test"));

          byte[] in = "ThisIsATest".getBytes(StandardCharsets.UTF_8);
          CompletionStage<Done> finished = Source.single(in).runWith(jmsSink, system);
          // #bytearray-sink

          // #bytearray-source
          Source<byte[], JmsConsumerControl> jmsSource =
              JmsConsumer.bytesSource(
                  JmsConsumerSettings.create(system, connectionFactory).withQueue("test"));

          CompletionStage<byte[]> result = jmsSource.take(1).runWith(Sink.head(), system);
          // #bytearray-source

          assertEquals(Done.getInstance(), finished.toCompletableFuture().get(3, TimeUnit.SECONDS));
          byte[] resultArray = result.toCompletableFuture().get(3, TimeUnit.SECONDS);
          assertEquals("ThisIsATest", new String(resultArray, StandardCharsets.UTF_8));
        });
  }

  @Test
  public void publishAndConsumeJmsMapMessage() throws Exception {
    withConnectionFactory(
        connectionFactory -> {
          // #map-sink
          // #map-source

          // #map-sink
          // #map-source
          // #map-sink
          Sink<Map<String, Object>, CompletionStage<Done>> jmsSink =
              JmsProducer.mapSink(
                  JmsProducerSettings.create(system, connectionFactory).withQueue("test"));

          Map<String, Object> in = new HashMap<>();
          in.put("string value", "value");
          in.put("int value", 42);
          in.put("double value", 43.0);
          in.put("short value", (short) 7);
          in.put("boolean value", true);
          in.put("long value", 7L);
          in.put("bytearray", "AStringAsByteArray".getBytes(StandardCharsets.UTF_8));
          in.put("byte", (byte) 1);

          CompletionStage<Done> finished = Source.single(in).runWith(jmsSink, system);
          // #map-sink

          // #map-source
          Source<Map<String, Object>, JmsConsumerControl> jmsSource =
              JmsConsumer.mapSource(
                  JmsConsumerSettings.create(system, connectionFactory).withQueue("test"));

          CompletionStage<Map<String, Object>> resultStage =
              jmsSource.take(1).runWith(Sink.head(), system);
          // #map-source

          Map<String, Object> resultMap =
              resultStage.toCompletableFuture().get(3, TimeUnit.SECONDS);

          assertEquals(resultMap.get("string value"), in.get("string value"));
          assertEquals(resultMap.get("int value"), in.get("int value"));
          assertEquals(resultMap.get("double value"), in.get("double value"));
          assertEquals(resultMap.get("short value"), in.get("short value"));
          assertEquals(resultMap.get("boolean value"), in.get("boolean value"));
          assertEquals(resultMap.get("long value"), in.get("long value"));
          assertEquals(resultMap.get("byte"), in.get("byte"));

          assertEquals(Done.getInstance(), finished.toCompletableFuture().get(3, TimeUnit.SECONDS));
          byte[] resultByteArray = (byte[]) resultMap.get("bytearray");
          assertEquals(new String(resultByteArray, StandardCharsets.UTF_8), "AStringAsByteArray");
        });
  }

  @Test
  public void publishAndConsumeJmsTextMessagesWithProperties() throws Exception {
    withConnectionFactory(
        connectionFactory -> {
          int expectedMessages = 2;

          // #create-jms-sink
          Sink<JmsTextMessage, CompletionStage<Done>> jmsSink =
              JmsProducer.sink(
                  JmsProducerSettings.create(producerConfig, connectionFactory).withQueue("test"));

          CompletionStage<Done> finished =
              Source.from(List.of("Message A", "Message B"))
                  .map(JmsTextMessage::create)
                  .runWith(jmsSink, system);
          // #create-jms-sink

          // #jms-source
          Source<jakarta.jms.Message, JmsConsumerControl> jmsSource =
              JmsConsumer.create(
                  JmsConsumerSettings.create(system, connectionFactory).withQueue("test"));

          Pair<JmsConsumerControl, CompletionStage<List<String>>> controlAndResult =
              jmsSource
                  .take(expectedMessages)
                  .map(
                      msg -> {
                        if (msg instanceof TextMessage t) {
                          return t.getText();
                        } else
                          throw new RuntimeException("unexpected message type " + msg.getClass());
                      })
                  .toMat(Sink.seq(), Keep.both())
                  .run(system);

          // #jms-source

          CompletionStage<List<String>> result = controlAndResult.second();
          List<String> outMessages = result.toCompletableFuture().get(3, TimeUnit.SECONDS);
          assertEquals(expectedMessages, outMessages.size(), "unexpected number of elements");
          // #jms-source
          JmsConsumerControl control = controlAndResult.first();
          control.shutdown();
          // #jms-source
        });
  }

  @Test
  public void publishAndConsumeJmsTextMessagesWithHeaders() throws Exception {
    withConnectionFactory(
        connectionFactory -> {
          Sink<JmsTextMessage, CompletionStage<Done>> jmsSink =
              JmsProducer.sink(
                  JmsProducerSettings.create(producerConfig, connectionFactory).withQueue("test"));

          // #create-messages-with-headers
          List<JmsTextMessage> msgsIn =
              createTestMessageList().stream()
                  .map(
                      jmsTextMessage ->
                          jmsTextMessage
                              .withHeader(JmsType.create("type"))
                              .withHeader(JmsCorrelationId.create("correlationId"))
                              .withHeader(JmsReplyTo.queue("test-reply"))
                              .withHeader(JmsTimeToLive.create(999, TimeUnit.SECONDS))
                              .withHeader(JmsPriority.create(2))
                              .withHeader(JmsDeliveryMode.create(DeliveryMode.NON_PERSISTENT)))
                  .toList();
          // #create-messages-with-headers

          Source.from(msgsIn).runWith(jmsSink, system);

          Source<Message, JmsConsumerControl> jmsSource =
              JmsConsumer.create(
                  JmsConsumerSettings.create(system, connectionFactory).withQueue("test"));

          CompletionStage<List<Message>> result =
              jmsSource.take(msgsIn.size()).runWith(Sink.seq(), system);

          List<Message> outMessages = result.toCompletableFuture().get(3, TimeUnit.SECONDS);
          int msgIdx = 0;
          for (Message outMsg : outMessages) {
            assertEquals(
                outMsg.getIntProperty("Number"),
                msgsIn.get(msgIdx).properties().get("Number").get());
            assertEquals(
                outMsg.getBooleanProperty("IsOdd"),
                msgsIn.get(msgIdx).properties().get("IsOdd").get());
            assertEquals(
                outMsg.getBooleanProperty("IsEven"),
                (msgsIn.get(msgIdx).properties().get("IsEven").get()));
            assertEquals(outMsg.getJMSType(), "type");
            assertEquals(outMsg.getJMSCorrelationID(), "correlationId");
            assertEquals(((ActiveMQQueue) outMsg.getJMSReplyTo()).getQueueName(), "test-reply");

            assertTrue(outMsg.getJMSExpiration() != 0);
            assertEquals(2, outMsg.getJMSPriority());
            assertEquals(DeliveryMode.NON_PERSISTENT, outMsg.getJMSDeliveryMode());
            msgIdx++;
          }
        });
  }

  // #custom-destination
  Function<jakarta.jms.Session, jakarta.jms.Destination> createQueue(String destinationName) {
    return (session) -> {
      ActiveMQSession amqSession = (ActiveMQSession) session;
      try {
        return amqSession.createQueue("my-" + destinationName);
      } catch (JMSException e) {
        throw new RuntimeException(e);
      }
    };
  }

  // #custom-destination

  @Test
  public void useCustomDestination() throws Exception {
    withConnectionFactory(
        connectionFactory -> {
          Sink<JmsTextMessage, CompletionStage<Done>> jmsSink =
              JmsProducer.sink(
                  JmsProducerSettings.create(producerConfig, connectionFactory)
                      .withDestination(new CustomDestination("custom", createQueue("custom"))));

          List<JmsTextMessage> msgsIn = createTestMessageList();

          Source.from(msgsIn).runWith(jmsSink, system);
          // #custom-destination

          Source<Message, JmsConsumerControl> jmsSource =
              JmsConsumer.create(
                  JmsConsumerSettings.create(system, connectionFactory)
                      .withDestination(new CustomDestination("custom", createQueue("custom"))));
          // #custom-destination

          CompletionStage<List<Message>> result =
              jmsSource.take(msgsIn.size()).runWith(Sink.seq(), system);
          List<Message> outMessages = result.toCompletableFuture().get(3, TimeUnit.SECONDS);
          assertEquals(10, outMessages.size());
        });
  }

  @Test
  public void publishJmsTextMessagesWithPropertiesAndConsumeThemWithASelector() throws Exception {
    withConnectionFactory(
        connectionFactory -> {
          Sink<JmsTextMessage, CompletionStage<Done>> jmsSink =
              JmsProducer.sink(
                  JmsProducerSettings.create(producerConfig, connectionFactory).withQueue("test"));

          List<JmsTextMessage> msgsIn = createTestMessageList();

          Source.from(msgsIn).runWith(jmsSink, system);

          // #source-with-selector
          Source<Message, JmsConsumerControl> jmsSource =
              JmsConsumer.create(
                  JmsConsumerSettings.create(system, connectionFactory)
                      .withQueue("test")
                      .withSelector("IsOdd = TRUE"));
          // #source-with-selector

          List<JmsTextMessage> oddMsgsIn =
              msgsIn.stream().filter(msg -> Integer.valueOf(msg.body()) % 2 == 1).toList();
          assertEquals(5, oddMsgsIn.size());

          CompletionStage<List<Message>> result =
              jmsSource.take(oddMsgsIn.size()).runWith(Sink.seq(), system);

          List<Message> outMessages = result.toCompletableFuture().get(4, TimeUnit.SECONDS);
          int msgIdx = 0;
          for (Message outMsg : outMessages) {
            assertEquals(
                outMsg.getIntProperty("Number"),
                oddMsgsIn.get(msgIdx).properties().get("Number").get());
            assertEquals(
                outMsg.getBooleanProperty("IsOdd"),
                oddMsgsIn.get(msgIdx).properties().get("IsOdd").get());
            assertEquals(
                outMsg.getBooleanProperty("IsEven"),
                (oddMsgsIn.get(msgIdx).properties().get("IsEven").get()));
            assertEquals(1, outMsg.getIntProperty("Number") % 2);
            msgIdx++;
          }
        });
  }

  @Test
  public void publishAndConsumeTopic() throws Exception {
    withConnectionFactory(
        connectionFactory -> {
          List<String> in = List.of("a", "b", "c", "d", "e", "f", "g", "h", "i", "j", "k");
          List<String> inNumbers = IntStream.range(0, 10).boxed().map(String::valueOf).toList();

          Sink<String, CompletionStage<Done>> jmsTopicSink =
              JmsProducer.textSink(
                  JmsProducerSettings.create(producerConfig, connectionFactory).withTopic("topic"));

          Sink<String, CompletionStage<Done>> jmsTopicSink2 =
              JmsProducer.textSink(
                  JmsProducerSettings.create(producerConfig, connectionFactory).withTopic("topic"));

          Source<String, JmsConsumerControl> jmsTopicSource =
              JmsConsumer.textSource(
                  JmsConsumerSettings.create(system, connectionFactory).withTopic("topic"));

          Source<String, JmsConsumerControl> jmsTopicSource2 =
              JmsConsumer.textSource(
                  JmsConsumerSettings.create(system, connectionFactory).withTopic("topic"));

          CompletionStage<List<String>> result =
              jmsTopicSource
                  .take(in.size() + inNumbers.size())
                  .runWith(Sink.seq(), system)
                  .thenApply(l -> l.stream().sorted().toList());

          CompletionStage<List<String>> result2 =
              jmsTopicSource2
                  .take(in.size() + inNumbers.size())
                  .runWith(Sink.seq(), system)
                  .thenApply(l -> l.stream().sorted().toList());

          Thread.sleep(500);

          CompletionStage<Done> finished = Source.from(in).runWith(jmsTopicSink, system);
          Source.from(inNumbers).runWith(jmsTopicSink2, system);

          assertEquals(
              Stream.concat(in.stream(), inNumbers.stream()).sorted().toList(),
              result.toCompletableFuture().get(5, TimeUnit.SECONDS));
          assertEquals(
              Stream.concat(in.stream(), inNumbers.stream()).sorted().toList(),
              result2.toCompletableFuture().get(5, TimeUnit.SECONDS));
        });
  }

  @Test
  public void sinkNormalCompletion() throws Exception {
    withConnectionFactory(
        connectionFactory -> {
          Sink<JmsTextMessage, CompletionStage<Done>> jmsSink =
              JmsProducer.sink(
                  JmsProducerSettings.create(producerConfig, connectionFactory)
                      .withQueue("test")
                      .withConnectionRetrySettings(
                          ConnectionRetrySettings.create(system).withMaxRetries(0)));

          List<JmsTextMessage> msgsIn = createTestMessageList();

          CompletionStage<Done> completionFuture = Source.from(msgsIn).runWith(jmsSink, system);
          Done completed = completionFuture.toCompletableFuture().get(3, TimeUnit.SECONDS);
          assertEquals(completed, Done.getInstance());
        });
  }

  @Test
  public void sinkExceptionalCompletion() throws Exception {
    withConnectionFactory(
        connectionFactory -> {
          Sink<JmsTextMessage, CompletionStage<Done>> jmsSink =
              JmsProducer.sink(
                  JmsProducerSettings.create(producerConfig, connectionFactory).withQueue("test"));

          CompletionStage<Done> completionFuture =
              Source.<JmsTextMessage>failed(new RuntimeException("Simulated error"))
                  .runWith(jmsSink, system);

          try {
            completionFuture.toCompletableFuture().get(3, TimeUnit.SECONDS);
            fail("Not completed exceptionally");
          } catch (ExecutionException e) {
            Throwable cause = e.getCause();
            assertTrue(cause instanceof RuntimeException);
          }
        });
  }

  @Test
  public void sinkExceptionalCompletionOnDisconnect() throws Exception {
    withServer(
        server -> {
          ConnectionFactory connectionFactory = new ActiveMQConnectionFactory(server.getVmURL());

          Sink<JmsTextMessage, CompletionStage<Done>> jmsSink =
              JmsProducer.sink(
                  JmsProducerSettings.create(producerConfig, connectionFactory)
                      .withQueue("test")
                      .withConnectionRetrySettings(
                          ConnectionRetrySettings.create(system).withMaxRetries(0)));

          List<JmsTextMessage> msgsIn = createTestMessageList();

          CompletionStage<Done> completionFuture =
              Source.from(msgsIn)
                  .mapAsync(
                      1,
                      m ->
                          CompletableFuture.supplyAsync(
                              () -> {
                                try {
                                  Thread.sleep(500);
                                } catch (InterruptedException e) {
                                  throw new RuntimeException(e);
                                }
                                return m;
                              }))
                  .runWith(jmsSink, system);

          try { // Make sure connection got started before stopping.
            Thread.sleep(500);
          } catch (InterruptedException e) {
            fail("Sleep interrupted.");
          }

          server.stop();

          try {
            completionFuture.toCompletableFuture().get(3, TimeUnit.SECONDS);
            fail("Not completed exceptionally");
          } catch (ExecutionException e) {
            Throwable cause = e.getCause();
            assertTrue(cause instanceof JMSException);
          }
        });
  }

  @Test
  public void browse() throws Exception {
    withConnectionFactory(
        connectionFactory -> {
          Sink<String, CompletionStage<Done>> jmsSink =
              JmsProducer.textSink(
                  JmsProducerSettings.create(producerConfig, connectionFactory).withQueue("test"));

          List<String> in = List.of("a", "b", "c", "d", "e", "f", "g", "h", "i", "j", "k");

          Source.from(in).runWith(jmsSink, system).toCompletableFuture().get();

          // #browse-source
          Source<jakarta.jms.Message, NotUsed> browseSource =
              JmsConsumer.browse(
                  JmsBrowseSettings.create(system, connectionFactory).withQueue("test"));

          CompletionStage<List<jakarta.jms.Message>> result =
              browseSource.runWith(Sink.seq(), system);
          // #browse-source

          List<String> resultText =
              result.toCompletableFuture().get().stream()
                  .map(message -> ((ActiveMQTextMessage) message).getText())
                  .toList();

          assertEquals(in, resultText);
        });
  }

  @Test
  public void publishAndConsumeDurableTopic() throws Exception {
    withServer(
        server -> {
          ConnectionFactory producerConnectionFactory =
              new ActiveMQConnectionFactory(server.getVmURL());
          // #create-connection-factory-with-client-id
          ConnectionFactory consumerConnectionFactory =
              new ActiveMQConnectionFactory(server.getVmURL());
          ((ActiveMQConnectionFactory) consumerConnectionFactory)
              .setClientID(getClass().getSimpleName());
          // #create-connection-factory-with-client-id

          List<String> in = List.of("a", "b", "c", "d", "e", "f", "g", "h", "i", "j", "k");

          Sink<String, CompletionStage<Done>> jmsTopicSink =
              JmsProducer.textSink(
                  JmsProducerSettings.create(producerConfig, producerConnectionFactory)
                      .withTopic("topic"));

          // #create-durable-topic-source
          Source<String, JmsConsumerControl> jmsTopicSource =
              JmsConsumer.textSource(
                  JmsConsumerSettings.create(system, consumerConnectionFactory)
                      .withDurableTopic("topic", "durable-test"));
          // #create-durable-topic-source

          // #run-durable-topic-source
          CompletionStage<List<String>> result =
              jmsTopicSource.take(in.size()).runWith(Sink.seq(), system);
          // #run-durable-topic-source

          Thread.sleep(500);

          Source.from(in).runWith(jmsTopicSink, system);

          assertEquals(in, result.toCompletableFuture().get(5, TimeUnit.SECONDS));
        });
  }

  @Test
  public void producerFlow() throws Exception {
    withConnectionFactory(
        connectionFactory -> {
          // #flow-producer
          Flow<JmsTextMessage, JmsTextMessage, JmsProducerStatus> flow =
              JmsProducer.flow(
                  JmsProducerSettings.create(system, connectionFactory).withQueue("test"));

          List<JmsTextMessage> input = createTestMessageList();

          CompletionStage<List<JmsTextMessage>> result =
              Source.from(input).via(flow).runWith(Sink.seq(), system);
          // #flow-producer

          assertEquals(input, result.toCompletableFuture().get());
        });
  }

  @Test
  public void directedProducerFlow() throws Exception {
    withConnectionFactory(
        connectionFactory -> {
          // #run-directed-flow-producer
          Flow<JmsTextMessage, JmsTextMessage, JmsProducerStatus> flowSink =
              JmsProducer.flow(
                  JmsProducerSettings.create(system, connectionFactory).withQueue("test"));

          List<JmsTextMessage> input = new ArrayList<>();
          for (Integer n : List.of(1, 2, 3, 4, 5, 6, 7, 8, 9, 10)) {
            String queueName = (n % 2 == 0) ? "even" : "odd";
            input.add(JmsTextMessage.create(n.toString()).toQueue(queueName));
          }

          Source.from(input).via(flowSink).runWith(Sink.seq(), system);
          // #run-directed-flow-producer

          CompletionStage<List<Integer>> even =
              JmsConsumer.textSource(
                      JmsConsumerSettings.create(system, connectionFactory).withQueue("even"))
                  .take(5)
                  .map(Integer::parseInt)
                  .runWith(Sink.seq(), system);

          CompletionStage<List<Integer>> odd =
              JmsConsumer.textSource(
                      JmsConsumerSettings.create(system, connectionFactory).withQueue("odd"))
                  .take(5)
                  .map(Integer::parseInt)
                  .runWith(Sink.seq(), system);

          assertEquals(List.of(1, 3, 5, 7, 9), odd.toCompletableFuture().get());
          assertEquals(List.of(2, 4, 6, 8, 10), even.toCompletableFuture().get());
        });
  }

  @Test
  public void failAfterRetry() throws Exception {
    withServer(
        server -> {
          ConnectionFactory connectionFactory = new ActiveMQConnectionFactory(server.getVmURL());
          server.stop();
          long startTime = System.currentTimeMillis();
          CompletionStage<List<Message>> result =
              JmsConsumer.create(
                      JmsConsumerSettings.create(system, connectionFactory)
                          .withConnectionRetrySettings(
                              ConnectionRetrySettings.create(system).withMaxRetries(4))
                          .withQueue("test"))
                  .runWith(Sink.seq(), system);

          CompletionStage<Try<List<Message>>> tryFuture =
              result.handle(
                  (l, e) -> {
                    if (l != null) return Success.apply(l);
                    else return Failure.apply(e);
                  });

          Try<List<Message>> tryResult = tryFuture.toCompletableFuture().get();
          long endTime = System.currentTimeMillis();

          assertTrue(endTime - startTime > 100L + 400L + 900L + 1600L, "Total retry is too short");
          assertTrue(tryResult.isFailure(), "Result must be a failure");
          Throwable exception = tryResult.failed().get();
          assertTrue(
              exception instanceof ConnectionRetryException,
              "Did not fail with a ConnectionRetryException");
          assertTrue(
              exception.getCause() instanceof JMSException,
              "Cause of failure is not a JMSException");
        });
  }

  @Test
  public void passThroughMessageEnvelopes() throws Exception {
    withConnectionFactory(
        connectionFactory -> {
          // #run-flexi-flow-producer
          Flow<JmsEnvelope<String>, JmsEnvelope<String>, JmsProducerStatus> jmsProducer =
              JmsProducer.flexiFlow(
                  JmsProducerSettings.create(producerConfig, connectionFactory).withQueue("test"));

          List<String> data = List.of("a", "b", "c", "d", "e", "f", "g", "h", "i", "j", "k");
          List<JmsEnvelope<String>> input = new ArrayList<>();
          for (String s : data) {
            String passThrough = s;
            input.add(JmsTextMessage.create(s, passThrough));
          }

          CompletionStage<List<String>> result =
              Source.from(input)
                  .via(jmsProducer)
                  .map(JmsEnvelope::passThrough)
                  .runWith(Sink.seq(), system);
          // #run-flexi-flow-producer
          assertEquals(data, result.toCompletableFuture().get());
        });
  }

  @Test
  public void passThroughEmptyMessageEnvelopes() throws Exception {
    withConnectionFactory(
        connectionFactory -> {
          Pair<JmsConsumerControl, CompletionStage<List<String>>> switchAndItems =
              JmsConsumer.textSource(
                      JmsConsumerSettings.create(system, connectionFactory).withQueue("test"))
                  .toMat(Sink.seq(), Keep.both())
                  .run(system);

          // #run-flexi-flow-pass-through-producer
          Flow<JmsEnvelope<String>, JmsEnvelope<String>, JmsProducerStatus> jmsProducer =
              JmsProducer.flexiFlow(
                  JmsProducerSettings.create(producerConfig, connectionFactory).withQueue("test"));

          List<String> data = List.of("a", "b", "c", "d", "e", "f", "g", "h", "i", "j", "k");
          List<JmsEnvelope<String>> input = new ArrayList<>();
          for (String s : data) {
            String passThrough = s;
            input.add(JmsPassThrough.create(passThrough));
          }

          CompletionStage<List<String>> result =
              Source.from(input)
                  .via(jmsProducer)
                  .map(JmsEnvelope::passThrough)
                  .runWith(Sink.seq(), system);
          // #run-flexi-flow-pass-through-producer

          assertEquals(data, result.toCompletableFuture().get());

          Thread.sleep(500);

          switchAndItems.first().shutdown();
          assertTrue(switchAndItems.second().toCompletableFuture().get().isEmpty());
        });
  }

  @Test
  public void requestReplyWithTempQueues() throws Exception {
    withConnectionFactory(
        connectionFactory -> {
          Connection connection = connectionFactory.createConnection();
          connection.start();
          Session session = connection.createSession(false, Session.AUTO_ACKNOWLEDGE);
          TemporaryQueue tempQueue = session.createTemporaryQueue();
          Destination tempQueueDest =
              org.apache.pekko.stream.connectors.jakartams.Destination.createDestination(tempQueue);
          String message = "ThisIsATest";
          Function<String, String> reverse = (in) -> new StringBuilder(in).reverse().toString();
          String correlationId = UUID.randomUUID().toString();

          Sink<JmsTextMessage, CompletionStage<Done>> toRespondSink =
              JmsProducer.sink(
                  JmsProducerSettings.create(system, connectionFactory).withQueue("test"));

          CompletionStage<Done> toRespondStreamCompletion =
              Source.single(
                      JmsTextMessage.create(message)
                          .withHeader(JmsCorrelationId.create(correlationId))
                          .withHeader(new JmsReplyTo(tempQueueDest)))
                  .runWith(toRespondSink, system);

          // #request-reply
          JmsConsumerControl respondStreamControl =
              JmsConsumer.create(
                      JmsConsumerSettings.create(system, connectionFactory).withQueue("test"))
                  .map(JmsMessageFactory::create)
                  .collectType(JmsTextMessage.class)
                  .map(
                      textMessage -> {
                        JmsTextMessage m = JmsTextMessage.create(reverse.apply(textMessage.body()));
                        for (JmsHeader h : textMessage.getHeaders())
                          if (h.getClass().equals(JmsReplyTo.class))
                            m = m.to(((JmsReplyTo) h).jmsDestination());
                          else if (h.getClass().equals(JmsCorrelationId.class)) m = m.withHeader(h);
                        return m;
                      })
                  .via(
                      JmsProducer.flow(
                          JmsProducerSettings.create(system, connectionFactory)
                              .withQueue("ignored")))
                  .to(Sink.ignore())
                  .run(system);
          // #request-reply

          // getting ConnectionRetryException when trying to listen using streams, assuming it's
          // because a different session can't listen on the original session's TemporaryQueue
          MessageConsumer consumer = session.createConsumer(tempQueue);
          Message msg = consumer.receive(5000);

          assertEquals(Done.done(), toRespondStreamCompletion.toCompletableFuture().get());

          assertTrue(msg instanceof TextMessage);
          assertEquals(correlationId, msg.getJMSCorrelationID());
          assertEquals(reverse.apply(message), ((TextMessage) msg).getText());

          respondStreamControl.shutdown();

          connection.close();
        });
  }

  private void withServer(ConsumerChecked<EmbeddedActiveMQResource> test) throws Exception {
    EmbeddedActiveMQResource server = new EmbeddedActiveMQResource();
    try {
      server.start();
      test.accept(server);
      Thread.sleep(500);
    } finally {
      if (server.getServer().getActiveMQServer().isStarted()) {
        server.stop();
      }
    }
  }

  private void withConnectionFactory(ConsumerChecked<ConnectionFactory> test) throws Exception {
    withServer(server -> test.accept(new ActiveMQConnectionFactory(server.getVmURL())));
  }

  @FunctionalInterface
  private interface ConsumerChecked<T> {
    void accept(T elt) throws Exception;
  }
}
