package ru.yandex.webmaster3.storage.util.yt;

import java.util.ArrayList;
import java.util.Comparator;
import java.util.List;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.PriorityBlockingQueue;
import java.util.concurrent.ThreadFactory;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;

import com.google.common.util.concurrent.ThreadFactoryBuilder;
import org.apache.commons.lang3.tuple.Pair;
import org.joda.time.Duration;
import org.joda.time.Instant;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 * @author aherman
 */
public class YtTransactionService {
    private static final Logger log = LoggerFactory.getLogger(YtTransactionService.class);
    private static final int TRANSACTION_SHARDS_COUNT = 10;

    private static final int MAX_RETRY_COUNT = 5;
    private static final int SLEEP_TIME = (int) TimeUnit.SECONDS.toMillis(10);
    private static final int MINIMUM_PING_INTERVAL = SLEEP_TIME * 3;

    private final YtService ytService;

    private final PriorityBlockingQueue<YtTransaction>[] transactions = new PriorityBlockingQueue[TRANSACTION_SHARDS_COUNT];

    private AtomicInteger transactionCellNumber = new AtomicInteger(0);

    private ExecutorService pingExecutorService;

    YtTransactionService(YtService ytService) {
        this.ytService = ytService;
        for (int i = 0; i < TRANSACTION_SHARDS_COUNT; i++) {
            transactions[i] = new PriorityBlockingQueue<>(16, Comparator.comparing(YtTransaction::getNextPingTime));
        }
    }

    void start() {
        ThreadFactory threadFactory = new ThreadFactoryBuilder()
                .setDaemon(true)
                .setNameFormat("yt-pinger-%d")
                .build();
        pingExecutorService = Executors.newFixedThreadPool(TRANSACTION_SHARDS_COUNT, threadFactory);
        for (int i = 0; i < TRANSACTION_SHARDS_COUNT; i++) {
            pingExecutorService.submit(new YtTransactionService.PingThread(transactions[i]));
        }
    }

    void stop() {
        pingExecutorService.shutdown();
        try {
            pingExecutorService.awaitTermination(5000, TimeUnit.MILLISECONDS);
        } catch (InterruptedException e) {
            log.error("Unable to stop pinger", e);
        }
    }

    private YtTransaction registerTransaction(String parentId, String cluster, String id, int timeoutMs) {
        Instant now = Instant.now();
        int pingIntervalMs = Math.max(timeoutMs / 10, MINIMUM_PING_INTERVAL);
        YtTransaction transaction = new YtTransaction(parentId, cluster, id, Duration.millis(timeoutMs), pingIntervalMs);
        transaction.ping(now, getNextPingTime(now, transaction));
        final int increment = transactionCellNumber.getAndIncrement();
        transactions[Math.abs(increment % 5)].put(transaction);
        return transaction;
    }

    private void unregisterTransaction(YtTransaction transaction) {
        for (int i = 0; i < TRANSACTION_SHARDS_COUNT; i++) {
            transactions[i].remove(transaction);
        }
    }

    TransactionBuilder newTransactionBuilder(String transactionId, String cluster) {
        return new TransactionBuilder(transactionId, cluster);
    }

    private Instant getNextPingTime(Instant now, YtTransaction transaction) {
        return now.plus(transaction.getPingIntervalMs());
    }

    private class PingThread implements Runnable {
        private final PriorityBlockingQueue<YtTransaction> priorityBlockingQueue;

        private PingThread(PriorityBlockingQueue<YtTransaction> priorityBlockingQueue) {
            this.priorityBlockingQueue = priorityBlockingQueue;
        }

