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

import com.datastax.driver.core.utils.UUIDs;
import org.apache.commons.io.IOUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.BufferedReader;
import java.io.Closeable;
import java.io.File;
import java.io.FileInputStream;
import java.io.FileOutputStream;
import java.io.FileReader;
import java.io.IOException;
import java.io.InputStream;
import java.io.InputStreamReader;
import java.io.OutputStream;
import java.nio.ByteBuffer;
import java.nio.IntBuffer;
import java.nio.file.Files;
import java.nio.file.StandardCopyOption;
import java.util.HashSet;
import java.util.Map;
import java.util.Set;
import java.util.SortedMap;
import java.util.TreeMap;
import java.util.UUID;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.locks.Lock;
import java.util.concurrent.locks.ReentrantLock;

import ru.yandex.webmaster3.core.util.MapUtils;

/**
 * @author avhaliullin
 */
public class WriteAheadLog implements Closeable {
    private static final Logger log = LoggerFactory.getLogger(WriteAheadLog.class);

    static final int MAX_CONTROL_SIZE_BYTES = 1024;
    private static final String CONTROL_NAME = "control";

    private static final String INDEX_NAME = "index";
    private static final String PART_NAME_PREFIX = "part-";
    private final byte[] lengthWriteByteBuff = new byte[4];

    private final IntBuffer lengthWriteIntBuff = ByteBuffer.wrap(lengthWriteByteBuff).asIntBuffer();

    private final File workDir;
    private final int minPartSize;

    private final ScheduledExecutorService scheduler = Executors.newSingleThreadScheduledExecutor();

    private final Lock controlLock = new ReentrantLock();
    private final Lock indexLock = new ReentrantLock();
    private final Lock dataLock = new ReentrantLock();

    private OutputStream controlWriter;
    private OutputStream indexWriter;
    private OutputStream dataWriter;

    private UUID currentPartId;
    private TreeMap<Long, UUID> index = new TreeMap<>();
    private long checkPoint;
    private long recordId;
    private volatile boolean closed = false;

    public WriteAheadLog(File workDir, int minPartSize) throws IOException {
        this(workDir, minPartSize, true);
    }

    /**
     * For testing
     */
    WriteAheadLog(File workDir, int minPartSize, boolean useMaintenance) throws IOException {
        this.minPartSize = minPartSize;
        if (!workDir.exists()) {
            if (!workDir.mkdirs()) {
                throw new IOException("Failed to create working dir " + workDir.getAbsolutePath());
            }
        }
        this.workDir = workDir;
        File controlFile = new File(workDir, CONTROL_NAME);
        File indexFile = new File(workDir, INDEX_NAME);

        if (controlFile.exists()) {
            this.checkPoint = readCheckPoint(controlFile);
        } else {
            this.checkPoint = 0L;
        }
        this.controlWriter = new FileOutputStream(controlFile, true);

        if (indexFile.exists()) {
            fillIndex(indexFile, index);
        }
        this.indexWriter = new FileOutputStream(indexFile, true);

        long partOffset;
        if (index.isEmpty()) {
            this.currentPartId = UUIDs.timeBased();
            partOffset = 0L;
            addNewPart(partOffset, currentPartId);
        } else {
            Map.Entry<Long, UUID> entry = index.lastEntry();
            this.currentPartId = entry.getValue();
            partOffset = entry.getKey();
        }

        File dataFile = new File(workDir, PART_NAME_PREFIX + currentPartId);
        if (dataFile.exists()) {
            recordId = partOffset + readData(dataFile, partOffset, null, 0L, -1, true);
        } else {
            recordId = partOffset;
        }
        this.dataWriter = new FileOutputStream(dataFile, true);
        if (useMaintenance) {
            scheduler.scheduleAtFixedRate(new MaintenanceTask(), 0, 30, TimeUnit.SECONDS);
        }
    }

