package com.atlassian.bitbucket.mesh.io;

import com.google.common.base.Throwables;

import javax.annotation.Nonnull;
import javax.annotation.Nullable;
import javax.annotation.WillNotClose;
import java.io.*;
import java.nio.channels.*;
import java.nio.file.FileAlreadyExistsException;
import java.nio.file.Files;
import java.nio.file.NoSuchFileException;
import java.nio.file.Path;
import java.security.SecureRandom;
import java.time.Duration;
import java.util.concurrent.locks.LockSupport;

import static java.nio.charset.StandardCharsets.UTF_8;
import static java.nio.file.StandardCopyOption.REPLACE_EXISTING;
import static java.nio.file.StandardOpenOption.CREATE_NEW;
import static java.nio.file.StandardOpenOption.WRITE;
import static java.util.Objects.requireNonNull;
import static org.apache.commons.lang3.StringUtils.isNotBlank;

/**
 * Utility class used to edit a file, complete with locking to prevent concurrent edits by other cluster nodes or
 * {@code git} processes.
 */
public class UnderLock {

    private static final SecureRandom RANDOM = new SecureRandom();

    private final Path file;
    private final Duration timeout;

    private UnderLock(@Nonnull Path file, @Nonnull Duration timeout) {
        this.file = file;
        this.timeout = timeout;
    }

    @Nonnull
    public static UnderLock forFile(@Nonnull Path file) {
        return new UnderLock(file, Duration.ZERO);
    }

    @Nonnull
    public static UnderLock forFile(@Nonnull Path file, @Nonnull Duration timeout) {
        return new UnderLock(requireNonNull(file, "file"), requireNonNull(timeout, "timeout"));
    }

    /**
     * Edits the file, complete with locking to prevent concurrent edits by other cluster nodes or {@code git}
     * processes.
     * <p>
     * If the file already exists, the {@code mutator} receives an {@code Reader} for it. Otherwise, the
     * first argument will be {@code null}. The {@code mutator} <i>always</i> receives an {@code Writer}, for
     * emitting the new contents.
     * <p>
     * The primary use case for this method is guarding an edit with a lock.
     *
     * @param mutator    a function to use for mutating the file
     * @param withBackup whether the old version of the file should be saved with a {@code .bak} extension
     * @throws IOException                  if the {@code mutator} fails, or if the file cannot be locked
     * @throws OverlappingFileLockException if the file can't be locked because it has already been locked
     */
    public void edit(@Nonnull IoBiConsumer<? super BufferedReader, ? super Writer> mutator,
                     boolean withBackup) throws IOException, OverlappingFileLockException {
        try (LockedFile locked = acquireLock(true)) {
            locked.edit(mutator, withBackup);
        }
    }

    /**
     * Edits the file, complete with locking to prevent concurrent edits by other cluster nodes or {@code git}
     * processes.
     * <p>
     * If the file already exists, the {@code mutator} receives an {@code InputStream} for it. Otherwise, the
     * first argument will be {@code null}. The {@code mutator} <i>always</i> receives an {@code OutputStream}, for
     * emitting the new contents.
     * <p>
     * The primary use case for this method is guarding an edit with a lock.
     *
     * @param mutator    a function to use for mutating the file
     * @param withBackup whether the old version of the file should be saved with a {@code .bak} extension
     * @throws IOException                  if the {@code mutator} fails, or if the file cannot be locked
     * @throws OverlappingFileLockException if the file can't be locked because it has already been locked
     */
    public void editRaw(@Nonnull IoBiConsumer<? super InputStream, ? super OutputStream> mutator,
                        boolean withBackup) throws IOException {
        requireNonNull(mutator, "mutator");

        try (LockedFile locked = acquireLock(true)) {
            locked.editRaw(mutator, withBackup);
        }
    }

