public final class MesosClientBuilder<Send, Receive> {

    private URI mesosUri;
    private Function<Class<?>, UserAgentEntry> applicationUserAgentEntry;
    private MessageCodec<Send> sendCodec;
    private MessageCodec<Receive> receiveCodec;
    private Send subscribe;
    private Function<Observable<Receive>, Observable<Optional<SinkOperation<Send>>>> streamProcessor;
    private Observable.Transformer<byte[], byte[]> backpressureTransformer;

    private MesosClientBuilder() {
        backpressureTransformer = observable -> observable;
    }

    /**
     * Create a new instance of MesosClientBuilder
     * @param <Send>       The type of objects that will be sent to Mesos
     * @param <Receive>    The type of objects are expected from Mesos
     * @return A new instance of MesosClientBuilder
     */
    @NotNull
    public static <Send, Receive> MesosClientBuilder<Send, Receive> newBuilder() {
        return new MesosClientBuilder<>();
    }

    /**
     * The {@link URI} that should be used to connect to Mesos. The following segments of the URI are used:
     * <ul>
     *     <li>hostname</li>
     *     <li>port</li>
     *     <li>username</li>
     *     <li>password</li>
     *     <li>path</li>
     * </ul>
     * @param mesosUri    Fully qualified URI to use to connect to mesos.
     * @return this builder (allowing for further chained calls)
     */
    @NotNull
    public MesosClientBuilder<Send, Receive> mesosUri(
        @NotNull final URI mesosUri
    ) {
        this.mesosUri = mesosUri;
        return this;
    }

    /**
     * Sets the function used to create a {@link UserAgentEntry} to be included in the {@code User-Agent} header
     * sent to Mesos for all requests.
     * @param applicationUserAgentEntry    Function to provide the {@link UserAgentEntry}
     * @return this builder (allowing for further chained calls)
     */
    @NotNull
    public MesosClientBuilder<Send, Receive> applicationUserAgentEntry(
        @NotNull final Function<Class<?>, UserAgentEntry> applicationUserAgentEntry
    ) {
        this.applicationUserAgentEntry = applicationUserAgentEntry;
        return this;
    }

    /**
     * Allows configuration of the codec used for the {@code Send} type.
     * @param sendCodec    {@link MessageCodec} for {@code Send} type
     * @return this builder (allowing for further chained calls)
     */
    @NotNull
    public MesosClientBuilder<Send, Receive> sendCodec(
        @NotNull final MessageCodec<Send> sendCodec
    ) {
        this.sendCodec = sendCodec;
        return this;
    }

    /**
     * Allows configuration of the codec used for the {@code Receive} type.
     * @param receiveCodec    {@link MessageCodec} for {@code Receive} type
     * @return this builder (allowing for further chained calls)
     */
    @NotNull
    public MesosClientBuilder<Send, Receive> receiveCodec(
        @NotNull final MessageCodec<Receive> receiveCodec
    ) {
        this.receiveCodec = receiveCodec;
        return this;
    }

    /**
     * @param subscribe     The {@code SUBSCRIBE} to be sent to Mesos when opening the event stream.
     * @return this builder (allowing for further chained calls)
     */
    @NotNull
    public MesosClientBuilder<Send, Receive> subscribe(
        @NotNull final Send subscribe
    ) {
        this.subscribe = subscribe;
        return this;
    }

    /**
     * This method provides the means for a user to define how the event stream will be processed.
     * <p>
     * The function passed to this method will function as the actual event processing code for the user.
     * <p>
     * The stream the user will source from is an {@link Observable} of {@code Receive}s. With this stream
     * source a number of functions can be applied to transform/interact/evaluate the stream.
     * <p>
     * The output of this function is the users reaction to each event represented as an
     * {@code Observable<Optional<SinkOperation<Send>>>}. If {@link Optional#isPresent()} the specified
     * {@link SinkOperation} will be processed.
     * <p>
     * For example, if you wanted to log all tasks that result in error:
     * <pre>{@code
     * events -> {
     *     final Observable<Optional<SinkOperation<Call>>> errorLogger = events
     *         .filter(event -> event.getType() == Event.Type.UPDATE && event.getUpdate().getStatus().getState() == TaskState.TASK_ERROR)
     *         .doOnNext(e -> LOGGER.warn("Task Error: {}", ProtoUtils.protoToString(e)))
     *         .map(e -> Optional.empty());
     *
     *     return errorLogger;
     * }
     * }</pre>
     * @param streamProcessing    The function that will be woven between the event spout and the call sink
     * @return this builder (allowing for further chained calls)
     */
    @NotNull
    public MesosClientBuilder<Send, Receive> processStream(
        @NotNull final Function<Observable<Receive>, Observable<Optional<SinkOperation<Send>>>> streamProcessing
    ) {
        this.streamProcessor = streamProcessing;
        return this;
    }

