package ru.yandex.webmaster3.storage.util.fs.persistence;

import org.apache.commons.io.IOUtils;
import org.joda.time.Duration;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Required;

import java.io.Closeable;
import java.io.File;
import java.io.IOException;
import java.util.ArrayList;
import java.util.List;
import java.util.Timer;
import java.util.TimerTask;
import java.util.concurrent.ArrayBlockingQueue;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.function.Consumer;

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

    private String writeAheadLogDir;
    private AsyncLogWriter asyncWriter;

    private volatile Consumer<T> writer;

    private final int intervalMin;
    private final int batchSize;

    public CHBatchWriter(int intervalMin, int batchSize) {
        this.intervalMin = intervalMin;
        this.batchSize = batchSize;
    }

    public void init() throws IOException {
        WriteAheadLog wal = WriteAheadLog.createQuarantineIfNeeded(new File(writeAheadLogDir), 10000);
        asyncWriter = new AsyncLogWriter(wal);
        writer = this::logAsync;
    }

    public void destroy() throws IOException {
        asyncWriter.close();
    }

    public abstract void logSyncInternal(List<T> items);

    public void write(T entity) {
        writer.accept(entity);
    }

    private void logAsync(T entity) {
        try {
            asyncWriter.log(entity);
        } catch (IOException e) {
            log.error("PANIC! Async writer failed with exception, falling back to in-memory damper", e);
            initFallbackWriter();
            writer.accept(entity);
            IOUtils.closeQuietly(asyncWriter);
        }
    }

    private synchronized void initFallbackWriter() {
        FallbackWritesDamper fallbackWritesDamper = new FallbackWritesDamper(new ArrayBlockingQueue<>(batchSize * 2));
        writer = data -> fallbackWritesDamper.log(data);
    }

    private class AsyncLogWriter implements Closeable {
        private final WriteAheadLog wal;
        private final ScheduledExecutorService scheduler = Executors.newSingleThreadScheduledExecutor();
        private final AtomicBoolean closed = new AtomicBoolean(false);

        public AsyncLogWriter(WriteAheadLog wal) {
            this.wal = wal;
            scheduler.scheduleWithFixedDelay(new BatchWriter(wal), 0, intervalMin, TimeUnit.MINUTES);
        }

        public void log(T entity) throws IOException {
            if (closed.get()) {
                throw new IllegalStateException("Already closed");
            }
            byte[] msg = serialize(entity);
            wal.log(msg);
        }

        @Override
        public void close() throws IOException {
            if (closed.compareAndSet(false, true)) {
                scheduler.shutdown();
                wal.close();
            }
        }
    }

    private class BatchWriter implements Runnable {
        private final Logger log = LoggerFactory.getLogger(BatchWriter.class);

        private final WriteAheadLog wal;
        private final List<T> events = new ArrayList<>(batchSize);
        private long lastReqId = 0L;

        public BatchWriter(WriteAheadLog wal) {
            this.wal = wal;
        }

        @Override
        public void run() {
            try {
                do {
                    events.clear();
                    wal.replay(
                            (reqId, b, offset, len) -> {
                                events.add(deserialize(b, offset, len));
                                lastReqId = reqId;
                            },
                            batchSize);
                    if (!events.isEmpty()) {
                        logSyncInternal(events);
                        wal.releaseIncluding(lastReqId);
                        log.info("Dumped {} entities", events.size());
                    }
                    if (Thread.interrupted()) {
                        log.error("Interrupted");
                        return;
                    }
                } while (events.size() >= batchSize);
            } catch (IOException e) {
                log.error("Failed to parse from wal", e);
            } catch (Exception e) {
                log.error("Failed to write batch", e);
            }
        }
    }

    private class FallbackWritesDamper {
        private final BlockingQueue<T> queue;
        private final Timer timer;

        FallbackWritesDamper(BlockingQueue<T> queue) {
            this.queue = queue;
            timer = new Timer(true);
            long periodMs = Duration.standardMinutes(intervalMin).getMillis();
            timer.schedule(new FallbackAsyncWriter(queue), periodMs, periodMs);
        }

        public void log(T data) {
            if (!queue.offer(data)) {
                log.warn("Dropping message " + data);
            }
        }

    }

    private class FallbackAsyncWriter extends TimerTask {
        private final BlockingQueue<T> queue;

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

        @Override
        public void run() {
            int size = queue.size();
            if (size > batchSize / 2) {
                List<T> data = new ArrayList<>(batchSize);
                do {
                    queue.drainTo(data, batchSize);
                    logSyncInternal(data);
                } while (data.size() >= batchSize);
            } else {
                List<T> data = new ArrayList<>(size);
                queue.drainTo(data, size);
                logSyncInternal(data);
            }
        }
    }

    protected abstract byte[] serialize(T entity);

    protected abstract T deserialize(byte[] data, int offset, int len);

    @Required
    public void setWriteAheadLogDir(String writeAheadLogDir) {
        this.writeAheadLogDir = writeAheadLogDir;
    }
}
