/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package org.apache.beam.sdk.io.rabbitmq;

import com.rabbitmq.client.AMQP;
import com.rabbitmq.client.Channel;
import com.rabbitmq.client.DefaultConsumer;
import com.rabbitmq.client.Envelope;
import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.util.List;
import java.util.stream.Collectors;
import java.util.stream.IntStream;

public class RabbitMqTestUtils {
  private RabbitMqTestUtils() {
    throw new UnsupportedOperationException(
        "RabbitMqTestUtils is a non-instantiable utility class");
  }

  /**
   * Generates a byte array to be used as a rabbit message payload given a record number.
   *
   * @param recordNum Arbitrary message number, typically the index of a for loop, used to construct
   *     the message payload.
   * @return The byte array message payload for a rabbitmq message, generated by appending the
   *     record number to the string "Test " and converting to UTF-8 bytes.
   * @see #generateRecords(int) for use in a for loop
   * @see #recordToString(byte[]) for reversing this payload back into a String
   */
  public static byte[] generateRecord(int recordNum) {
    return ("Test " + recordNum).getBytes(StandardCharsets.UTF_8);
  }

  /**
   * Produces payloads for {@code numRecords} messages utilizing {@link #generateRecord(int)}.
   *
   * @param numRecords the number of messages to produce
   * @return a list of length {@code numRecords} of distinct message payloads
   */
  public static List<byte[]> generateRecords(int numRecords) {
    return IntStream.range(0, numRecords)
        .mapToObj(RabbitMqTestUtils::generateRecord)
        .collect(Collectors.toList());
  }

  /**
   * @param record a byte array used as a rabbit message payload
   * @return the String representation produced by treating the payload as bytes of UTF-8 characters
   */
  public static String recordToString(byte[] record) {
    return new String(record, StandardCharsets.UTF_8);
  }

  /**
   * A simple RabbitMQ {@code Consumer} that stores all received messages in the
   * constructor-supplied List.
   */
  static class TestConsumer extends DefaultConsumer {

    private final List<String> received;

    public TestConsumer(Channel channel, List<String> received) {
      super(channel);
      this.received = received;
    }

    @Override
    public void handleDelivery(
        String consumerTag, Envelope envelope, AMQP.BasicProperties properties, byte[] body)
        throws IOException {
      received.add(recordToString(body));
    }
  }
}