    /**
     * Instructs the HTTP byte[] stream to be composed with reactive pull backpressure such that
     * a burst of incoming Mesos messages is handled by an unbounded buffer rather than a
     * MissingBackpressureException.
     *
     * As an example, this may be necessary for Mesos schedulers that launch large numbers
     * of tasks at a time and then request reconciliation.
     *
     * @return this builder (allowing for further chained calls)
     * @see <a href="http://reactivex.io/documentation/operators/backpressure.html">ReactiveX operators documentation: backpressure operators</a>
     */
    @NotNull
    public MesosClientBuilder<Send, Receive> onBackpressureBuffer(
    ) {
        this.backpressureTransformer = observable -> observable.onBackpressureBuffer();
        return this;
    }


    /**
     * Instructs the HTTP byte[] stream to be composed with reactive pull backpressure such that
     * a burst of incoming Mesos messages is handled by a bounded buffer rather than a
     * MissingBackpressureException. If the buffer is overflown, a {@link java.nio.BufferOverflowException}
     * is thrown.
     *
     * As an example, this may be necessary for Mesos schedulers that launch large numbers
     * of tasks at a time and then request reconciliation.
     *
     * @param capacity number of slots available in the buffer.
     * @return this builder (allowing for further chained calls)
     * @see <a href="http://reactivex.io/documentation/operators/backpressure.html">ReactiveX operators documentation: backpressure operators</a>
     */
    @NotNull
    public MesosClientBuilder<Send, Receive> onBackpressureBuffer(
        final long capacity
    ) {
        this.backpressureTransformer = observable -> observable.onBackpressureBuffer(capacity);
        return this;
    }

    /**
     * Instructs the HTTP byte[] stream to be composed with reactive pull backpressure such that
     * a burst of incoming Mesos messages is handled by a bounded buffer rather than a
     * MissingBackpressureException. If the buffer is overflown, your own custom onOverflow callback
     * will be invoked, and the overflow will mitigate the issue based on the {@link BackpressureOverflow.Strategy}
     * that you select.
     *
     * <ul>
     *     <li>{@link BackpressureOverflow#ON_OVERFLOW_ERROR} (default) will {@code onError} dropping all undelivered items,
     *     unsubscribing from the source, and notifying the producer with {@code onOverflow}. </li>
     *     <li>{@link BackpressureOverflow#ON_OVERFLOW_DROP_LATEST} will drop any new items emitted by the producer while
     *     the buffer is full, without generating any {@code onError}.  Each drop will however invoke {@code onOverflow}
     *     to signal the overflow to the producer.</li>
     *     <li>{@link BackpressureOverflow#ON_OVERFLOW_DROP_OLDEST} will drop the oldest items in the buffer in order to make
     *     room for newly emitted ones. Overflow will not generate an{@code onError}, but each drop will invoke
     *     {@code onOverflow} to signal the overflow to the producer.</li>
     * </ul>
     *
     * As an example, this may be necessary for Mesos schedulers that launch large numbers
     * of tasks at a time and then request reconciliation.
     *
     * @param capacity number of slots available in the buffer.
     * @param onOverflow action to execute if an item needs to be buffered, but there are no available slots.  Null is allowed.
     * @param strategy how should the {@code Observable} react to buffer overflows.
     * @return this builder (allowing for further chained calls)
     * @see <a href="http://reactivex.io/documentation/operators/backpressure.html">ReactiveX operators documentation: backpressure operators</a>
     */
    @NotNull
    public MesosClientBuilder<Send, Receive> onBackpressureBuffer(
            final long capacity,
            @Nullable final Action0 onOverflow,
            @NotNull final BackpressureOverflow.Strategy strategy
            ) {
        this.backpressureTransformer = observable -> observable.onBackpressureBuffer(capacity, onOverflow, strategy);
        return this;
    }

