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

import com.google.common.base.Stopwatch;
import com.google.common.collect.Iterators;
import com.google.common.collect.Range;
import com.google.common.io.CountingInputStream;
import org.apache.commons.io.IOUtils;
import org.apache.commons.lang3.time.DurationFormatUtils;
import org.joda.time.Duration;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import ru.yandex.webmaster3.core.util.RetryUtils;
import ru.yandex.webmaster3.storage.util.ProgressLogInputStream;

import java.io.BufferedInputStream;
import java.io.File;
import java.io.FileInputStream;
import java.io.IOException;
import java.io.InputStream;
import java.util.ArrayList;
import java.util.Collections;
import java.util.Iterator;
import java.util.List;
import java.util.concurrent.ArrayBlockingQueue;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.Executor;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.TimeUnit;
import java.util.function.Function;

/**
 * @author aherman
 */
public class AsyncTableReader<T> {
    private static final Logger log = LoggerFactory.getLogger(AsyncTableReader.class);

    private final YtCypressServiceImpl cypressService;
    private final YtPath path;
    private final YtTableRange range;
    private final YtTableReadDriver<T> tableReadDriver;

    private int retryCount = 1;
    private long partSize = 0L;
    private boolean isRandomOrderParts = false;
    private Executor executor;
    private boolean defaultExecutor;
    private int cacheQueueSize = 3;
    private String cacherThreadName = "yt-cacher";
    private YtTableRange readRange = YtTableRange.index(Range.closed(0L, 0L));
    private boolean needLock = true;

    public AsyncTableReader(YtCypressService cypressService, YtPath path, Range<Long> range,
                            YtTableReadDriver<T> tableReadDriver) {
        this(cypressService, path, YtTableRange.index(range), tableReadDriver);
    }

    public AsyncTableReader(YtCypressService cypressService, YtPath path, YtTableRange range,
                            YtTableReadDriver<T> tableReadDriver) {
        this.cypressService = (YtCypressServiceImpl) cypressService;
        this.path = path;
        this.range = range;
        this.tableReadDriver = tableReadDriver;
    }

    public AsyncTableReader(YtCypressService cypressService, YtPath path, Range<Long> range, YtRowMapper<T> mapper) {
        this(cypressService, path, YtTableRange.index(range), mapper);
    }

    public AsyncTableReader(YtCypressService cypressService, YtPath path, YtTableRange range, YtRowMapper<T> mapper) {
        this(cypressService, path, range, mapper, YtMissingValueMode.DEFAULT);
    }

    public AsyncTableReader(YtCypressService cypressService, YtPath path, Range<Long> range,
                            YtRowMapper<T> mapper, YtMissingValueMode missingValueMode) {
        this(cypressService, path, YtTableRange.index(range), mapper, missingValueMode);
    }

    public AsyncTableReader(YtCypressService cypressService, YtPath path, YtTableRange range,
                            YtRowMapper<T> mapper, YtMissingValueMode missingValueMode) {
        this(cypressService, path, range, YtTableReadDriver.createDSVDriver(mapper, missingValueMode));
    }

    public AsyncTableReader<T> withRetry(int retryCount) {
        this.retryCount = retryCount;
        return this;
    }

    public AsyncTableReader<T> splitInParts(long partSize) {
        this.partSize = partSize;
        return this;
    }

    public AsyncTableReader<T> splitInRandomOrderParts(long partSize) {
        this.partSize = partSize;
        this.isRandomOrderParts = true;
        return this;
    }

    public AsyncTableReader<T> inExecutor(Executor executor, String cacherThreadName) {
        this.executor = executor;
        this.cacherThreadName = cacherThreadName;
        return this;
    }

    public AsyncTableReader<T> withThreadName(String cacherThreadName) {
        this.cacherThreadName = cacherThreadName;
        return this;
    }

    public AsyncTableReader<T> needLock(boolean needLock) {
        this.needLock = needLock;
        return this;
    }

    public YtTableRange getReadRange() {
        return readRange;
    }