    /**
     * Acquires a lock on the file, returning a {@link LockedFile} that can be used to release the lock.
     *
     * @return a {@link LockedFile} that when closed, releases the lock
     * @throws IOException                  if an I/O error occurred when acquiring the lock
     * @throws OverlappingFileLockException if the lock could not be acquired within the specified timeout
     */
    @Nonnull
    public LockedFile lock() throws IOException, OverlappingFileLockException {
        return acquireLock(true);
    }

    /**
     * Checks whether the lock is currently held, and if not, creates the lock file with the provided {@code content}
     * as the file content.
     * <p>
     * The effect of this operation is that any attempt to acquire an exclusive lock or edit the target file
     * under lock will fail, which is useful for situations where the lock should survive node restarts,
     * (e.g. pausing of garbage collection).
     *
     * @throws IOException                  if an I/O error occurred when acquiring the lock
     * @throws OverlappingFileLockException if the lock could not be acquired within the specified timeout
     */
    public void lockIndefinitely(@Nullable String content) throws IOException, OverlappingFileLockException {
        //Take out a lock where the lock file is not deleted when the lock is released
        try (LockedFile locked = acquireLock(false)) {
            if (isNotBlank(content)) {
                //write to the lock file itself
                locked.edit((ignored, writer) -> writer.write(content), false);
                //but make sure the file is unchanged
                locked.rollBack();
            }
        }
    }

    /**
     * Releases a {@link #lock previously acquired lock}, including stale lock files, or locks that were created
     * through {@link #lockIndefinitely(String)}. If the lock file does not exist, nothing happens.
     *
     * @throws IOException                  if the lock cannot be released
     * @throws OverlappingFileLockException if another process or thread is actively holding the lock
     */
    public void releaseIfNotHeld() throws IOException {
        releaseIfNotHeld(null);
    }

    /**
     * Releases a {@link #lock previously acquired lock}, including stale lock files, or locks that were created
     * through {@link #lockIndefinitely(String)}. If the lock is successfully released, the provided cleanup task
     * is run. If the lock file does not exist the cleanup task is not run and nothing happens.
     *
     * @throws IOException                  if the lock cannot be released
     * @throws OverlappingFileLockException if another process or thread is actively holding the lock
     */
    public void releaseIfNotHeld(@Nullable Runnable cleanupTask) throws IOException {
        Path lockFile = getLockFile();
        try (FileChannel lockChannel = FileChannel.open(lockFile, WRITE)) {
            //Verify that no other thread or process is currently holding an _active_ lock on the lockChannel
            try (FileLock lock = lockChannel.tryLock()) {
                if (lock == null) {
                    //The lock is already held by another process or thread
                    throw new OverlappingFileLockException();
                }
                if (cleanupTask != null) {
                    cleanupTask.run();
                }
                Files.delete(lockFile);
            }
        } catch (FileNotFoundException | NoSuchFileException e) {
            //This lock was not held (anymore). Nothing left to do
        }
    }

    /**
     * Acquires a lock on the file, and then writes the contents of the provided {@link ReadableByteChannel} to it.
     * Contents are written to the lock first and then renamed into place to produce an (effectively) atomic write
     * even on NFS.
     *
     * @param source the contents to write
     * @throws IOException                  if the file cannot be locked, or the contents cannot be written
     * @throws OverlappingFileLockException if the file can't be locked because it has already been locked
     */
    public void write(@Nonnull @WillNotClose ReadableByteChannel source) throws IOException {
        try (LockedFile locked = acquireLock(true)) {
            locked.write(source);
        }
    }

    private static InputStream newInputStream(Path file) throws IOException {
        try {
            return Files.newInputStream(file);
        } catch (FileNotFoundException | NoSuchFileException e) {
            return null;
        }
    }

    private static BufferedReader newReader(Path file) throws IOException {
        try {
            return Files.newBufferedReader(file, UTF_8);
        } catch (FileNotFoundException | NoSuchFileException e) {
            return null;
        }
    }

