/*
 * 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.runners.samza.adapter;

import java.io.IOException;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.NoSuchElementException;
import java.util.stream.Collectors;
import javax.annotation.Nullable;
import org.apache.beam.runners.samza.adapter.TestSourceHelpers.ElementEvent;
import org.apache.beam.runners.samza.adapter.TestSourceHelpers.Event;
import org.apache.beam.runners.samza.adapter.TestSourceHelpers.ExceptionEvent;
import org.apache.beam.runners.samza.adapter.TestSourceHelpers.LatchEvent;
import org.apache.beam.runners.samza.adapter.TestSourceHelpers.NoElementEvent;
import org.apache.beam.runners.samza.adapter.TestSourceHelpers.SourceBuilder;
import org.apache.beam.runners.samza.adapter.TestSourceHelpers.WatermarkEvent;
import org.apache.beam.sdk.coders.Coder;
import org.apache.beam.sdk.coders.SerializableCoder;
import org.apache.beam.sdk.io.UnboundedSource;
import org.apache.beam.sdk.options.PipelineOptions;
import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
import org.joda.time.Instant;

/**
 * A unbounded source that can be used for test purposes.
 *
 * @param <T> element type
 */
public class TestUnboundedSource<T> extends UnboundedSource<T, TestCheckpointMark> {
  // each list of events is a split
  private final List<List<Event<T>>> events;

  public static <T> Builder<T> createBuilder() {
    return new Builder<>();
  }

  public static <T> SplittableBuilder<T> createSplits(int numSplits) {
    return new SplittableBuilder<>(numSplits);
  }

  private TestUnboundedSource(List<List<Event<T>>> events) {
    this.events = Collections.unmodifiableList(new ArrayList<>(events));
  }

  @Override
  public List<? extends UnboundedSource<T, TestCheckpointMark>> split(
      int desiredNumSplits, PipelineOptions options) throws Exception {
    return events
        .stream()
        .map(ev -> new TestUnboundedSource<>(Collections.singletonList(ev)))
        .collect(Collectors.toList());
  }

  @Override
  public UnboundedReader<T> createReader(
      PipelineOptions options, @Nullable TestCheckpointMark checkpointMark) throws IOException {
    assert events.size() == 1;

    return new Reader(events.get(0), checkpointMark);
  }

  @Override
  public Coder<TestCheckpointMark> getCheckpointMarkCoder() {
    return SerializableCoder.of(TestCheckpointMark.class);
  }

  @Override
  public void validate() {}

  /** A builder used to populate the events emitted by {@link TestUnboundedSource}. */
  public static class Builder<T> extends SourceBuilder<T, TestUnboundedSource<T>> {

    @Override
    public TestUnboundedSource<T> build() {
      return new TestUnboundedSource<>(Collections.singletonList(getEvents()));
    }
  }

  /**
   * A SplittableBuilder supports multiple splits and each split {@link TestUnboundedSource} can be
   * built separately from the above Builder.
   */
  public static class SplittableBuilder<T> extends SourceBuilder<T, TestUnboundedSource<T>> {
    private final List<Builder<T>> builders = new ArrayList<>();

    private SplittableBuilder(int splits) {
      while (splits != 0) {
        builders.add(new Builder<T>());
        --splits;
      }
    }

    @Override
    public TestUnboundedSource<T> build() {
      List<List<Event<T>>> events = new ArrayList<>();
      builders.forEach(builder -> events.add(builder.getEvents()));
      return new TestUnboundedSource<>(events);
    }

    public Builder<T> forSplit(int split) {
      return builders.get(split);
    }
  }

  private class Reader extends UnboundedReader<T> {
    private final List<Event<T>> events;
    private Instant curTime = BoundedWindow.TIMESTAMP_MIN_VALUE;
    private Instant watermark = BoundedWindow.TIMESTAMP_MIN_VALUE;
    private boolean started;
    private int index = -1;
    private int offset;

    private Reader(List<Event<T>> events, TestCheckpointMark checkpointMark) {
      this.events = events;
      this.offset = checkpointMark == null ? -1 : checkpointMark.checkpoint;
    }

    @Override
    public boolean start() throws IOException {
      if (started) {
        throw new IllegalStateException("Start called when reader was already started");
      }
      started = true;
      return advance();
    }

    @Override
    public boolean advance() throws IOException {
      if (!started) {
        throw new IllegalStateException("Advance called when reader was not started");
      }

      for (++index; index < events.size(); ++index) {
        final Event<T> event = events.get(index);
        if (event instanceof ExceptionEvent) {
          throw ((ExceptionEvent<T>) event).exception;
        } else if (event instanceof LatchEvent) {
          try {
            ((LatchEvent) event).latch.await();
          } catch (InterruptedException e) {
            // Propagate interrupt
            Thread.currentThread().interrupt();
          }
        } else if (event instanceof WatermarkEvent) {
          watermark = ((WatermarkEvent) event).watermark;
        } else if (event instanceof NoElementEvent) {
          return false;
        } else {
          curTime = ((ElementEvent<T>) event).timestamp;
          ++offset;
          return true;
        }
      }

      return false;
    }

    @Override
    public T getCurrent() throws NoSuchElementException {
      if (!started) {
        throw new NoSuchElementException();
      }

      final Event<T> event = events.get(index);
      assert event instanceof ElementEvent;
      return ((ElementEvent<T>) event).element;
    }

    @Override
    public Instant getCurrentTimestamp() throws NoSuchElementException {
      return curTime;
    }

    @Override
    public Instant getWatermark() {
      return watermark;
    }

    @Override
    public CheckpointMark getCheckpointMark() {
      return TestCheckpointMark.of(offset);
    }

    @Override
    public void close() throws IOException {}

    @Override
    public UnboundedSource<T, ?> getCurrentSource() {
      return TestUnboundedSource.this;
    }
  }
}