    public TableIterator<T> read() throws YtException {
        if (executor == null) {
            executor = Executors.newSingleThreadExecutor();
            defaultExecutor = true;
        }
        readRange = fixRange(range);
        if (needLock && cypressService.transaction != null && !cypressService.transaction.hasLock(path)) {
            cypressService.ytService.lock(cypressService.transaction, path, YtLockMode.SNAPSHOT);
        }
        BlockingQueue<YtReadAndSaveTableCommand.TableCache> queue = new ArrayBlockingQueue<>(cacheQueueSize);
        Iterator<YtTableRange> subranges;
        if (partSize == 0) {
            subranges = Iterators.forArray(readRange);
        } else {
            if (isRandomOrderParts) {
                subranges = generateRandomOrderSubranges(readRange, partSize);
            } else {
                subranges = generateSubranges(readRange, partSize);
            }
        }
        Function<YtTableRange, File> cacheFileSupplier = r -> {
            String cacheFileName = "download_" + path.getName() + "_"
                    + r.toString().replace("/", "_") + "_" + System.currentTimeMillis();
            File cacheFolder = cypressService.ytService.getCacheFolder(path);
            return new File(cacheFolder, cacheFileName);
        };

        OutputFormat outputFormat = tableReadDriver.getFormat();

        CacheQueue<YtReadAndSaveTableCommand.TableCache> cacheQueue = new CacheQueue<>(queue);
        Runnable cacheRunnable = new YtCacheRunnable(
                cypressService,
                path,
                subranges,
                cacheQueue,
                cacheFileSupplier,
                retryCount,
                outputFormat,
                cacherThreadName
        );

        executor.execute(cacheRunnable);
        if (defaultExecutor && executor instanceof ExecutorService) {
            ((ExecutorService) executor).shutdown();
        }
        TableIteratorImpl<T> tableIterator = new TableIteratorImpl<>(path, cacheQueue, tableReadDriver);
        tableIterator.setFullRange(readRange);
        return tableIterator;
    }

    private YtTableRange fixRange(YtTableRange range) throws YtException {
        if (range instanceof YtTableRange.IndexRange) {
            Range<Long> indexRange = ((YtTableRange.IndexRange) range).getRange();
            long upperBound;
            if (indexRange.hasUpperBound()) {
                upperBound = indexRange.upperEndpoint();
            } else {
                YtNode node = cypressService.getNode(path);
                YtTable table = (YtTable) node;
                upperBound = table.getRowCount();
            }
            if (indexRange.hasLowerBound()) {
                indexRange = Range.closed(indexRange.lowerEndpoint(), upperBound);
            } else {
                indexRange = Range.closed(0L, upperBound);
            }
            return YtTableRange.index(indexRange);
        } else {
            return range;
        }
    }

    private static class CacheQueue<T> {
        private final BlockingQueue<T> queue;

        private volatile boolean readerDied = false;
        private volatile boolean readerStopped = false;

        private CacheQueue(BlockingQueue<T> queue) {
            this.queue = queue;
        }
    }

    public interface TableIterator<T> extends InterruptableIterator<T> {
        YtTableRange getFullRange();

        YtTableRange getCurrentRange();
    }

    private static class TableIteratorImpl<T> implements TableIterator<T> {
        private final YtPath path;
        private final CacheQueue<YtReadAndSaveTableCommand.TableCache> cacheQueue;
        private final YtTableReadDriver<T> tableReadDriver;
        private final Stopwatch readerStopwatch = Stopwatch.createUnstarted();

        private long row = 0;
        private YtTableRange fullRange;

        private InterruptableIterator<T> parser;
        private YtReadAndSaveTableCommand.TableCache currentTablePart;
        private InputStream is;
        private boolean stopped = false;

        public TableIteratorImpl(YtPath path, CacheQueue<YtReadAndSaveTableCommand.TableCache> cacheQueue,
                                 YtTableReadDriver<T> tableReadDriver) {
            this.path = path;
            this.cacheQueue = cacheQueue;
            this.tableReadDriver = tableReadDriver;
        }

        @Override
        public YtTableRange getFullRange() {
            return fullRange;
        }

        @Override
        public YtTableRange getCurrentRange() {
            return currentTablePart.getRange();
        }

        @Override
        public long getRow() {
            return row;
        }

        void setFullRange(YtTableRange fullRange) {
            this.fullRange = fullRange;
        }

        @Override
        public boolean hasNext() throws InterruptedException, IOException, YtException {
            if (stopped) {
                return false;
            }

            if (Thread.interrupted()) {
                throw new InterruptedException("Table iterator interrupted");
            }

            if (parser != null && parser.hasNext()) {
                return true;
            }

            log.debug("Close streamParser");
            IOUtils.closeQuietly(is);

            if (currentTablePart != null) {
                log.info("Stop processing: {} {}",
                        path.toYtPath(currentTablePart.getRange()),
                        DurationFormatUtils.formatDurationHMS(readerStopwatch.elapsed(TimeUnit.MILLISECONDS))
                );

                if (currentTablePart.getFile().exists()) {
                    currentTablePart.getFile().delete();
                }
            }

            YtReadAndSaveTableCommand.TableCache nextPart;
            while (true) {
                nextPart = cacheQueue.queue.poll(10, TimeUnit.MILLISECONDS);
                if (cacheQueue.readerDied) {
                    throw new YtException("Reader died, unable to read table: " + path.toYtPath(fullRange));
                }
                if (nextPart != null) {
                    break;
                }

                if (cacheQueue.readerStopped && cacheQueue.queue.isEmpty()) {
                    stopped = true;
                    return false;
                }
            }
            currentTablePart = nextPart;

            if (readerStopwatch.isRunning()) {
                readerStopwatch.stop();
                readerStopwatch.reset();
            }
            readerStopwatch.start();
            log.debug("Start processing: {}", path.toYtPath(currentTablePart.getRange()));
            is = new BufferedInputStream(new FileInputStream(currentTablePart.getFile()));
            InputStream uncompressedIs = is;
            if (currentTablePart.isCompressed()) {
                uncompressedIs = new YtSnappyCompressedStream(is);
            }
            CountingInputStream cis = new CountingInputStream(uncompressedIs);
            ProgressLogInputStream pis = new ProgressLogInputStream(cis,
                    "Read " + path.getName()
                            + "["
                            + currentTablePart.getRange().renderInPath().orElse("")
                            + "]"
            );
            log.debug("Open streamParser");
            parser = tableReadDriver.createParser(pis);
            return parser.hasNext();
        }

