/**
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you 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.twitter.distributedlog.service.balancer;

import com.google.common.base.Optional;
import com.google.common.util.concurrent.RateLimiter;
import com.twitter.distributedlog.client.monitor.MonitorServiceClient;
import com.twitter.distributedlog.service.ClientUtils;
import com.twitter.distributedlog.service.DLSocketAddress;
import com.twitter.distributedlog.service.DistributedLogClient;
import com.twitter.distributedlog.service.DistributedLogClientBuilder;
import com.twitter.util.Await;
import com.twitter.util.Function;
import com.twitter.util.Future;
import com.twitter.util.FutureEventListener;
import org.apache.commons.lang3.tuple.Pair;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.Serializable;
import java.net.SocketAddress;
import java.util.ArrayList;
import java.util.Collections;
import java.util.Comparator;
import java.util.LinkedList;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * A balancer balances ownerships with a cluster of targets
 */
public class ClusterBalancer implements Balancer {

    static final Logger logger = LoggerFactory.getLogger(ClusterBalancer.class);

    /**
     * Represent a single host. Ordered by number of streams in desc order.
     */
    static class Host {

        final SocketAddress address;
        final Set<String> streams;
        final DistributedLogClientBuilder clientBuilder;
        DistributedLogClient client = null;
        MonitorServiceClient monitor = null;

        Host(SocketAddress address, Set<String> streams,
             DistributedLogClientBuilder clientBuilder) {
            this.address = address;
            this.streams = streams;
            this.clientBuilder = clientBuilder;
        }

        private void initializeClientsIfNeeded() {
            if (null == client) {
                Pair<DistributedLogClient, MonitorServiceClient> clientPair =
                        createDistributedLogClient(address, clientBuilder);
                client = clientPair.getLeft();
                monitor = clientPair.getRight();
            }
        }

        synchronized DistributedLogClient getClient() {
            initializeClientsIfNeeded();
            return client;
        }

        synchronized MonitorServiceClient getMonitor() {
            initializeClientsIfNeeded();
            return monitor;
        }

        synchronized void close() {
            if (null != client) {
                client.close();
            }
        }

        @Override
        public String toString() {
            StringBuilder sb = new StringBuilder();
            sb.append("Host(").append(address).append(")");
            return sb.toString();
        }
    }

    static class HostComparator implements Comparator<Host>, Serializable {
        private static final long serialVersionUID = 7984973796525102538L;

        @Override
        public int compare(Host h1, Host h2) {
            return h2.streams.size() - h1.streams.size();
        }
    }

    protected final DistributedLogClientBuilder clientBuilder;
    protected final DistributedLogClient client;
    protected final MonitorServiceClient monitor;

    public ClusterBalancer(DistributedLogClientBuilder clientBuilder) {
        this(clientBuilder, ClientUtils.buildClient(clientBuilder));
    }

    ClusterBalancer(DistributedLogClientBuilder clientBuilder,
                    Pair<DistributedLogClient, MonitorServiceClient> clientPair) {
        this.clientBuilder = clientBuilder;
        this.client = clientPair.getLeft();
        this.monitor = clientPair.getRight();
    }

    /**
     * Build a new distributedlog client to a single host <i>host</i>.
     *
     * @param host
     *          host to access
     * @return distributedlog clients
     */
    static Pair<DistributedLogClient, MonitorServiceClient> createDistributedLogClient(
            SocketAddress host, DistributedLogClientBuilder clientBuilder) {
        DistributedLogClientBuilder newBuilder =
                DistributedLogClientBuilder.newBuilder(clientBuilder).host(host);
        return ClientUtils.buildClient(newBuilder);
    }

    @Override
    public void balanceAll(String source,
                           int rebalanceConcurrency, /* unused */
                           Optional<RateLimiter> rebalanceRateLimiter) {
        balance(0, 0.0f, rebalanceConcurrency, Optional.of(source), rebalanceRateLimiter);
    }