    private static FileChannel openLockWithHardLink(Path lockFile) throws IOException {
        //File creation via NFS does not appear to be atomic. If the lock is highly contested,
        //it is possible that lock files are visible before their 'mode' or 'last modified timestamp'
        //is written. This breaks the detection of stale lock files, as they rely on a correct
        //'lastModified' timestamp.
        //
        //Instead of creating the lock file directly, the file is created in several steps:
        //
        // 1. A temporary file is created
        // 2. A FileChannel is opened for the temporary file
        // 3. A hard link with the name of the desired lock file, linking to the temporary file,
        //    is created
        // 4. The temporary file is deleted
        // 5. The FileChannel of #2 is returned
        //
        //If any of these operations fail, a fallback of directly opening the lock file is used.
        //
        //Returning of the FileChannel for the temporary file works as intended because hard links
        //share locking semantics and file attributes between each link.

        //Step #1
        long n = RANDOM.nextLong();
        n = (n == Long.MIN_VALUE) ? 0 : Math.abs(n);
        Path temporaryLockFile = lockFile.resolveSibling("tmp-" + Long.toString(n, 16) + "-" + lockFile.getFileName());
        //Step #2
        //DELETE_ON_CLOSE on Linux deletes the file immediately after creation, before returning
        //the FileChannel. As a result it can't be used.
        FileChannel lockChannel = FileChannel.open(temporaryLockFile, CREATE_NEW, WRITE);

        Exception caughtException = null;
        try {
            //Step #3
            Files.createLink(lockFile, temporaryLockFile);
        } catch (FileAlreadyExistsException e) {
            //FileAlreadyExistsException means the lock could not be acquired - this will be handled
            //in the 'catch' block in 'withLock()'
            caughtException = e;
            throw e;
        } catch (UnsupportedOperationException | IOException e) {
            //Hard linking may not be supported; fall back to locking without hard linking
            try {
                return FileChannel.open(lockFile, CREATE_NEW, WRITE);
            } catch (IOException suppressed) {
                //If the fallback fails as well, leave the exception from Files.createLink as the main
                //exception and add this one as the suppressed one.
                e.addSuppressed(suppressed);
            }
            caughtException = e;
            throw e;
        } finally {
            //If an exception is being handled, the channel should be closed before returning
            if (caughtException != null) {
                try {
                    lockChannel.close();
                } catch (IOException suppressed) {
                    caughtException.addSuppressed(suppressed);
                }
            }
            try {
                //Step #4
                Files.deleteIfExists(temporaryLockFile);
            } catch (IOException suppressed) {
                if (caughtException == null) {
                    throw suppressed;
                }
                caughtException.addSuppressed(suppressed);
            }
        }
        //Step #5
        return lockChannel;
    }

    private LockedFile acquireLock(boolean deleteOnClose) throws IOException {
        Path lockFile = getLockFile();
        FileChannel lockChannel = null;
        FileLock lock = null;

        try {
            //Try to acquire the file lock
            long deadline = System.nanoTime() + timeout.toNanos();
            long sleepNanos = 100_000L; // wait 0.1ms on the first failure to acquire the lock
            while (lock == null) {
                try {
                    lockChannel = openLockWithHardLink(lockFile);
                    lock = lockChannel.tryLock(0L, Long.MAX_VALUE, false);
                } catch (OverlappingFileLockException e) {
                    // lock could not be acquired
                } catch (FileAlreadyExistsException e) {
                    //Since the lock file is opened with CREATE_NEW, it's far more likely locking will fail
                    //because the file already exists than because the FileLock couldn't be acquired
                    if (!e.getFile().equals(lockFile.toString())) {
                        throw e;
                    }
                }

                if (lock == null) {
                    long curTime = System.nanoTime();
                    if (curTime > deadline) {
                        //Failed to acquire the file lock without timeout
                        throw new OverlappingFileLockException();
                    }
                    LockSupport.parkNanos(Math.min(deadline - curTime, sleepNanos));
                    sleepNanos = Math.min(50_000_000L, sleepNanos * 2); //Next loop wait longer, up to 50ms
                }
            }

            //The file lock was successfully acquired
            return new DefaultLockedFile(lockChannel, lock, lockFile, deleteOnClose);
        } catch (IOException | RuntimeException exception) {
            try {
                if (lockChannel != null) {
                    lockChannel.close();
                }
            } catch (IOException e) {
                exception.addSuppressed(e);
            }
            throw exception;
        }
    }