        @Override
        public T next() throws InterruptedException, IOException, YtException {
            row++;
            return parser.next();
        }

        @Override
        public void close() throws IOException {
            log.debug("Close iterator");
            IOUtils.closeQuietly(parser);
            IOUtils.closeQuietly(is);
        }
    }

    private static class YtCacheRunnable implements Runnable {
        private final YtCypressServiceImpl cypressService;
        private final YtPath path;
        private final Iterator<YtTableRange> subranges;
        private final CacheQueue<YtReadAndSaveTableCommand.TableCache> cacheQueue;
        private final Function<YtTableRange, File> cacheFileSupplier;
        private final int retryCount;
        private final OutputFormat outputFormat;
        private final String cacherThreadName;

        public YtCacheRunnable(YtCypressServiceImpl cypressService, YtPath path, Iterator<YtTableRange> subranges,
                               CacheQueue<YtReadAndSaveTableCommand.TableCache> cacheQueue,
                               Function<YtTableRange, File> cacheFileSupplier, int retryCount, OutputFormat outputFormat,
                               String cacherThreadName) {
            this.cypressService = cypressService;
            this.path = path;
            this.subranges = subranges;
            this.cacheQueue = cacheQueue;
            this.cacheFileSupplier = cacheFileSupplier;
            this.retryCount = retryCount;
            this.outputFormat = outputFormat;
            this.cacherThreadName = cacherThreadName;
        }

        @Override
        public void run() {
            String previousName = Thread.currentThread().getName();
            Thread.currentThread().setName(cacherThreadName);
            log.info("Start async table cacher");
            try {
                while (subranges.hasNext() && !Thread.interrupted()) {
                    YtTableRange subrange = subranges.next();
                    Stopwatch downloadTime = Stopwatch.createStarted();
                    File cacheFile = cacheFileSupplier.apply(subrange);
                    YtReadAndSaveTableCommand.TableCache tableCache;
                    try {
                        tableCache = RetryUtils
                                .query(RetryUtils.expBackoff(retryCount, Duration.standardSeconds(15)), () -> {
                                    return cypressService
                                            .readAndSaveTable(path, subrange, true, outputFormat,
                                                    cacheFile);
                                });
                    } catch (InterruptedException e) {
                        cacheQueue.readerDied = true;
                        break;
                    }

                    downloadTime.stop();
                    log.info("Stop caching: {} {}",
                            path.toYtPath(subrange),
                            DurationFormatUtils.formatDurationHMS(downloadTime.elapsed(TimeUnit.MILLISECONDS))
                    );
                    if (!tableCache.isEmpty()) {
                        cacheQueue.queue.put(tableCache);
                    }
                }
            } catch (InterruptedException e) {
                cacheQueue.readerDied = true;
                log.debug("Interrupted");
            } catch (Exception e) {
                cacheQueue.readerDied = true;
                log.error("Unhandled exception in cacher", e);
            } catch (Throwable e) {
                cacheQueue.readerDied = true;
                log.error("Unhandled error in cacher", e);
                throw e;
            }
            cacheQueue.readerStopped = true;
            log.info("Stop cacher");
            Thread.currentThread().setName(previousName);
        }
    }

    private static Iterator<YtTableRange> generateSubranges(YtTableRange range, long subrangeSize) {
        if (range instanceof YtTableRange.IndexRange) {
            Range<Long> indexRange = ((YtTableRange.IndexRange) range).getRange();
            return new Iterator<YtTableRange>() {
                long nextStart = indexRange.lowerEndpoint();

                @Override
                public boolean hasNext() {
                    return nextStart < indexRange.upperEndpoint();
                }

                @Override
                public YtTableRange next() {
                    long nextEnd = Math.min(nextStart + subrangeSize, indexRange.upperEndpoint());
                    Range<Long> result = Range.closedOpen(nextStart, nextEnd);
                    nextStart = nextEnd;
                    return YtTableRange.index(result);
                }
            };
        } else {
            return Iterators.forArray(range);
        }
    }

    private static Iterator<YtTableRange> generateRandomOrderSubranges(YtTableRange range, long subrangeSize) {
        List<YtTableRange> ranges = new ArrayList<>();
        generateSubranges(range, subrangeSize)
                .forEachRemaining(ranges::add);

        Collections.shuffle(ranges);

        return ranges.iterator();
    }
}