    @Override
    public void balance(int rebalanceWaterMark,
                        double rebalanceTolerancePercentage,
                        int rebalanceConcurrency, /* unused */
                        Optional<RateLimiter> rebalanceRateLimiter) {
        Optional<String> source = Optional.absent();
        balance(rebalanceWaterMark, rebalanceTolerancePercentage, rebalanceConcurrency, source, rebalanceRateLimiter);
    }

    public void balance(int rebalanceWaterMark,
                        double rebalanceTolerancePercentage,
                        int rebalanceConcurrency,
                        Optional<String> source,
                        Optional<RateLimiter> rebalanceRateLimiter) {
        Map<SocketAddress, Set<String>> distribution = monitor.getStreamOwnershipDistribution();
        if (distribution.size() <= 1) {
            return;
        }
        SocketAddress sourceAddr = null;
        if (source.isPresent()) {
            sourceAddr = DLSocketAddress.parseSocketAddress(source.get());
            logger.info("Balancer source is {}", sourceAddr);
            if (!distribution.containsKey(sourceAddr)) {
                return;
            }
        }
        // Get the list of hosts ordered by number of streams in DESC order
        List<Host> hosts = new ArrayList<Host>(distribution.size());
        for (Map.Entry<SocketAddress, Set<String>> entry : distribution.entrySet()) {
            Host host = new Host(entry.getKey(), entry.getValue(), clientBuilder);
            hosts.add(host);
        }
        Collections.sort(hosts, new HostComparator());
        try {

            // find the host to move streams from.
            int hostIdxMoveFrom = -1;
            if (null != sourceAddr) {
                for (Host host : hosts) {
                    ++hostIdxMoveFrom;
                    if (sourceAddr.equals(host.address)) {
                        break;
                    }
                }
            }

            // compute the average load.
            int totalStream = 0;
            for (Host host : hosts) {
                totalStream += host.streams.size();
            }
            double averageLoad;
            if (hostIdxMoveFrom >= 0) {
                averageLoad = ((double) totalStream / (hosts.size() - 1));
            } else {
                averageLoad = ((double) totalStream / hosts.size());
            }

            int moveFromLowWaterMark;
            int moveToHighWaterMark = Math.max(1, (int) (averageLoad + averageLoad * rebalanceTolerancePercentage / 100.0f));

            if (hostIdxMoveFrom >= 0) {
                moveFromLowWaterMark = Math.max(0, rebalanceWaterMark);
                moveStreams(
                        hosts,
                        new AtomicInteger(hostIdxMoveFrom), moveFromLowWaterMark,
                        new AtomicInteger(hosts.size() - 1), moveToHighWaterMark,
                        rebalanceRateLimiter);
                moveRemainingStreamsFromSource(hosts.get(hostIdxMoveFrom), hosts, rebalanceRateLimiter);
            } else {
                moveFromLowWaterMark = Math.max((int) Math.ceil(averageLoad), rebalanceWaterMark);
                AtomicInteger moveFrom = new AtomicInteger(0);
                AtomicInteger moveTo = new AtomicInteger(hosts.size() - 1);
                while (moveFrom.get() < moveTo.get()) {
                    moveStreams(hosts, moveFrom, moveFromLowWaterMark, moveTo, moveToHighWaterMark, rebalanceRateLimiter);
                    moveFrom.incrementAndGet();
                }
            }
        } finally {
            for (Host host : hosts) {
                host.close();
            }
        }
    }