    private Path getLockFile() {
        return file.resolveSibling(file.getFileName() + ".lock");
    }

    public interface LockedFile extends AutoCloseable {

        /**
         * Releases the lock and updates the locked file with the changes staged through {@link #edit},
         * {@link #editRaw} or {@link #write} (unless they have been {@link #rollBack() rolled back}).
         *
         * @throws IOException if an I/O error occurs while applying staged changes or releasing the lock
         */
        @Override
        void close() throws IOException;

        /**
         * Edits the file, under lock, to prevent concurrent edits by other threads or processes.
         * <p>
         * If the file already exists, the {@code mutator} receives an {@code Reader} for it. Otherwise, the
         * first argument will be {@code null}. The {@code mutator} <i>always</i> receives an {@code Writer}, for
         * emitting the new contents.
         * <p>
         * The changes to the file content are staged and will not be visible until the lock is released.
         * If {@link #rollBack()} is called before the lock is {@link #close()} released, the staged changes are
         * discarded and the file contents remain unchanged.
         * <p>
         * The primary use case for this method is guarding an edit with a lock.
         *
         * @param mutator    a function to use for mutating the file
         * @param withBackup whether the old version of the file should be saved with a {@code .bak} extension
         * @throws IOException           if the {@code mutator} fails, or if the file cannot be locked
         * @throws IllegalStateException if called after the lock has been {@link #close() released}
         */
        void edit(@Nonnull IoBiConsumer<? super BufferedReader, ? super Writer> mutator,
                  boolean withBackup) throws IOException;

        /**
         * Edits the file, under lock, to prevent concurrent edits by other threads or processes.
         * <p>
         * If the file already exists, the {@code mutator} receives an {@code InputStream} for it. Otherwise, the
         * first argument will be {@code null}. The {@code mutator} <i>always</i> receives an {@code OutputStream}, for
         * emitting the new contents.
         * <p>
         * The changes to the file content are staged and will not be visible until the lock is released.
         * If {@link #rollBack()} is called before the lock is {@link #close()} released, the staged changes are
         * discarded and the file contents remain unchanged.
         * <p>
         * The primary use case for this method is guarding an edit with a lock.
         *
         * @param mutator    a function to use for mutating the file
         * @param withBackup whether the old version of the file should be saved with a {@code .bak} extension
         * @throws IOException           if the {@code mutator} fails, or if the file cannot be locked
         * @throws IllegalStateException if called after the lock has been {@link #close() released}
         */
        void editRaw(@Nonnull IoBiConsumer<? super InputStream, ? super OutputStream> mutator,
                     boolean withBackup) throws IOException;

        /**
         * Rolls back any changes to the file content staged through {@link #edit(IoBiConsumer, boolean)},
         * {@link #editRaw(IoBiConsumer, boolean)} or {@link #write(ReadableByteChannel)}
         *
         * @throws IllegalStateException if called after the lock has been {@link #close() released}
         */
        void rollBack();

        /**
         * Acquires a lock on the file, and then writes the contents of the provided {@link ReadableByteChannel} to it.
         * Contents are written to the lock first and then renamed into place to produce an (effectively) atomic write
         * even on NFS.
         * <p>
         * The changes to the file content are staged and will not be visible until the lock is released.
         * If {@link #rollBack()} is called before the lock is {@link #close()} released, the staged changes are
         * discarded and the file contents remain unchanged.
         *
         * @param source the contents to write
         * @throws IOException           if the file cannot be locked, or the contents cannot be written
         * @throws IllegalStateException if called after the lock has been {@link #close() released}
         */
        void write(@Nonnull @WillNotClose ReadableByteChannel source) throws IOException;
    }

