/*
 * Copyright 2015 Data Artisans GmbH
 *
 * Licensed 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 com.dataartisans.flink.dataflow.translation.wrappers.streaming.io;

import com.dataartisans.flink.dataflow.FlinkPipelineRunner;
import com.google.cloud.dataflow.sdk.coders.Coder;
import com.google.cloud.dataflow.sdk.io.UnboundedSource;
import com.google.cloud.dataflow.sdk.options.PipelineOptions;
import com.google.common.base.Preconditions;
import org.apache.flink.streaming.api.functions.source.RichParallelSourceFunction;

import javax.annotation.Nullable;
import java.util.List;

/**
 * A wrapper translating Flink Sources implementing the {@link RichParallelSourceFunction} interface, into
 * unbounded Beam sources (see {@link UnboundedSource}).
 * */
public class UnboundedFlinkSource<T, C extends UnboundedSource.CheckpointMark> extends UnboundedSource<T, C> {

	private final PipelineOptions options;
	private final RichParallelSourceFunction<T> flinkSource;

	public UnboundedFlinkSource(PipelineOptions pipelineOptions, RichParallelSourceFunction<T> source) {
		if(!pipelineOptions.getRunner().equals(FlinkPipelineRunner.class)) {
			throw new RuntimeException("Flink Sources are supported only when running with the FlinkPipelineRunner.");
		}
		options = Preconditions.checkNotNull(pipelineOptions);
		flinkSource = Preconditions.checkNotNull(source);
		validate();
	}

	public RichParallelSourceFunction<T> getFlinkSource() {
		return this.flinkSource;
	}

	@Override
	public List<? extends UnboundedSource<T, C>> generateInitialSplits(int desiredNumSplits, PipelineOptions options) throws Exception {
		throw new RuntimeException("Flink Sources are supported only when running with the FlinkPipelineRunner.");
	}

	@Override
	public UnboundedReader<T> createReader(PipelineOptions options, @Nullable C checkpointMark) {
		throw new RuntimeException("Flink Sources are supported only when running with the FlinkPipelineRunner.");
	}

	@Nullable
	@Override
	public Coder<C> getCheckpointMarkCoder() {
		throw new RuntimeException("Flink Sources are supported only when running with the FlinkPipelineRunner.");
	}


	@Override
	public void validate() {
		Preconditions.checkNotNull(options);
		Preconditions.checkNotNull(flinkSource);
		if(!options.getRunner().equals(FlinkPipelineRunner.class)) {
			throw new RuntimeException("Flink Sources are supported only when running with the FlinkPipelineRunner.");
		}
	}

	@Override
	public Coder<T> getDefaultOutputCoder() {
		throw new RuntimeException("Flink Sources are supported only when running with the FlinkPipelineRunner.");
	}
}

@@ -675,68 +675,8 @@
 o;%0A%0A
-import com.dataartisans.flink.dataflow.FlinkPipelineRunner;%0A
 impo
@@ -941,28 +941,16 @@
 .source.
-RichParallel
 SourceFu
@@ -1083,28 +1083,16 @@
  %7B@link 
-RichParallel
 SourceFu
@@ -1132,12 +1132,16 @@
 ded 
-Beam
+Dataflow
  sou
@@ -1217,19 +1217,17 @@
 Source%3CT
-, C
+%3E
  extends
@@ -1246,97 +1246,46 @@
 urce
-.CheckpointMark%3E extends UnboundedSource%3CT, C%3E %7B%0A%0A%09private final PipelineOptions options;
+%3CT, UnboundedSource.CheckpointMark%3E %7B%0A
 %0A%09pr
@@ -1296,28 +1296,16 @@
 e final 
-RichParallel
 SourceFu
@@ -1361,386 +1361,85 @@
 rce(
-PipelineOptions pipelineOptions, RichParallelSourceFunction%3CT%3E source) %7B%0A%09%09if(!pipelineOptions.getRunner().equals(FlinkPipelineRunner.class)) %7B%0A%09%09%09throw new RuntimeException(%22Flink Sources are supported only when running with the FlinkPipelineRunner.%22);%0A%09%09%7D%0A%09%09options = Preconditions.checkNotNull(pipelineOptions);%0A%09%09flinkSource = Preconditions.checkNotNull(source);%0A%09%09validate(
+SourceFunction%3CT%3E source) %7B%0A%09%09flinkSource = Preconditions.checkNotNull(source
 );%0A%09
@@ -1453,20 +1453,8 @@
 lic 
-RichParallel
 Sour
@@ -1570,17 +1570,46 @@
 urce%3CT, 
-C
+UnboundedSource.CheckpointMark
 %3E%3E gener
@@ -1891,16 +1891,29 @@
 llable C
+heckpointMark
  checkpo
@@ -2072,17 +2072,46 @@
 c Coder%3C
-C
+UnboundedSource.CheckpointMark
 %3E getChe
@@ -2244,17 +2244,16 @@
 %22);%0A%09%7D%0A%0A
-%0A
 %09@Overri
@@ -2286,430 +2286,500 @@
  %7B%0A%09
-%09Preconditions.checkNotNull(options);%0A%09%09Preconditions.checkNotNull(f
+%7D%0A%0A%09@Override%0A%09public Coder%3CT%3E getDefaultOutputCoder() %7B%0A%09%09// The coder is specified in the F
 link
-S
+ s
 ource
-);
 %0A%09%09
-if(!options.getRunner().equals(FlinkPipelineRunner.class)) %7B%0A%09%09%09throw new RuntimeException(%22
+return null;%0A%09%7D%0A%0A%09/**%0A%09 * Creates a new unbounded source from a 
 Flink 
-S
+s
 ource
-s are supported only when running with the FlinkPipelineRunner.%22);%0A%09%09%7D%0A%09%7D%0A%0A%09@Override%0A%09public Coder%3CT%3E getDefaultOutputCoder() %7B%0A%09%09throw new RuntimeException(%22F
+.%0A%09 * @param flinkSource The Flink source function%0A%09 * @param %3CT%3E The type that the source function produces.%0A%09 * @return The wrapped source function.%0A%09 */%0A%09public static %3CT%3E UnboundedSource%3CT, UnboundedSource.CheckpointMark%3E of(SourceFunction%3CT%3E f
 link
- 
 Source
-s are supported only when running with the FlinkPipelineRunner.%22
+) %7B%0A%09%09return new UnboundedFlinkSource%3C%3E(flinkSource
 );%0A%09

fix UnboundedFlinkSource wrapper