    /**
     * Builds the instance of {@link MesosClient} that has been configured by this builder.
     * All items are expected to have non-null values, if any item is null an exception will be thrown.
     * @return The configured {@link MesosClient}
     */
    @NotNull
    public final MesosClient<Send, Receive> build() {
        return new MesosClient<>(
            checkNotNull(mesosUri),
            checkNotNull(applicationUserAgentEntry),
            checkNotNull(sendCodec),
            checkNotNull(receiveCodec),
            checkNotNull(subscribe),
            checkNotNull(streamProcessor),
            checkNotNull(backpressureTransformer)
        );
    }

} public final class SinkOperation<T> {
    @NotNull
    private final T thingToSink;
    @NotNull
    private final Action1<Throwable> onError;
    @NotNull
    private final Action0 onCompleted;

    /**
     * This constructor is considered an internal API and should not be used directly, instead use one of the
     * factory methods defined in {@link SinkOperations}.
     * @param thingToSink    The {@link T} to send to Mesos
     * @param onCompleted    The callback invoked when HTTP 202 is returned by Mesos
     * @param onError        The callback invoked for an HTTP 400 or 500 status code returned by Mesos
     */
    SinkOperation(
        @NotNull final T thingToSink,
        @NotNull final Action0 onCompleted,
        @NotNull final Action1<Throwable> onError
    ) {
        this.thingToSink = checkNotNull(thingToSink, "argument thingToSink can not be null");
        this.onCompleted = checkNotNull(onCompleted, "argument onCompleted can not be null");
        this.onError = checkNotNull(onError, "argument onError can not be null");
    }

    public void onCompleted() {
        onCompleted.call();
    }

    public void onError(final Throwable e) {
        onError.call(e);
    }

    @NotNull
    public T getThingToSink() {
        return thingToSink;
    }
} public final class SinkOperations {

    @NotNull
    private static final Action1<Throwable> ERROR_NO_OP = (t) -> {};
    @NotNull
    private static final Action0 COMPLETED_NO_OP = () -> {};

    private SinkOperations() {}

    /**
     * Creates a new {@link SinkOperation}.
     * @param thing          The message to be sent to Mesos.
     * @param onCompleted    The callback to be invoked upon a 202 response from Mesos.
     * @param onError        The callback to be invoked upon a 4xx or 5xx response from Mesos.
     * @param <T>            The type of the message to be sent to Mesos.
     * @return  A new {@link SinkOperation} that can be sent to Mesos.
     */
    @NotNull
    public static <T> SinkOperation<T> create(
        @NotNull final T thing,
        @NotNull final Action0 onCompleted,
        @NotNull final Action1<Throwable> onError
    ) {
        return new SinkOperation<>(thing, onCompleted, onError);
    }
    /**
     * Creates a new {@link SinkOperation}.
     * @param thing          The message to be sent to Mesos.
     * @param onCompleted    The callback to be invoked upon a 202 response from Mesos.
     * @param <T>            The type of the message to be sent to Mesos.
     * @return  A new {@link SinkOperation} that can be sent to Mesos.
     */
    @NotNull
    public static <T> SinkOperation<T> create(
        @NotNull final T thing,
        @NotNull final Action0 onCompleted
    ) {
        return create(thing, onCompleted, ERROR_NO_OP);
    }

    /**
     * Creates a new {@link SinkOperation}.
     * @param thing          The message to be sent to Mesos.
     * @param onError        The callback to be invoked upon a 4xx or 5xx response from Mesos.
     * @param <T>            The type of the message to be sent to Mesos.
     * @return  A new {@link SinkOperation} that can be sent to Mesos.
     */
    @NotNull
    public static <T> SinkOperation<T> create(
        @NotNull final T thing,
        @NotNull final Action1<Throwable> onError
    ) {
        return create(thing, COMPLETED_NO_OP, onError);
    }

    /**
     * Creates a new {@link SinkOperation}.
     * @param thing          The message to be sent to Mesos.
     * @param <T>            The type of the message to be sent to Mesos.
     * @return  A new {@link SinkOperation} that can be sent to Mesos.
     */
    @NotNull
    public static <T> SinkOperation<T> create(@NotNull final T thing) {
        return create(thing, COMPLETED_NO_OP, ERROR_NO_OP);
    }

    /**
     * Creates a new {@link SinkOperation}.
     * @param thing          The message to be sent to Mesos.
     * @param onCompleted    The callback to be invoked upon a 202 response from Mesos.
     * @param onError        The callback to be invoked upon a 4xx or 5xx response from Mesos.
     * @param <T>            The type of the message to be sent to Mesos.
     * @return  A new {@link SinkOperation} that can be sent to Mesos.
     * @see #create(Object, Action0, Action1)
     */
    @NotNull
    public static <T> SinkOperation<T> sink(
        @NotNull final T thing,
        @NotNull final Action0 onCompleted,
        @NotNull final Action1<Throwable> onError
    ) {
        return create(thing, onCompleted, onError);
    }

    /**
     * Creates a new {@link SinkOperation}.
     * @param thing          The message to be sent to Mesos.
     * @param onCompleted    The callback to be invoked upon a 202 response from Mesos.
     * @param <T>            The type of the message to be sent to Mesos.
     * @return  A new {@link SinkOperation} that can be sent to Mesos.
     * @see #create(Object, Action0)
     */
    @NotNull
    public static <T> SinkOperation<T> sink(
        @NotNull final T thing,
        @NotNull final Action0 onCompleted
    ) {
        return create(thing, onCompleted);
    }

    /**
     * Creates a new {@link SinkOperation}.
     * @param thing          The message to be sent to Mesos.
     * @param onError        The callback to be invoked upon a 4xx or 5xx response from Mesos.
     * @param <T>            The type of the message to be sent to Mesos.
     * @return  A new {@link SinkOperation} that can be sent to Mesos.
     * @see #create(Object, Action1)
     */
    @NotNull
    public static <T> SinkOperation<T> sink(
        @NotNull final T thing,
        @NotNull final Action1<Throwable> onError
    ) {
        return create(thing, onError);
    }

    /**
     * Creates a new {@link SinkOperation}.
     * @param thing          The message to be sent to Mesos.
     * @param <T>            The type of the message to be sent to Mesos.
     * @return  A new {@link SinkOperation} that can be sent to Mesos.
     * @see #create(Object)
     */
    @NotNull
    public static <T> SinkOperation<T> sink(
        @NotNull final T thing
    ) {
        return create(thing);
    }


} public final class ProtoUtils {
    private static final Pattern PROTO_TO_STRING = Pattern.compile(" *\\n *");

    private ProtoUtils() {}

    /**
     * Utility method to make logging of protos a little better than what is done by default.
     * @param message    The proto message to be "stringified"
     * @return  A string representation of the provided {@code message} surrounded with curly braces ('{}') and all
     *          new lines replaced with a space.
     */
    @NotNull
    public static String protoToString(@NotNull final Object message) {
        return "{ " + PROTO_TO_STRING.matcher(message.toString()).replaceAll(" ").trim() + " }";
    }

} public final class ProtobufMesosClientBuilder {
    private ProtobufMesosClientBuilder() {}

    /**
     * @return  An initial {@link MesosClientBuilder} that will use protobuf
     *          for the {@link org.apache.mesos.v1.scheduler.Protos.Call Call} and
     *          {@link org.apache.mesos.v1.scheduler.Protos.Event Event} messages.
     */
    @NotNull
    public static MesosClientBuilder<Protos.Call, Protos.Event> schedulerUsingProtos() {
        return MesosClientBuilder.<Protos.Call, Protos.Event>newBuilder()
            .sendCodec(ProtobufMessageCodecs.SCHEDULER_CALL)
            .receiveCodec(ProtobufMessageCodecs.SCHEDULER_EVENT)
            ;
    }

    /**
     * @return  An initial {@link MesosClientBuilder} that will use protobuf
     *          for the {@link org.apache.mesos.v1.executor.Protos.Call Call} and
     *          {@link org.apache.mesos.v1.executor.Protos.Event Event} messages.
     */
    @NotNull
    public static MesosClientBuilder<
        org.apache.mesos.v1.executor.Protos.Call,
        org.apache.mesos.v1.executor.Protos.Event
        > executorUsingProtos() {
        return MesosClientBuilder.<org.apache.mesos.v1.executor.Protos.Call, org.apache.mesos.v1.executor.Protos.Event>newBuilder()
            .sendCodec(ProtobufMessageCodecs.EXECUTOR_CALL)
            .receiveCodec(ProtobufMessageCodecs.EXECUTOR_EVENT)
            ;
    }
} public final class SchedulerCalls {

    private SchedulerCalls() {}

    /**
     * Utility method to more succinctly construct a {@link Call Call} of type {@link Type#ACKNOWLEDGE ACKNOWLEDGE}.
     * <p>
     *
     * @param frameworkId    The {@link Protos.FrameworkID} to be set on the {@link Call}
     * @param uuid           The {@link Protos.TaskStatus#getUuid() uuid} from the
     *                       {@link org.apache.mesos.v1.scheduler.Protos.Event.Update#getStatus() TaskStatus} received from Mesos.
     * @param agentId        The {@link Protos.TaskStatus#getAgentId() agentId} from the
     *                       {@link org.apache.mesos.v1.scheduler.Protos.Event.Update#getStatus() TaskStatus} received from Mesos.
     * @param taskId         The {@link Protos.TaskStatus#getTaskId() taskId} from the
     *                       {@link org.apache.mesos.v1.scheduler.Protos.Event.Update#getStatus() TaskStatus} received from Mesos.
     * @return  A {@link Call} with a configured {@link Acknowledge}.
     */
    @NotNull
    public static Call ackUpdate(
        @NotNull final Protos.FrameworkID frameworkId,
        @NotNull final ByteString uuid,
        @NotNull final Protos.AgentID agentId,
        @NotNull final Protos.TaskID taskId
    ) {
        return newBuilder()
            .setFrameworkId(frameworkId)
            .setType(Type.ACKNOWLEDGE)
            .setAcknowledge(
                Acknowledge.newBuilder()
                    .setUuid(uuid)
                    .setAgentId(agentId)
                    .setTaskId(taskId)
                    .build()
            )
            .build();
    }

    /**
     * Utility method to more succinctly construct a {@link Call Call} of type {@link Type#DECLINE DECLINE}.
     * <p>
     *
     * @param frameworkId    The {@link Protos.FrameworkID FrameworkID} to be set on the {@link Call}
     * @param offerIds       A list of {@link Protos.OfferID OfferID} from the
     *                       {@link Protos.Offer}s received from Mesos.
     * @return  A {@link Call} with a configured {@link Decline}.
     */
    @NotNull
    public static Call decline(@NotNull final Protos.FrameworkID frameworkId, @NotNull final List<Protos.OfferID> offerIds) {
        return newBuilder()
            .setFrameworkId(frameworkId)
            .setType(Type.DECLINE)
            .setDecline(
                Decline.newBuilder()
                    .addAllOfferIds(offerIds)
            )
            .build();
    }

    /**
     * Utility method to more succinctly construct a {@link Call Call} of type {@link Type#SUBSCRIBE SUBSCRIBE}.
     * <p>
     *
     * @param frameworkId               The frameworkId to set on the {@link Protos.FrameworkInfo FrameworkInfo} and
     *                                  {@link Call Call} messages.
     * @param user                      The user to set on the {@link Protos.FrameworkInfo FrameworkInfo} message.
     * @param frameworkName             The name to set on the {@link Protos.FrameworkInfo FrameworkInfo} message.
     * @param failoverTimeoutSeconds    The failoverTimeoutSeconds to set on the
     *                                  {@link Protos.FrameworkInfo FrameworkInfo} message.
     * @return An {@link Call Call} of type {@link Type#SUBSCRIBE SUBSCRIBE} with the configured
     * {@link Subscribe Subscribe} sub-message.
     */
    @NotNull
    public static Call subscribe(
        @NotNull final String frameworkId,
        @NotNull final String user,
        @NotNull final String frameworkName,
        final long failoverTimeoutSeconds
    ) {
        final Protos.FrameworkID frameworkID = Protos.FrameworkID.newBuilder().setValue(frameworkId).build();
        return subscribe(frameworkID, user, frameworkName, failoverTimeoutSeconds);
    }

    /**
     * Utility method to more succinctly construct a {@link Call Call} of type {@link Type#SUBSCRIBE SUBSCRIBE}.
     * <p>
     *
     * @param frameworkId               The frameworkId to set on the {@link Protos.FrameworkInfo FrameworkInfo} and
     *                                  {@link Call Call} messages.
     * @param user                      The user to set on the {@link Protos.FrameworkInfo FrameworkInfo} message.
     * @param frameworkName             The name to set on the {@link Protos.FrameworkInfo FrameworkInfo} message.
     * @param failoverTimeoutSeconds    The failoverTimeoutSeconds to set on the
     *                                  {@link Protos.FrameworkInfo FrameworkInfo} message.
     * @return An {@link Call Call} of type {@link Type#SUBSCRIBE SUBSCRIBE} with the configured
     * {@link Subscribe Subscribe} sub-message.
     */
    @NotNull
    public static Call subscribe(
        @NotNull final Protos.FrameworkID frameworkId,
        @NotNull final String user,
        @NotNull final String frameworkName,
        final long failoverTimeoutSeconds
    ) {
        final Protos.FrameworkInfo frameworkInfo = Protos.FrameworkInfo.newBuilder()
            .setId(frameworkId)
            .setUser(user)
            .setName(frameworkName)
            .setFailoverTimeout(failoverTimeoutSeconds)
            .build();
        return subscribe(frameworkId, frameworkInfo);
    }

    /**
     * Utility method to more succinctly construct a {@link Call Call} of type {@link Type#SUBSCRIBE SUBSCRIBE}.
     * <p>
     *
     * @param frameworkId   The frameworkId to set on the {@link Call Call} messages.
     * @param frameworkInfo The frameworkInfo to set on the {@link Subscribe Subscribe} sub-message.
     * @return An {@link Call Call} of type {@link Type#SUBSCRIBE SUBSCRIBE} with the configured
     * {@link Subscribe Subscribe} sub-message.
     */
    @NotNull
    public static Call subscribe(
        @NotNull final Protos.FrameworkID frameworkId,
        @NotNull final Protos.FrameworkInfo frameworkInfo
    ) {
        return newBuilder()
            .setFrameworkId(frameworkId)
            .setType(Type.SUBSCRIBE)
            .setSubscribe(
                Subscribe.newBuilder()
                    .setFrameworkInfo(frameworkInfo)
            )
            .build();
    }
} @NotNull
public static <T> SinkOperation<T> sink(
    @NotNull final T thing,
    @NotNull final Action0 onCompleted,
    @NotNull final Action1<Throwable> onError
) {
    return create(thing, onCompleted, onError);
} @NotNull
public static Call decline(@NotNull final Protos.FrameworkID frameworkId, @NotNull final List<Protos.OfferID> offerIds) {
    return newBuilder()
        .setFrameworkId(frameworkId)
        .setType(Type.DECLINE)
        .setDecline(
            Decline.newBuilder()
                .addAllOfferIds(offerIds)
        )
        .build();
} @NotNull
public static Call subscribe(
    @NotNull final String frameworkId,
    @NotNull final String user,
    @NotNull final String frameworkName,
    final long failoverTimeoutSeconds
) {
    final Protos.FrameworkID frameworkID = Protos.FrameworkID.newBuilder().setValue(frameworkId).build();
    return subscribe(frameworkID, user, frameworkName, failoverTimeoutSeconds);
} @NotNull
public static Function<Class<?>, UserAgentEntry> userAgentEntryForMavenArtifact(@NotNull final String groupId, @NotNull final String artifactId) {
    return (Class<?> c) -> {
        final Properties props = loadProperties(c, String.format("/META-INF/maven/%s/%s/pom.properties", groupId, artifactId));
        return new UserAgentEntry(props.getProperty("artifactId", artifactId), props.getProperty("version", "unknown-version"));
    };
}
import com.mesosphere.mesos.rx.java.MesosClientBuilder;
import com.mesosphere.mesos.rx.java.SinkOperation;
import com.mesosphere.mesos.rx.java.SinkOperations;
import com.mesosphere.mesos.rx.java.protobuf.ProtoUtils;
import com.mesosphere.mesos.rx.java.protobuf.ProtobufMesosClientBuilder;
import com.mesosphere.mesos.rx.java.protobuf.SchedulerCalls;
import org.apache.mesos.v1.Protos;
import org.apache.mesos.v1.Protos.*;
import org.apache.mesos.v1.scheduler.Protos.Call;
import org.apache.mesos.v1.scheduler.Protos.Event;
import org.jetbrains.annotations.NotNull;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import rx.Observable;
import java.net.URI;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.UUID;
import static com.google.common.base.Preconditions.checkState;
import static com.google.common.collect.Lists.newArrayList;
import static com.mesosphere.mesos.rx.java.SinkOperations.sink;
import static com.mesosphere.mesos.rx.java.protobuf.SchedulerCalls.decline;
import static com.mesosphere.mesos.rx.java.protobuf.SchedulerCalls.subscribe;
import static com.mesosphere.mesos.rx.java.util.UserAgentEntries.userAgentEntryForMavenArtifact;
import static java.util.stream.Collectors.groupingBy;
import static rx.Observable.from;
import static rx.Observable.just;
/*
 *    Copyright (C) 2015 Mesosphere, Inc
 *
 * 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.mesosphere.mesos.rx.java.example.framework.sleepy;




/**
 * A relatively simple Mesos framework that launches {@code sleep $SLEEP_SECONDS} tasks for offers it receives.
 * This framework uses the Mesos HTTP Scheduler API.
 * @see <a href="https://github.com/mesosphere/mesos-rxjava/blob/master/mesos-rxjava-example/mesos-rxjava-example-framework/src/main/java/com/mesosphere/mesos/rx/java/example/framework/sleepy/Sleepy.java">Sleepy.java</a>
 */