    private class DefaultLockedFile implements LockedFile {

        private final boolean deleteOnClose;
        private final FileLock lock;
        private final FileChannel lockChannel;
        private final Path lockFile;

        private volatile boolean closed;
        private volatile boolean staged;
        private volatile boolean withBackup;

        private DefaultLockedFile(FileChannel lockChannel, FileLock lock, Path lockFile, boolean deleteOnClose) {
            this.deleteOnClose = deleteOnClose;
            this.lockChannel = lockChannel;
            this.lock = lock;
            this.lockFile = lockFile;
        }

        @Override
        public void close() throws IOException {
            if (closed) {
                return;
            }

            Exception exception = null;
            if (staged) {
                try {
                    //_While still holding the lock_, move the current file out of the way to prepare for unlocking
                    if (withBackup) {
                        Files.move(file, file.resolveSibling(file.getFileName() + ".bak"), REPLACE_EXISTING);
                    } else {
                        //Files.move's semantics are not guaranteed if the target file exists, so it needs
                        //to be explicitly deleted first to ensure consistent unlocking behavior
                        Files.deleteIfExists(file);
                    }

                    //Publish the staged content
                    Files.move(lockFile, file);
                } catch (IOException | RuntimeException e) {
                    exception = e;
                }
            }

            try {
                lock.close();
            } catch (IOException | RuntimeException e) {
                if (exception == null) {
                    exception = e;
                } else {
                    exception.addSuppressed(e);
                }
            }

            try {
                lockChannel.close();
            } catch (IOException | RuntimeException e) {
                if (exception == null) {
                    exception = e;
                } else {
                    exception.addSuppressed(e);
                }
            }

            if (deleteOnClose) {
                try {
                    Files.deleteIfExists(lockFile);
                } catch (IOException | RuntimeException e) {
                    if (exception == null) {
                        exception = e;
                    } else {
                        exception.addSuppressed(e);
                    }
                }
            }
            if (exception != null) {
                Throwables.propagateIfPossible(exception, IOException.class);
            }
            closed = true;
        }

        @Override
        public void edit(@Nonnull IoBiConsumer<? super BufferedReader, ? super Writer> mutator,
                         boolean withBackup) throws IOException {
            checkLockHeld();
            //Stage the updated file contents in the lock file
            Writer lockWriter = Channels.newWriter(lockChannel, UTF_8.newEncoder(), 8192);

            try (BufferedReader reader = newReader(file)) { //try-with-resources handles null automatically
                this.withBackup = withBackup && reader != null;
                mutator.accept(reader, lockWriter);
                staged = true;
            }
            //Flush the lockWriter but don't close it because it would close the underlying lockChannel
            //and release the lock
            lockWriter.flush();
        }

        @Override
        public void editRaw(@Nonnull IoBiConsumer<? super InputStream, ? super OutputStream> mutator,
                            boolean withBackup) throws IOException {
            checkLockHeld();
            //Stage the updated file contents in the lock file
            OutputStream lockOutput = Channels.newOutputStream(lockChannel);

            try (InputStream inputStream = newInputStream(file)) { //try-with-resources handles null automatically
                this.withBackup = withBackup && inputStream != null;

                mutator.accept(inputStream, lockOutput);

                staged = true;
            }
            //Flush the lockOutput but don't close it because it would close the underlying lockChannel
            //and release the lock
            lockOutput.flush();
        }

        @Override
        public void rollBack() {
            checkLockHeld();
            staged = false;
        }

        @Override
        public void write(@Nonnull @WillNotClose ReadableByteChannel source) throws IOException {
            checkLockHeld();
            lockChannel.transferFrom(source, 0L, Long.MAX_VALUE);
            staged = true;
        }

        private void checkLockHeld() {
            if (closed) {
                throw new IllegalStateException("Lock on " + file + " has already been released");
            }
        }
    }
}