        @Override
        public void run() {
            log.info("Pinger started");
            while (!Thread.interrupted()) {
                Instant now = Instant.now();
                while (true) {
                    YtTransaction headTransaction = priorityBlockingQueue.peek();
                    if (headTransaction == null || headTransaction.getNextPingTime().isAfter(now)) {
                        break;
                    }
                    if (headTransaction.isClosed()) {
                        priorityBlockingQueue.remove(headTransaction);
                        continue;
                    }
                    try {
                        if (headTransaction.getPingRetryCount().incrementAndGet() > MAX_RETRY_COUNT) {
                            log.error("Max retry reached: {}", headTransaction);
                            headTransaction.close();
                            continue;
                        }
                        final YtResult<Void> result = ytService.pingTxInfo(headTransaction);
                        if (result.getStatus() == YtStatus.YT_404_UNKNOWN_ACTION) {
                            log.trace("Invalid status: {}, Error: {}", result.getStatus(), result.getError());
                            headTransaction.close();
                            continue;
                        } else if (result.isError()) {
                            log.error("Invalid status: {}, Error: {}", result.getStatus(), result.getError());
                        }

                        updateNextPingTime(now, headTransaction);
                    } catch (YtException e) {
                        log.error("Unable to ping transaction: {}", headTransaction, e);
                        headTransaction.getPingRetryCount().incrementAndGet();
                    }
                }
                try {
                    Thread.sleep(SLEEP_TIME);
                } catch (InterruptedException e) {
                    log.warn("Pinger interrupted");
                    break;
                }
            }
            log.info("Pinger stopped");
        }

        private void updateNextPingTime(Instant now, YtTransaction headTransaction) {
            priorityBlockingQueue.remove(headTransaction);
            Instant nextPingTime = getNextPingTime(now, headTransaction);
            headTransaction.ping(now, nextPingTime);
            priorityBlockingQueue.put(headTransaction);
        }
    }

    public class TransactionBuilder {
        private final String parentTransactionId;
        private final String cluster;

        private final List<Pair<YtPath, YtLockMode>> lockPath = new ArrayList<>();
        private int timeoutMs = (int) MINIMUM_PING_INTERVAL * 2;

        public TransactionBuilder(String parentTransactionId, String cluster) {
            this.parentTransactionId = parentTransactionId;
            this.cluster = cluster;
        }

        public TransactionBuilder withLock(YtPath path, YtLockMode mode) {
            lockPath.add(Pair.of(path, mode));
            return this;
        }

        public TransactionBuilder withTimeout(int timeout, TimeUnit timeUnit) {
            if (timeUnit == TimeUnit.NANOSECONDS || timeUnit == TimeUnit.MICROSECONDS) {
                throw new IllegalArgumentException("Time unit must be greater or equal than millisecond");
            }
            if (timeUnit.toHours(timeout) > 6) {
                throw new IllegalArgumentException("Timeout is too long");
            }
            this.timeoutMs = (int) timeUnit.toMillis(timeout);

            return this;
        }

        private <T> T execute(TransactionProcess process, TransactionProcessWithResult<T> processWithResult) throws YtException {
            String id = ytService.startTx(cluster, timeoutMs, parentTransactionId);
            YtTransaction transaction = registerTransaction(parentTransactionId, cluster, id, timeoutMs);
            try {
                for (Pair<YtPath, YtLockMode> lock : lockPath) {
                    ytService.lock(transaction, lock.getLeft(), lock.getRight());
                }
            } catch (Exception exp) {
                unregisterTransaction(transaction);
                throw new YtException("Exception in transaction: " + transaction, exp);
            }
            try {

                YtCypressService cs = new YtCypressServiceImpl(ytService, transaction);
                if (process != null) {
                    if (process.run(cs)) {
                        ytService.commitTx(transaction);
                    } else {
                        ytService.abortTx(transaction);
                    }
                    return null;
                } else {
                    T result = processWithResult.run(cs);
                    ytService.commitTx(transaction);
                    return result;
                }
            } catch (InterruptedException e) {
                log.error("Transaction interrupted: " + transaction, e);
                throw new YtException("Transaction interrupted: " + transaction, e);
            } catch (Exception e) {
                log.error("Exception in transaction, abort: " + transaction, e);
                try {
                    ytService.abortTx(transaction);
                } catch (YtException e1) {
                    log.error("Unable to abort transaction " + transaction, e1);
                    throw new YtException("Unable to abort transaction: " + transaction, e);
                }
                throw new YtException("Exception in transaction: " + transaction, e);
            } finally {
                unregisterTransaction(transaction);
            }
        }

        public void execute(TransactionProcess process) throws YtException {
            execute(process, null);
        }

        public <T> T query(TransactionProcessWithResult<T> process) throws YtException {
            return execute(null, process);
        }
    }

    public interface TransactionProcess {
        boolean run(YtCypressService cypressService) throws YtException, InterruptedException;
    }

    public interface TransactionProcessWithResult<T> {
        T run(YtCypressService cypressService) throws YtException, InterruptedException;
    }
}