public final class Sleepy {
    private static final Logger LOGGER = LoggerFactory.getLogger(Sleepy.class);

    /**
     * <pre>{@code
     * Usage: java -cp <application-jar> com.mesosphere.mesos.rx.java.example.framework.sleepy.Sleepy <mesos-uri> <cpus-per-task> <mesos-resource-role>
     * mesos-uri            The fully qualified URI to the Mesos Master. (http://localhost:5050/api/v1/scheduler)
     * cpus-per-task        The number of CPUs each task should claim from an offer.
     * mesos-resources-role The resource role to use when registering with mesos and evaluating offers
     * }</pre>
     * @param args    Application arguments {@code mesos-uri}, {@code cpus-per-task} and {@code mesos-resource-role}.
     */
    public static void main(final String... args) {
        try {
            _main("sleepy-" + UUID.randomUUID(), args);
        } catch (Throwable e) {
            LOGGER.error("Unhandled exception caught at main", e);
            System.exit(1);
        }
    }

    static void _main(final String fwId, final String... args) throws Throwable {
        if (args.length != 3) {
            final String className = Sleepy.class.getCanonicalName();
            System.err.println("Usage: java -cp <application-jar> " + className + " <mesos-uri> <cpus-per-task> <mesos-resource-role>");
            System.exit(1);
        }

        final URI mesosUri = URI.create(args[0]);
        final double cpusPerTask = Double.parseDouble(args[1]);
        final String role = args[2];
        checkState(role != null && !role.trim().isEmpty(), "<mesos-resource-role> must not be empty");
        final FrameworkID frameworkID = FrameworkID.newBuilder().setValue(fwId).build();
        final State<FrameworkID, TaskID, TaskState> stateObject = new State<>(frameworkID, role.trim(), cpusPerTask, 16);

        final MesosClientBuilder<Call, Event> clientBuilder = ProtobufMesosClientBuilder.schedulerUsingProtos()
            .mesosUri(mesosUri)
            .applicationUserAgentEntry(userAgentEntryForMavenArtifact("com.mesosphere.mesos.rx.java.example", "mesos-rxjava-example-framework"));

        final Call subscribeCall = subscribe(
            stateObject.getFwId(),
            Protos.FrameworkInfo.newBuilder()
                .setId(stateObject.getFwId())
                .setUser(Optional.ofNullable(System.getenv("user")).orElse("root")) // https://issues.apache.org/jira/browse/MESOS-3747
                .setName("sleepy")
                .setFailoverTimeout(0)
                .setRole(stateObject.getResourceRole())
                .build()
        );

        final Observable<State<FrameworkID, TaskID, TaskState>> stateObservable = just(stateObject).repeat();

        clientBuilder
            .subscribe(subscribeCall)
            .processStream(unicastEvents -> {
                final Observable<Event> events = unicastEvents.share();

                final Observable<Optional<SinkOperation<Call>>> offerEvaluations = events
                    .filter(event -> event.getType() == Event.Type.OFFERS)
                    .flatMap(event -> from(event.getOffers().getOffersList()))
                    .zipWith(stateObservable, Tuple2::create)
                    .map(Sleepy::handleOffer)
                    .map(Optional::of);

                final Observable<Optional<SinkOperation<Call>>> updateStatusAck = events
                    .filter(event -> event.getType() == Event.Type.UPDATE && event.getUpdate().getStatus().hasUuid())
                    .zipWith(stateObservable, Tuple2::create)
                    .doOnNext((Tuple2<Event, State<FrameworkID, TaskID, TaskState>> t) -> {
                        final Event event = t._1;
                        final State<FrameworkID, TaskID, TaskState> state = t._2;
                        final TaskStatus status = event.getUpdate().getStatus();
                        state.put(status.getTaskId(), status.getState());
                    })
                    .map((Tuple2<Event, State<FrameworkID, TaskID, TaskState>> t) -> {
                        final TaskStatus status = t._1.getUpdate().getStatus();
                        return SchedulerCalls.ackUpdate(t._2.getFwId(), status.getUuid(), status.getAgentId(), status.getTaskId());
                    })
                    .map(SinkOperations::create)
                    .map(Optional::of);

                final Observable<Optional<SinkOperation<Call>>> errorLogger = events
                    .filter(event -> event.getType() == Event.Type.ERROR || (event.getType() == Event.Type.UPDATE && event.getUpdate().getStatus().getState() == TaskState.TASK_ERROR))
                    .doOnNext(e -> LOGGER.warn("Task Error: {}", ProtoUtils.protoToString(e)))
                    .map(e -> Optional.empty());

                return offerEvaluations.mergeWith(updateStatusAck).mergeWith(errorLogger);
            });

        clientBuilder.build().openStream().await();
    }

