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

import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.CancellationException;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.ThreadLocalRandom;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicLong;
import java.util.concurrent.locks.ReentrantLock;

import lombok.Getter;
import org.apache.commons.lang3.tuple.Pair;
import org.joda.time.Duration;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import ru.yandex.kikimr.persqueue.LogbrokerClientAsyncFactory;
import ru.yandex.kikimr.persqueue.auth.Credentials;
import ru.yandex.kikimr.persqueue.compression.CompressionCodec;
import ru.yandex.kikimr.persqueue.producer.AsyncProducer;
import ru.yandex.kikimr.persqueue.producer.async.AsyncProducerConfig;
import ru.yandex.kikimr.persqueue.producer.transport.message.inbound.ProducerInitResponse;
import ru.yandex.kikimr.persqueue.producer.transport.message.inbound.ProducerWriteResponse;
import ru.yandex.kikimr.persqueue.proxy.ProxyBalancer;
import ru.yandex.webmaster3.core.metrics.externals.AbstractExternalAPIService;
import ru.yandex.webmaster3.core.metrics.externals.ExternalDependencyMethod;
import ru.yandex.webmaster3.core.security.tvm.TVMTokenService;
import ru.yandex.webmaster3.core.util.JavaMethodWitness;

/**
 * @author avhaliullin
 */
public class LogbrokerClient extends AbstractExternalAPIService implements LogbrokerWriter {
    private static final Duration DEFAULT_WAIT_TIMEOUT = Duration.standardSeconds(20);
    private static final Logger log = LoggerFactory.getLogger(LogbrokerClient.class);

    private final TVMTokenService logBrokerTvmTokenService;
    private final String topic;
    private final String source;

    private LogbrokerClientAsyncFactory factory;
    private AsyncProducer producer;

    @Getter
    private int partitionNumber;

    private final ReentrantLock initLock = new ReentrantLock();
    private long writeSeqNo = 1;
    private final AtomicLong initCount = new AtomicLong(1);

    public LogbrokerClient(TVMTokenService logBrokerTvmTokenService, String lbHost, int lbPort, String topic, String hostName) {
        this.logBrokerTvmTokenService = logBrokerTvmTokenService;
        this.factory = new LogbrokerClientAsyncFactory(new ProxyBalancer(lbHost, lbPort));
        this.topic = topic;
        this.source = (hostName + "-" + ThreadLocalRandom.current().nextLong());
    }

    public LogbrokerClient(TVMTokenService logBrokerTvmTokenService, LogbrokerClientAsyncFactory factory, String topic, String source) {
        this.logBrokerTvmTokenService = logBrokerTvmTokenService;
        this.factory = factory;
        this.topic = topic;
        this.source = source;
    }

    public void init() {
        initProducer(initCount.get());
    }

    private void initProducer(long expectedInitCount) {
        initLock.lock();
        try {
            long currentInitCount = initCount.get();
            if (expectedInitCount < currentInitCount) {
                // уже инициализировали
                log.info("Producer already initialized, expected count: {}, actual count: {}",
                        expectedInitCount, currentInitCount);
                return;
            }

            try {
                if (producer != null) {
                    producer.close();
                }
            } catch (Exception e) {
                log.error("Exception during closing lb producer", e);
            } finally {
                producer = null;
            }

            while (!doInitProducerLocked()) {
                log.error("Sleeping before next producer init attempt");
                try {
                    Thread.sleep(3000);
                } catch (InterruptedException e) {
                    // ignore
                }
            }

            initCount.incrementAndGet();
        } finally {
            initLock.unlock();
        }
    }

    private boolean doInitProducerLocked() {
        AsyncProducerConfig config = AsyncProducerConfig.builder(topic, source.getBytes())
                .setCodec(CompressionCodec.GZIP)
                .setCredentialsProvider(() -> Credentials.tvm(logBrokerTvmTokenService.getToken()))
                .build();

        CompletableFuture<Pair<AsyncProducer, ProducerInitResponse>> initFuture = factory.asyncProducer(config)
                .thenCompose(c -> c.init().thenApply(response -> Pair.of(c, response)));

        Pair<AsyncProducer, ProducerInitResponse> producerAndResponse = null;
        try {
            log.info("Initializing producer topic: {}, source: {}, partitionNumber: {}", topic, source, partitionNumber);
            producerAndResponse = initFuture.get();
        } catch (CancellationException | ExecutionException | InterruptedException e) {
            log.error("LB producer init failed", e);
        }

        if (producerAndResponse != null) {
            producer = producerAndResponse.getLeft();
            writeSeqNo = producerAndResponse.getRight().getMaxSeqNo() + 1;
            partitionNumber = producerAndResponse.getRight().getPartition();
            log.info("Producer initialized topic: {}, source: {}, partitionNumber: {}", topic, source, partitionNumber);
            return true;
        } else {
            log.error("Failed to init producer");
            return false;
        }
    }

    private CompletableFuture<ProducerWriteResponse> writeData(byte[] data) {
        initLock.lock();
        try {
            // Эту операцию нужно делать под глобальным локом, чтобы:
            // 1) каждый вызов producer.write получал возрастающие writeSeqNo
            // 2) не было конфликта с инициализацией writeSeqNo в initProducer
            return producer.write(data, writeSeqNo++);
        } finally {
            initLock.unlock();
        }
    }

    private List<CompletableFuture<ProducerWriteResponse>> writeData(List<byte[]> data) {
        initLock.lock();
        try {
            List<CompletableFuture<ProducerWriteResponse>> result = new ArrayList<>();
            for (byte[] item : data) {
                result.add(producer.write(item, writeSeqNo++));
            }
            return result;
        } finally {
            initLock.unlock();
        }
    }


    @ExternalDependencyMethod("writesync")
    @Override
    public void write(byte[] data) throws Exception {
        write(data, DEFAULT_WAIT_TIMEOUT);
    }

    @ExternalDependencyMethod("writesync")
    @Override
    public void write(List<byte[]> data) throws Exception {
        write(data, DEFAULT_WAIT_TIMEOUT);
    }

    @ExternalDependencyMethod("writesync")
    public void write(List<byte[]> items, Duration timeout) throws Exception {
        trackExecution(new JavaMethodWitness() {
        }, ALL_ERRORS_INTERNAL, () -> {
            try {
                List<CompletableFuture<ProducerWriteResponse>> respFutures = writeData(items);

                for (var respFuture : respFutures) {
                    ProducerWriteResponse response = respFuture.get(timeout.getMillis(), TimeUnit.MILLISECONDS);
                    if (response.isAlreadyWritten()) {
                        throw new RuntimeException("LB reported message already written");
                    }
                }
            } catch (ExecutionException e) {
                log.error("Restarting producer due to error", e);
                initProducer(initCount.get());
                throw e;
            }
        });
    }

    @ExternalDependencyMethod("writesync")
    public void write(byte[] data, Duration timeout) throws Exception {
        trackExecution(new JavaMethodWitness() {
        }, ALL_ERRORS_INTERNAL, () -> {
            CompletableFuture<ProducerWriteResponse> respFuture = writeData(data);
            try {
                ProducerWriteResponse response = respFuture.get(timeout.getMillis(), TimeUnit.MILLISECONDS);
                if (response.isAlreadyWritten()) {
                    throw new RuntimeException("LB reported message already written");
                }
            } catch (ExecutionException e) {
                log.error("Restarting producer due to error", e);
                initProducer(initCount.get());

                throw e;
            }
        });
    }

    public void close() {
        producer.close();
        log.info("Producer closed topic: {},source:{}, partitionNumber: {}", topic, source, partitionNumber);
    }
}
