/**
 * Copyright 2016 Netflix, 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 io.reactivesocket.client;

import io.reactivesocket.ReactiveSocket;
import io.reactivesocket.ReactiveSocketConnector;
import io.reactivesocket.ReactiveSocketFactory;
import io.reactivesocket.client.filter.*;
import org.reactivestreams.Publisher;
import org.reactivestreams.Subscriber;
import org.reactivestreams.Subscription;

import java.net.SocketAddress;
import java.util.*;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.ThreadFactory;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.function.Function;
import java.util.stream.Collectors;

public class ClientBuilder {
    private static AtomicInteger counter = new AtomicInteger(0);
    private final String name;

    private final ScheduledExecutorService executor;

    private final long requestTimeout;
    private final TimeUnit requestTimeoutUnit;

    private final long connectTimeout;
    private final TimeUnit connectTimeoutUnit;

    private final double backupQuantile;

    private final int retries;

    private final ReactiveSocketConnector<SocketAddress> connector;
    private final Function<Throwable, Boolean> retryThisException;

    private final Publisher<List<SocketAddress>> source;

    private ClientBuilder(
        String name,
        ScheduledExecutorService executor,
        long requestTimeout, TimeUnit requestTimeoutUnit,
        long connectTimeout, TimeUnit connectTimeoutUnit,
        double backupQuantile,
        int retries, Function<Throwable, Boolean> retryThisException,
        ReactiveSocketConnector<SocketAddress> connector,
        Publisher<List<SocketAddress>> source
    ) {
        this.name = name;
        this.executor = executor;
        this.requestTimeout = requestTimeout;
        this.requestTimeoutUnit = requestTimeoutUnit;
        this.connectTimeout = connectTimeout;
        this.connectTimeoutUnit = connectTimeoutUnit;
        this.backupQuantile = backupQuantile;
        this.retries = retries;
        this.connector = connector;
        this.retryThisException = retryThisException;
        this.source = source;
    }

    public ClientBuilder withRequestTimeout(long timeout, TimeUnit unit) {
        return new ClientBuilder(
            name,
            executor,
            timeout, unit,
            connectTimeout, connectTimeoutUnit,
            backupQuantile,
            retries, retryThisException,
            connector,
            source
        );
    }

    public ClientBuilder withConnectTimeout(long timeout, TimeUnit unit) {
        return new ClientBuilder(
            name,
            executor,
            requestTimeout, requestTimeoutUnit,
            timeout, unit,
            backupQuantile,
            retries, retryThisException,
            connector,
            source
        );
    }

    public ClientBuilder withBackupRequest(double quantile) {
        return new ClientBuilder(
            name,
            executor,
            requestTimeout, requestTimeoutUnit,
            connectTimeout, connectTimeoutUnit,
            quantile,
            retries, retryThisException,
            connector,
            source
        );
    }

    public ClientBuilder withExecutor(ScheduledExecutorService executor) {
        return new ClientBuilder(
            name,
            executor,
            requestTimeout, requestTimeoutUnit,
            connectTimeout, connectTimeoutUnit,
            backupQuantile,
            retries, retryThisException,
            connector,
            source
        );
    }

    public ClientBuilder withConnector(ReactiveSocketConnector<SocketAddress> connector) {
        return new ClientBuilder(
            name,
            executor,
            requestTimeout, requestTimeoutUnit,
            connectTimeout, connectTimeoutUnit,
            backupQuantile,
            retries, retryThisException,
            connector,
            source
        );
    }

    public ClientBuilder withSource(Publisher<List<SocketAddress>> source) {
        return new ClientBuilder(
            name,
            executor,
            requestTimeout, requestTimeoutUnit,
            connectTimeout, connectTimeoutUnit,
            backupQuantile,
            retries, retryThisException,
            connector,
            source
        );
    }

    public ClientBuilder withRetries(int nbOfRetries, Function<Throwable, Boolean> retryThisException) {
        return new ClientBuilder(
            name,
            executor,
            requestTimeout, requestTimeoutUnit,
            connectTimeout, connectTimeoutUnit,
            backupQuantile,
            nbOfRetries, retryThisException,
            connector,
            source
        );
    }

    public ReactiveSocket build() {
        if (source == null) {
            throw new IllegalStateException("Please configure the source!");
        }
        if (connector == null) {
            throw new IllegalStateException("Please configure the connector!");
        }

        ReactiveSocketConnector<SocketAddress> filterConnector = connector
            .chain(socket -> new TimeoutSocket(socket, requestTimeout, requestTimeoutUnit, executor))
            .chain(DrainingSocket::new);

        Publisher<List<ReactiveSocketFactory<SocketAddress>>> factories =
            sourceToFactory(source, filterConnector);

        ReactiveSocket socket = new LoadBalancer(factories);
        if (0.0 < backupQuantile && backupQuantile < 1.0) {
            socket = new BackupRequestSocket(socket, backupQuantile, executor);
        }
        if (retries > 0) {
            socket = new RetrySocket(socket, retries, t -> true);
        }
        return socket;
    }

    private Publisher<List<ReactiveSocketFactory<SocketAddress>>> sourceToFactory(
        Publisher<List<SocketAddress>> source,
        ReactiveSocketConnector<SocketAddress> connector
    ) {
        return subscriber ->
            source.subscribe(new Subscriber<List<SocketAddress>>() {
                private Map<SocketAddress, ReactiveSocketFactory<SocketAddress>> current;

                @Override
                public void onSubscribe(Subscription s) {
                    subscriber.onSubscribe(s);
                    current = new HashMap<>();
                }

                @Override
                public void onNext(List<SocketAddress> socketAddresses) {
                    socketAddresses.stream()
                        .filter(sa -> !current.containsKey(sa))
                        .map(connector::toFactory)
                        .map(factory -> factory.chain(TimeoutFactory.asChainFunction(connectTimeout, connectTimeoutUnit,
                                                                                     executor)))
                        .map(FailureAwareFactory::new)
                        .forEach(factory -> current.put(factory.remote(), factory));

                    Set<SocketAddress> addresses = new HashSet<>(socketAddresses);
                    Iterator<Map.Entry<SocketAddress, ReactiveSocketFactory<SocketAddress>>> it =
                        current.entrySet().iterator();
                    while (it.hasNext()) {
                        SocketAddress sa = it.next().getKey();
                        if (! addresses.contains(sa)) {
                            it.remove();
                        }
                    }

                    List<ReactiveSocketFactory<SocketAddress>> factories =
                        current.values().stream().collect(Collectors.toList());
                    subscriber.onNext(factories);
                }

                @Override
                public void onError(Throwable t) { subscriber.onError(t); }

                @Override
                public void onComplete() { subscriber.onComplete(); }
            });
    }

    public static ClientBuilder instance() {
        return new ClientBuilder(
            "rs-loadbalancer-" + counter.incrementAndGet(),
            Executors.newScheduledThreadPool(4, new ThreadFactory() {
                @Override
                public Thread newThread(Runnable r) {
                    Thread thread = new Thread(r);
                    thread.setName("reactivesocket-scheduler-thread");
                    thread.setDaemon(true);
                    return thread;
                }
            }),
            1, TimeUnit.SECONDS,
            10, TimeUnit.SECONDS,
            0.99,
            3, t -> true,
            null,
            null
        );
    }
}


@@ -3491,363 +3491,8 @@
  %7D%0A%0A
-    public ClientBuilder withBackupRequest(double quantile) %7B%0A        return new ClientBuilder(%0A            name,%0A            executor,%0A            requestTimeout, requestTimeoutUnit,%0A            connectTimeout, connectTimeoutUnit,%0A            quantile,%0A            retries, retryThisException,%0A            connector,%0A            source%0A        );%0A    %7D%0A%0A
     
@@ -3558,24 +3558,24 @@
 executor) %7B%0A
+
         retu
@@ -4631,416 +4631,8 @@
  %7D%0A%0A
-    public ClientBuilder withRetries(int nbOfRetries, Function%3CThrowable, Boolean%3E retryThisException) %7B%0A        return new ClientBuilder(%0A            name,%0A            executor,%0A            requestTimeout, requestTimeoutUnit,%0A            connectTimeout, connectTimeoutUnit,%0A            backupQuantile,%0A            nbOfRetries, retryThisException,%0A            connector,%0A            source%0A        );%0A    %7D%0A%0A
     
@@ -5264,31 +5264,14 @@
     
-ReactiveSocket socket =
+return
  new
@@ -5300,284 +5300,8 @@
 s);%0A
-        if (0.0 %3C backupQuantile && backupQuantile %3C 1.0) %7B%0A            socket = new BackupRequestSocket(socket, backupQuantile, executor);%0A        %7D%0A        if (retries %3E 0) %7B%0A            socket = new RetrySocket(socket, retries, t -%3E true);%0A        %7D%0A        return socket;%0A
     

Disable Backup-Request / Retries. (#102)