    public long log(byte[] record) throws IOException {
        assertClosed();
        long start = 0L;
        if (log.isDebugEnabled()) {
            start = System.nanoTime();
        }
        dataLock.lock();
        try {
            lengthWriteIntBuff.put(0, record.length);
            dataWriter.write(lengthWriteByteBuff);
            dataWriter.write(record, 0, record.length);
            dataWriter.flush();
            if (log.isDebugEnabled()) {
                log.debug("WA logging took {}micros", (System.nanoTime() - start) / 1000);
            }
            return recordId++;
        } finally {
            dataLock.unlock();
        }
    }

    public void releaseExcluding(long before) throws IOException {
        assertClosed();
        controlLock.lock();
        try {
            updateControlRecord(controlWriter, before);
            checkPoint = before;
        } finally {
            controlLock.unlock();
        }
    }

    public void releaseIncluding(long beforeIncluding) throws IOException {
        releaseExcluding(beforeIncluding + 1);
    }

    public int replay(WALRecordConsumer consumer, int limit) throws IOException {
        assertClosed();
        if (limit < 0) {
            throw new IllegalArgumentException("Limit should be > 0, actual " + limit);
        }
        int processed = 0;
        long checkPoint;
        controlLock.lock();
        try {
            checkPoint = this.checkPoint;
        } finally {
            controlLock.unlock();
        }
        indexLock.lock();
        try {
            SortedMap<Long, UUID> indexes = getIndexFromOffset(checkPoint);
            for (Map.Entry<Long, UUID> entry : indexes.entrySet()) {
                boolean strictConsistency = !currentPartId.equals(entry.getValue());
                processed += readData(new File(workDir, PART_NAME_PREFIX + entry.getValue().toString()), entry.getKey(),
                        consumer, checkPoint, limit - processed, strictConsistency);
                if (processed >= limit) {
                    return processed;
                }
            }
        } finally {
            indexLock.unlock();
        }
        return processed;
    }

