package com.atlassian.bitbucket.mesh.git;

import com.atlassian.bitbucket.mesh.git.exception.LfsChecksumValidationFailedException;
import com.atlassian.bitbucket.mesh.git.exception.LfsStorageException;
import org.apache.commons.codec.binary.Hex;
import org.apache.commons.codec.digest.DigestUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.annotation.Nonnull;
import java.io.IOException;
import java.nio.ByteBuffer;
import java.nio.channels.WritableByteChannel;
import java.nio.file.FileAlreadyExistsException;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.StandardOpenOption;
import java.security.MessageDigest;

import static java.util.Objects.requireNonNull;

/**
 * Used for writing LFS objects to the filesystem.
 *
 * The incoming object is first uploaded to a temporary file. This protects the server from being in a state where
 * an object exists on disk but the checksum does not match the filename. Once all writes are complete the caller
 * must call {@link #close()} at which time the completed checksum will be compared to the expected one, and if they
 * match the object will be atomically moved to the correct location thus making it available for download by clients.
 */
public class DefaultLfsObjectWriter implements LfsObjectWriter {

    private static final Logger log = LoggerFactory.getLogger(DefaultLfsObjectWriter.class);

    private final Path objectFile;
    private final String oid;
    private final WritableByteChannel outputChannel;
    private final MessageDigest messageDigest;
    private final Path tempFile;

    /**
     * @param oid           the OID that the client supplied
     * @param objectFile    the actual path for the object file
     * @throws IOException if an exception happens while writing the {@code objectFile}
     */
    public DefaultLfsObjectWriter(@Nonnull String oid, @Nonnull Path objectFile) throws IOException {
        this.oid = requireNonNull(oid, "oid");
        this.objectFile = requireNonNull(objectFile, "objectFile");

        messageDigest = DigestUtils.getSha256Digest();

        // The incoming object is first uploaded to a temporary file. This protects the server from being in a
        // state where an object exists on disk but the checksum does not match the filename.
        // The temporary file is created in the same directory where the real object will go, using a similar-
        // but-randomized filename to ensure concurrent uploads of the same object use distinct files
        //
        // The temporary file path will be of the form: <hierarchy>/<sha256[0:1]>/<sha256[2:63]>-<random>.tmp
        tempFile = Files.createTempFile(objectFile.getParent(), objectFile.getFileName() + "-", ".tmp");
        outputChannel = Files.newByteChannel(tempFile, StandardOpenOption.WRITE);
    }

    @Override
    public void abort() {
        cleanup();
    }

    @Override
    public void close() {
        if (!outputChannel.isOpen()) {
            log.error("close() called after abort, error or previous close");
            return;
        }

        try {
            outputChannel.close();
            String checksum = Hex.encodeHexString(messageDigest.digest());
            if (checksum.equalsIgnoreCase(oid)) {
                try {
                    Files.move(tempFile, objectFile);
                } catch (FileAlreadyExistsException e) {
                    // The same object has been uploaded on another connection. Both are guaranteed to have the
                    // same contents, so the collision is not an error
                }
            } else {
                log.warn("Upload failed - OID {} does not match object checksum {}", oid, checksum);
                throw new LfsChecksumValidationFailedException(oid, checksum);
            }
        } catch (IOException e) {
            log.error("Error closing/finalising LFS object write", e);
            throw new LfsStorageException(e);
        } finally {
            cleanup();
        }
    }

    @Override
    public void write(@Nonnull ByteBuffer buffer) {
        if (!outputChannel.isOpen()) {
            log.error("write() called after abort, error or close");
            return;
        }

        try {
            buffer.mark();
            outputChannel.write(buffer);
            buffer.reset();
            messageDigest.update(buffer);
        } catch (IOException e) {
            log.error("Error writing LFS object content", e);
            cleanup();
            throw new LfsStorageException(e);
        }
    }

    private void cleanup() {
        try {
            outputChannel.close();
            Files.deleteIfExists(tempFile);
        } catch (IOException e) {
            log.error("Error cleaning up after LFS object write", e);
        }
    }
}