    void moveStreams(List<Host> hosts,
                     AtomicInteger hostIdxMoveFrom,
                     int moveFromLowWaterMark,
                     AtomicInteger hostIdxMoveTo,
                     int moveToHighWaterMark,
                     Optional<RateLimiter> rateLimiter) {
        if (hostIdxMoveFrom.get() < 0 || hostIdxMoveFrom.get() >= hosts.size()
                || hostIdxMoveTo.get() < 0 || hostIdxMoveTo.get() >= hosts.size()
                || hostIdxMoveFrom.get() >= hostIdxMoveTo.get()) {
            return;
        }

        if (logger.isDebugEnabled()) {
            logger.debug("Moving streams : hosts = {}, from = {}, to = {} : from_low_water_mark = {}, to_high_water_mark = {}",
                         new Object[] { hosts, hostIdxMoveFrom.get(), hostIdxMoveTo.get(), moveFromLowWaterMark, moveToHighWaterMark });
        }

        Host hostMoveFrom = hosts.get(hostIdxMoveFrom.get());
        int numStreamsOnFromHost = hostMoveFrom.streams.size();
        if (numStreamsOnFromHost <= moveFromLowWaterMark) {
            // do nothing
            return;
        }

        int numStreamsToMove = numStreamsOnFromHost - moveFromLowWaterMark;
        LinkedList<String> streamsToMove = new LinkedList<String>(hostMoveFrom.streams);
        Collections.shuffle(streamsToMove);

        if (logger.isDebugEnabled()) {
            logger.debug("Try to move {} streams from host {} : streams = {}",
                         new Object[] { numStreamsToMove, hostMoveFrom.address, streamsToMove });
        }

        while (numStreamsToMove-- > 0 && !streamsToMove.isEmpty()) {
            if (rateLimiter.isPresent()) {
                rateLimiter.get().acquire();
            }

            // pick a host to move
            Host hostMoveTo = hosts.get(hostIdxMoveTo.get());
            while (hostMoveTo.streams.size() >= moveToHighWaterMark) {
                int hostIdx = hostIdxMoveTo.decrementAndGet();
                logger.info("move to host : {}, from {}", hostIdx, hostIdxMoveFrom.get());
                if (hostIdx <= hostIdxMoveFrom.get()) {
                    return;
                } else {
                    hostMoveTo = hosts.get(hostIdx);
                    if (logger.isDebugEnabled()) {
                        logger.debug("Target host to move moved to host {} @ {}",
                                hostIdx, hostMoveTo);
                    }
                }
            }

            // pick a stream
            String stream = streamsToMove.remove();

            // move the stream
            if (moveStream(stream, hostMoveFrom, hostMoveTo)) {
                hostMoveFrom.streams.remove(stream);
                hostMoveTo.streams.add(stream);
            }
        }

    }

    void moveRemainingStreamsFromSource(Host source,
                                        List<Host> hosts,
                                        Optional<RateLimiter> rateLimiter) {
        LinkedList<String> streamsToMove = new LinkedList<String>(source.streams);
        Collections.shuffle(streamsToMove);

        if (logger.isDebugEnabled()) {
            logger.debug("Try to move remaining streams from {} : {}", source, streamsToMove);
        }

        int hostIdx = hosts.size() - 1;

        while (!streamsToMove.isEmpty()) {
            if (rateLimiter.isPresent()) {
                rateLimiter.get().acquire();
            }

            Host target = hosts.get(hostIdx);
            if (!target.address.equals(source.address)) {
                String stream = streamsToMove.remove();
                // move the stream
                if (moveStream(stream, source, target)) {
                    source.streams.remove(stream);
                    target.streams.add(stream);
                }
            }
            --hostIdx;
            if (hostIdx < 0) {
                hostIdx = hosts.size() - 1;
            }
        }
    }

    private boolean moveStream(String stream, Host from, Host to) {
        try {
            doMoveStream(stream, from, to);
            return true;
        } catch (Exception e) {
            return false;
        }
    }

    private void doMoveStream(final String stream, final Host from, final Host to) throws Exception {
        logger.info("Moving stream {} from {} to {}.",
                    new Object[] { stream, from.address, to.address });
        Await.result(from.getClient().release(stream).flatMap(new Function<Void, Future<Void>>() {
            @Override
            public Future<Void> apply(Void result) {
                logger.info("Released stream {} from {}.", stream, from.address);
                return to.getMonitor().check(stream).addEventListener(new FutureEventListener<Void>() {

                    @Override
                    public void onSuccess(Void value) {
                        logger.info("Moved stream {} from {} to {}.",
                                    new Object[] { stream, from.address, to.address });
                    }

                    @Override
                    public void onFailure(Throwable cause) {
                        logger.info("Failed to move stream {} from {} to {} : ",
                                    new Object[] { stream, from.address, to.address, cause });
                    }
                });
            }
        }));
    }

    @Override
    public void close() {
        client.close();
    }
}