    /**
     * Use ONLY from maintenance thread and for unit tests
     */
    void maintenance() {
        long start = System.nanoTime();
        try {
            // ротируем кусок лога
            long lastPartOffset = index.lastKey();
            if (recordId - lastPartOffset >= minPartSize) {
                UUID newPartId = UUIDs.timeBased();
                log.info("Creating new part {}", newPartId);
                File file = new File(workDir, PART_NAME_PREFIX + newPartId.toString());

                FileOutputStream out = new FileOutputStream(file);

                Closeable closeOnFinish = out;
                indexLock.lock();
                dataLock.lock();
                try {
                    long curOffset = recordId;
                    addNewPart(curOffset, newPartId);
                    closeOnFinish = dataWriter;
                    dataWriter = out;
                    currentPartId = newPartId;
                } finally {
                    dataLock.unlock();
                    indexLock.unlock();
                    IOUtils.closeQuietly(closeOnFinish);
                }
                log.info("Log part file switched");
            }

        } catch (Exception e) {
            log.error("Failed to create new part file", e);
        }

        try {
            // перезаписываем control (по факту всегда используется только последний чекпоинт)
            if (new File(workDir, CONTROL_NAME).length() > MAX_CONTROL_SIZE_BYTES) {
                log.info("Optimizing control file");
                File tmpControl = new File(workDir, CONTROL_NAME + ".tmp");
                OutputStream newControlWriter = new FileOutputStream(tmpControl);
                Closeable closeOnFinish = newControlWriter;
                controlLock.lock();
                try {
                    updateControlRecord(newControlWriter, checkPoint);
                    Files.move(
                            tmpControl.toPath(),
                            new File(workDir, CONTROL_NAME).toPath(),
                            StandardCopyOption.REPLACE_EXISTING, StandardCopyOption.ATOMIC_MOVE
                    );
                    closeOnFinish = controlWriter;
                    controlWriter = newControlWriter;
                } finally {
                    controlLock.unlock();
                    IOUtils.closeQuietly(closeOnFinish);
                }
                log.info("Control file rewritten");
            }
        } catch (Exception e) {
            log.error("Failed to rewrite control file", e);
        }

        try {
            // Удаляем лишние чанки из индексного файла
            SortedMap<Long, UUID> indexRecordsToLeave = getIndexFromOffset(checkPoint);
            // index изменяется только в потоке обслуживания, поэтому лок можно брать только в момент подмены значений
            if (indexRecordsToLeave.size() * 2 < index.size()) {
                log.info("Optimizing index file");
                File tmpIndex = new File(workDir, INDEX_NAME + ".tmp");
                TreeMap<Long, UUID> newIndex = new TreeMap<>(indexRecordsToLeave);

                OutputStream newIndexWriter = new FileOutputStream(tmpIndex);
                Closeable closeOnFinish = newIndexWriter;
                try {
                    for (Map.Entry<Long, UUID> entry : newIndex.entrySet()) {
                        writeToIndex(newIndexWriter, entry.getKey(), entry.getValue());
                    }
                    Files.move(
                            tmpIndex.toPath(),
                            new File(workDir, INDEX_NAME).toPath(),
                            StandardCopyOption.REPLACE_EXISTING, StandardCopyOption.ATOMIC_MOVE
                    );
                    indexLock.lock();
                    try {
                        closeOnFinish = indexWriter;
                        indexWriter = newIndexWriter;
                        index = newIndex;
                    } finally {
                        indexLock.unlock();
                    }
                    log.info("Optimized index");
                } finally {
                    IOUtils.closeQuietly(closeOnFinish);
                }
            }
        } catch (Exception e) {
            log.error("Failed to optimize index file");
        }

        try {
            Set<UUID> activeParts = new HashSet<>(index.values());
            File[] children = workDir.listFiles();
            if (children == null) {
                log.error("workDir.listFiles() returned null");
            } else {
                for (File file : children) {
                    if (file.getName().startsWith(PART_NAME_PREFIX)) {
                        UUID id;
                        try {
                            id = UUID.fromString(file.getName().substring(PART_NAME_PREFIX.length()));
                        } catch (IllegalArgumentException e) {
                            log.warn("Unknown file in wal directory {}", file.getAbsolutePath());
                            continue;
                        }
                        if (!activeParts.contains(id)) {
                            log.info("Deleting redundant log part {}", id);
                            if (!file.delete()) {
                                log.warn("Failed to delete file {}", file.getAbsolutePath());
                            }
                        }
                    } else {
                        if (file.getName().endsWith(".tmp")) {
                            log.info("Deleting tmp file {}", file.getAbsolutePath());
                            if (!file.delete()) {
                                log.warn("Failed to delete file {}", file.getAbsolutePath());
                            }
                        }
                    }
                }
            }
        } catch (Exception e) {
            log.error("Failed to cleanup files");
        }
        log.debug("Finished maintenance in {}micros", (System.nanoTime() - start) / 1000);

    }

    private class MaintenanceTask implements Runnable {
        @Override
        public void run() {
            maintenance();
        }
    }

    @Override
    public void close() throws IOException {
        if (!closed) {
            closed = true;
            scheduler.shutdown();
            IOUtils.closeQuietly(indexWriter);
            IOUtils.closeQuietly(dataWriter);
            IOUtils.closeQuietly(controlWriter);
        }
    }

    private void assertClosed() {
        if (closed) {
            throw new IllegalStateException("Log is closed");
        }
    }

    private void addNewPart(long offset, UUID id) throws IOException {
        if (!index.isEmpty() && index.lastKey() >= offset) {
            Map.Entry<Long, UUID> lastEntry = index.lastEntry();
            throw new IllegalStateException("Trying to add part " + id + " with offset " + offset +
                    " while last offset is " + lastEntry.getKey() + " " + lastEntry.getValue());
        }
        index.put(offset, id);
        writeToIndex(indexWriter, offset, id);
    }

