package com.firefly.net.tcp.flex.client;

import com.firefly.net.tcp.codec.flex.stream.FlexConnection;
import com.firefly.net.tcp.flex.exception.ConnectionException;
import com.firefly.utils.Assert;
import com.firefly.utils.CollectionUtils;
import com.firefly.utils.concurrent.Scheduler;
import com.firefly.utils.concurrent.Schedulers;
import com.firefly.utils.lang.AbstractLifeCycle;
import com.firefly.utils.lang.HostPort;
import com.firefly.utils.retry.RetryTaskBuilder;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.*;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.stream.Collectors;

import static com.firefly.utils.retry.RetryStrategies.ifException;
import static com.firefly.utils.retry.RetryStrategies.ifResult;
import static com.firefly.utils.retry.StopStrategies.afterExecute;
import static com.firefly.utils.retry.WaitStrategies.exponentialWait;

/**
 * @author Pengtao Qiu
 */
public class FlexConnectionManager extends AbstractLifeCycle {

    protected static final Logger log = LoggerFactory.getLogger("firefly-system");

    private ConcurrentMap<HostPort, FlexConnection> concurrentMap = new ConcurrentHashMap<>();
    private final MultiplexingClient client;
    private final List<HostPort> hostPorts;
    private final AtomicInteger index = new AtomicInteger(0);
    private final Scheduler scheduler = Schedulers.createScheduler();
    private volatile List<HostPort> activatedList;

    public FlexConnectionManager(MultiplexingClient client, Set<String> hostPorts) {
        Assert.notEmpty(hostPorts);
        this.client = client;
        this.hostPorts = convert(hostPorts);
        start();
    }

    public FlexConnection removeConnection(HostPort hostPort) {
        return concurrentMap.remove(hostPort);
    }

    public FlexConnection putConnection(HostPort hostPort, FlexConnection connection) {
        return concurrentMap.put(hostPort, connection);
    }

    public FlexConnection getConnection() {
        FlexConnection ret = RetryTaskBuilder.<FlexConnection>newTask()
                .retry(ifResult(Objects::isNull))
                .stop(afterExecute(hostPorts.size()))
                .wait(exponentialWait(10, TimeUnit.MILLISECONDS))
                .task(() -> {
                    int i = Math.abs(index.getAndAdd(1)) % activatedList.size();
                    return getConnection(activatedList.get(i));
                })
                .call();
        if (ret == null) {
            throw new ConnectionException("Can not get connection");
        }
        return ret;
    }

    private FlexConnection getConnection(HostPort hostPort) {
        try {
            FlexConnection connection = concurrentMap.computeIfAbsent(hostPort, this::createConnection);
            if (connection == null) {
                concurrentMap.remove(hostPort);
                return null;
            }
            if (connection.isOpen()) {
                return connection;
            } else {
                FlexConnection newConnection = createConnection(hostPort);
                if (newConnection != null) {
                    concurrentMap.put(hostPort, newConnection);
                    return newConnection;
                } else {
                    concurrentMap.remove(hostPort);
                    return null;
                }
            }
        } catch (Exception e) {
            log.error("get connection exception", e);
            return null;
        } finally {
            activatedList = new ArrayList<>(concurrentMap.keySet());
        }
    }

    private FlexConnection createConnection(HostPort hostPort) {
        return RetryTaskBuilder.<FlexConnection>newTask()
                .retry(ifException(ex -> ex != null && ex.getCause() instanceof TimeoutException))
                .stop(afterExecute(10))
                .wait(exponentialWait(10, TimeUnit.MILLISECONDS))
                .task(() -> this.connect(hostPort))
                .call();
    }

    private FlexConnection connect(HostPort hostPort) {
        try {
            return client.connect(hostPort.getHost(), hostPort.getPort()).get(5, TimeUnit.SECONDS);
        } catch (Exception e) {
            log.error("Connect " + hostPort + " exception", e);
            throw new ConnectionException("connection exception", e);
        }
    }

    public List<HostPort> getHostPorts() {
        return hostPorts;
    }

    private List<HostPort> convert(Set<String> urls) {
        return Collections.unmodifiableList(
                urls.stream().map(HostPort::new)
                    .sorted(Comparator.comparing(HostPort::toString))
                    .collect(Collectors.toList()));
    }

    @Override
    protected void init() {
        if (!CollectionUtils.isEmpty(hostPorts)) {
            hostPorts.forEach(hostPort -> {
                try {
                    FlexConnection connection = createConnection(hostPort);
                    if (connection != null) {
                        concurrentMap.put(hostPort, connection);
                    }
                } catch (Exception e) {
                    log.error("Connect " + hostPort + " exception", e);
                }
            });
            activatedList = new ArrayList<>(concurrentMap.keySet());
        }
        scheduler.scheduleWithFixedDelay(() -> activatedList = new ArrayList<>(hostPorts), 5, 5, TimeUnit.SECONDS);
    }

    @Override
    protected void destroy() {

    }
}

@@ -765,24 +765,60 @@
 micInteger;%0A
+import java.util.function.Supplier;%0A
 import java.
@@ -1656,16 +1656,82 @@
 tedList;
+%0A    private volatile Supplier%3CList%3CHostPort%3E%3E resetActivatedList;
 %0A%0A    pu
@@ -1959,257 +1959,438 @@
 lic 
-FlexConnection removeConnection(HostPort hostPort) %7B%0A        return concurrentMap.remove(hostPort);%0A    %7D%0A%0A    public FlexConnection putConnection(HostPort hostPort, FlexConnection connection) %7B%0A        return concurrentMap.put(hostPort, connection)
+List%3CHostPort%3E getActivatedList() %7B%0A        return activatedList;%0A    %7D%0A%0A    public void setActivatedList(List%3CHostPort%3E activatedList) %7B%0A        this.activatedList = activatedList;%0A    %7D%0A%0A    public Supplier%3CList%3CHostPort%3E%3E getResetActivatedList() %7B%0A        return resetActivatedList;%0A    %7D%0A%0A    public void setResetActivatedList(Supplier%3CList%3CHostPort%3E%3E resetActivatedList) %7B%0A        this.resetActivatedList = resetActivatedList
 ;%0A  
@@ -2595,25 +2595,29 @@
 Execute(
-hostPorts
+activatedList
 .size())
@@ -5779,16 +5779,33 @@
 edDelay(
+%0A                
 () -%3E ac
@@ -5809,32 +5809,154 @@
  activatedList =
+ Optional.ofNullable(resetActivatedList).map(Supplier::get)%0A                                              .orElseGet(() -%3E
  new ArrayList%3C%3E
@@ -5966,17 +5966,34 @@
 stPorts)
-,
+),%0A               
  5, 5, T

test connection manager