    @NotNull
    private static SinkOperation<Call> handleOffer(@NotNull final Tuple2<Offer, State<FrameworkID, TaskID, TaskState>> t) {
        final Offer offer = t._1;
        final State<FrameworkID, TaskID, TaskState> state = t._2;
        final int offerCount = state.getOfferCounter().incrementAndGet();
        final String desiredRole = state.getResourceRole();

        final FrameworkID frameworkId = state.getFwId();
        final AgentID agentId = offer.getAgentId();
        final List<OfferID> ids = newArrayList(offer.getId());

        final Map<String, List<Resource>> resources = offer.getResourcesList()
            .stream()
            .collect(groupingBy(Resource::getName));
        final List<Resource> cpuList = resources.get("cpus");
        final List<Resource> memList = resources.get("mem");
        if (
            cpuList != null && !cpuList.isEmpty()
                && memList != null && !memList.isEmpty()
                && cpuList.size() == memList.size()
            ) {
            final List<TaskInfo> tasks = newArrayList();
            for (int i = 0; i < cpuList.size(); i++) {
                final Resource cpus = cpuList.get(i);
                final Resource mem = memList.get(i);

                if (desiredRole.equals(cpus.getRole()) && desiredRole.equals(mem.getRole())) {
                    double availableCpu = cpus.getScalar().getValue();
                    double availableMem = mem.getScalar().getValue();
                    final double cpusPerTask = state.getCpusPerTask();
                    final double memMbPerTask = state.getMemMbPerTask();
                    while (availableCpu >= cpusPerTask && availableMem >= memMbPerTask) {
                        availableCpu -= cpusPerTask;
                        availableMem -= memMbPerTask;
                        final String taskId = String.format("task-%d-%d", offerCount, state.getTotalTaskCounter().incrementAndGet());
                        tasks.add(sleepTask(agentId, taskId, cpus.getRole(), cpusPerTask, mem.getRole(), memMbPerTask));
                    }
                }
            }

            if (!tasks.isEmpty()) {
                LOGGER.info("Launching {} tasks", tasks.size());
                return sink(
                    sleep(frameworkId, ids, tasks),
                    () -> tasks.forEach(task -> state.put(task.getTaskId(), TaskState.TASK_STAGING)),
                    (e) -> LOGGER.warn("", e)
                );
            } else {
                return sink(decline(frameworkId, ids));