package io.galeb.router.handlers;

import io.galeb.router.client.ExtendedLoadBalancingProxyClient;
import io.galeb.router.client.hostselectors.HostSelector;
import io.galeb.router.client.hostselectors.HostSelectorAlgorithm;
import io.galeb.router.client.hostselectors.HostSelectorInitializer;
import io.galeb.router.services.ExternalData;
import io.undertow.client.UndertowClient;
import io.undertow.server.HttpHandler;
import io.undertow.server.HttpServerExchange;
import io.undertow.server.handlers.ResponseCodeHandler;
import io.undertow.server.handlers.proxy.LoadBalancingProxyClient;
import io.undertow.server.handlers.proxy.ProxyHandler;
import io.undertow.util.Headers;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.context.ApplicationContext;
import org.zalando.boot.etcd.EtcdNode;

import java.net.URI;
import java.util.concurrent.atomic.AtomicBoolean;

import static io.galeb.router.services.ExternalData.POOLS_KEY;

public class PoolHandler implements HttpHandler {

    private final Logger logger = LoggerFactory.getLogger(this.getClass());

    private final HttpHandler defaultHandler;
    private final ExternalData data;
    private final ApplicationContext context;

    private ExtendedProxyHandler proxyHandler = null;
    private String poolname = null;
    private final AtomicBoolean loaded = new AtomicBoolean(false);
    private HostSelectorInitializer hostSelectorInicializer = new HostSelectorInitializer();
    private final ExtendedLoadBalancingProxyClient proxyClient = new ExtendedLoadBalancingProxyClient(UndertowClient.getInstance(), exchange -> {
                                                                        // we always create a new connection for upgrade requests
                                                                        return exchange.getRequestHeaders().contains(Headers.UPGRADE);
                                                                    }, hostSelectorInicializer)
                                                                    .setConnectionsPerThread(2000);

    public PoolHandler(final ApplicationContext context, final ExternalData externalData) {
        this.context = context;
        this.data = externalData;
        this.defaultHandler = buildPoolHandler();
    }

    @Override
    public void handleRequest(HttpServerExchange exchange) throws Exception {
        if (proxyHandler != null) {
            proxyHandler.handleRequest(exchange);
        } else {
            defaultHandler.handleRequest(exchange);
        }
    }

    PoolHandler setPoolName(String aPoolname) {
        poolname = aPoolname;
        return this;
    }

    private HttpHandler buildPoolHandler() {
        return exchange -> {
            synchronized (loaded) {
                loaded.set(true);
                if (poolname != null) {
                    logger.info("creating pool " + poolname);
                    addTargets();
                    defineHostSelector();
                    proxyHandler = context.getBean(ExtendedProxyHandler.class);
                    proxyHandler.setProxyClientAndDefaultHandler(proxyClient, badGatewayExchange -> badGatewayExchange.setStatusCode(502));
                    proxyHandler.handleRequest(exchange);
                    return;
                }
                ResponseCodeHandler.HANDLE_500.handleRequest(exchange);
            }
        };
    }

    private void defineHostSelector() {
        if (poolname != null) {
            final String hostSelectorKeyName = POOLS_KEY + "/" + poolname + "/loadbalance";
            final EtcdNode hostSelectorNode = data.node(hostSelectorKeyName);
            final HostSelector hostSelector;
            if (hostSelectorNode.getKey() != null) {
                String hostSelectorName = hostSelectorNode.getValue();
                hostSelector = HostSelectorAlgorithm.valueOf(hostSelectorName).getHostSelector();
                hostSelectorInicializer.setHostSelector(hostSelector);
            } else {
                hostSelector = hostSelectorInicializer.getHostSelector();
            }
            logger.info("[Pool " + poolname + "] HostSelector: " + hostSelector.getClass().getSimpleName());
        }
    }

    private void addTargets() {
        if (poolname != null) {
            final String poolNameKey = POOLS_KEY + "/" + poolname + "/targets";
            for (EtcdNode etcdNode : data.listFrom(poolNameKey)) {
                String value = etcdNode.getValue();
                URI uri = URI.create(value);
                proxyClient.addHost(uri);
                logger.info("added target " + value);
            }
        }
    }
}

@@ -1214,16 +1214,80 @@
 context;
+%0A    private final ExtendedLoadBalancingProxyClient proxyClient;
 %0A%0A    pr
@@ -1449,16 +1449,22 @@
 private 
+final 
 HostSele
@@ -1540,821 +1540,493 @@
 ();%0A
+%0A
     p
-rivate final ExtendedLoadBalancingProxyClient proxyClient = new ExtendedLoadBalancingProxyClient(UndertowClient.getInstance(), exchange -%3E %7B%0A                                                                        // we always create a new connection for upgrade requests%0A                                                                        return exchange.getRequestHeaders().contains(Headers.UPGRADE);%0A                                                                    %7D, hostSelectorInicializer)%0A                                                                    .setConnectionsPerThread(2000);%0A%0A    public PoolHandler(final ApplicationContext context, final ExternalData externalData) %7B%0A        this.context = context;%0A        this.data = externalData;%0A        this.defaultHandler = buildPoolHandler(
+ublic PoolHandler(final ApplicationContext context, final ExternalData externalData) %7B%0A        this.context = context;%0A        this.data = externalData;%0A        this.defaultHandler = buildPoolHandler();%0A        this.proxyClient =%0A                new ExtendedLoadBalancingProxyClient(UndertowClient.getInstance(),%0A                        exchange -%3E exchange.getRequestHeaders().contains(Headers.UPGRADE), hostSelectorInicializer)%0A                    .setConnectionsPerThread(2000
 );%0A 

improve readability