package ru.yandex.webmaster3.core.logbroker.writer;

import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ThreadLocalRandom;
import java.util.stream.Collectors;
import java.util.stream.IntStream;

import ru.yandex.kikimr.persqueue.LogbrokerClientAsyncFactory;
import ru.yandex.kikimr.persqueue.proxy.ProxyBalancer;
import ru.yandex.webmaster3.core.security.tvm.TVMTokenService;

/**
 * ishalaru
 * 28.02.2020
 * Данный класс пересылает сообщения в разные топики, при этом гарании что при вызове из разных потоков будет так же нет, в целом это лучше блокировок.
 **/
public class LogbrokerRoundRobinParitionClient implements LogbrokerWriter {
    private int partitionCount = 1;
    private Map<Integer, LogbrokerClient> clients;
    private final TVMTokenService logBrokerTvmTokenService;
    private final String topic;
    private final String hostName;
    private final LogbrokerClientAsyncFactory factory;
    private int lastPartition = 0;


    public LogbrokerRoundRobinParitionClient(TVMTokenService logBrokerTvmTokenService,
                                             String topic,
                                             String hostName,
                                             String lbHost,
                                             int lbPort,
                                             int partitionCount) {
        this.logBrokerTvmTokenService = logBrokerTvmTokenService;
        this.topic = topic;
        this.hostName = hostName;
        this.factory = new LogbrokerClientAsyncFactory(new ProxyBalancer(lbHost, lbPort));
        this.partitionCount = partitionCount;
    }

    public void init() {
        clients = new HashMap<>();
        Set<Integer> neededPartitions = IntStream.range(0, partitionCount).boxed().collect(Collectors.toSet());
        long value = ThreadLocalRandom.current().nextLong();
        while (!neededPartitions.isEmpty()) {
            String source = hostName + "-" + value;
            final LogbrokerClient logbrokerClient =
                    new LogbrokerClient(logBrokerTvmTokenService, factory, topic, source);
            logbrokerClient.init();
            if (neededPartitions.contains(logbrokerClient.getPartitionNumber())) {
                logbrokerClient.setBeanName("logbrokerRoundRobinPartitionClient-" + logbrokerClient.getPartitionNumber());
                clients.put(logbrokerClient.getPartitionNumber(), logbrokerClient);
                neededPartitions.remove(logbrokerClient.getPartitionNumber());
            } else {
                logbrokerClient.close();
            }
            value++;
        }
    }


    public void write(byte[] data) throws Exception {
        clients.get(lastPartition).write(data);
        lastPartition = (lastPartition + 1) % partitionCount;
    }

    public void write(List<byte[]> data) throws Exception {
        clients.get(lastPartition).write(data);
        lastPartition = (lastPartition + 1) % partitionCount;
    }
}