    private static void writeToIndex(OutputStream out, long offset, UUID partId) throws IOException {
        out.write((offset + "\t" + partId + "\n").getBytes());
        out.flush();
    }

    private SortedMap<Long, UUID> getIndexFromOffset(long offset) {
        Long floor = index.floorKey(offset);
        if (floor == null) {
            return index;
        } else {
            return MapUtils.tailMap(index, floor);
        }
    }

    private static long readCheckPoint(File file) throws IOException {
        long checkPoint = 0L;
        BufferedReader r = new BufferedReader(new FileReader(file));
        String line;
        while ((line = r.readLine()) != null) {
            if (!line.isEmpty()) {
                checkPoint = Long.parseLong(line);
            }
        }
        return checkPoint;
    }

    private static void fillIndex(File file, TreeMap<Long, UUID> index) throws IOException {
        try (InputStream in = new FileInputStream(file)) {
            BufferedReader r = new BufferedReader(new InputStreamReader(in));
            String line;
            while ((line = r.readLine()) != null) {
                if (!line.isEmpty()) {
                    String[] parts = line.split("\t");
                    long offset = Long.parseLong(parts[0]);
                    UUID id = UUID.fromString(parts[1]);
                    index.put(offset, id);
                }
            }
        }
    }

    private static int readData(File file, long offset, WALRecordConsumer consumer, long reportFrom, int limit,
                                boolean strictCorrectness) throws IOException {
        int readRecords = 0;
        int reportedRecords = 0;
        if (file.exists()) {
            byte[] buff = new byte[1024];
            IntBuffer IB = ByteBuffer.wrap(buff).asIntBuffer();
            try (FileInputStream in = new FileInputStream(file)) {
                int len;
                while (((len = in.read(buff, 0, 4)) > 0) && (limit < 0 || reportedRecords < limit)) {
                    if (len != 4) {
                        throw new IOException("Corrupted log part file " + file.getAbsolutePath() +
                                ": expected 4-bytes data length, but found " + len + " bytes");
                    }
                    int recordLength = IB.get(0);
                    if (recordLength > buff.length) {
                        buff = new byte[recordLength];
                        IB = ByteBuffer.wrap(buff).asIntBuffer();
                    }
                    int actualLen = in.read(buff, 0, recordLength);
                    if (actualLen < recordLength) {
                        if (strictCorrectness) {
                            throw new IOException("Corrupted log part file " + file.getAbsolutePath() +
                                    ": expected " + recordLength + "-bytes data, but found " + actualLen + " bytes");
                        } else {
                            return reportedRecords;
                        }
                    }
                    if (consumer != null && (offset + readRecords) >= reportFrom) {
                        consumer.consume(offset + readRecords, buff, 0, actualLen);
                        reportedRecords++;
                    } else if (limit < 0) {
                        reportedRecords++;
                    }
                    readRecords++;
                }
            }
        }
        return reportedRecords;
    }

    private static void updateControlRecord(OutputStream out, long recordId) throws IOException {
        out.write((Long.toString(recordId) + "\n").getBytes());
        out.flush();
    }

    public static WriteAheadLog createQuarantineIfNeeded(File workDir, int minPartSize) throws IOException {
        try {
            return new WriteAheadLog(workDir, minPartSize);
        } catch (IOException e) {
            long now = System.currentTimeMillis();
            File qWorkDir = new File(workDir.getParent(), workDir.getName() + ".quarantine." + now);
            log.error("Log was corrupted, will quarantine into " + qWorkDir.getAbsolutePath(), e);
            Files.move(
                    workDir.toPath(),
                    qWorkDir.toPath(),
                    StandardCopyOption.ATOMIC_MOVE
            );
            Files.createDirectory(workDir.toPath());
            return new WriteAheadLog(workDir, minPartSize);
        }
    }
}
