package com.atlassian.bitbucket.mesh.transaction;

import com.atlassian.bitbucket.mesh.junit.TemporaryDirectory;
import com.atlassian.bitbucket.mesh.rpc.v1.RpcProposedRefUpdate;
import com.atlassian.bitbucket.mesh.rpc.v1.RpcTransaction;
import com.atlassian.bitbucket.mesh.rpc.v1.RpcUpdateRefsRequest;
import com.atlassian.bitbucket.mesh.transaction.TransactionLog.Result;
import com.google.common.collect.ImmutableList;
import com.google.protobuf.ByteString;
import org.junit.Rule;
import org.junit.Test;

import javax.annotation.Nonnull;
import java.io.IOException;
import java.nio.file.Files;
import java.nio.file.Path;
import java.util.List;
import java.util.stream.Collectors;
import java.util.stream.IntStream;
import java.util.stream.Stream;

import static com.atlassian.bitbucket.mesh.transaction.RpcTransactionMatcher.tx;
import static org.apache.commons.lang3.StringUtils.leftPad;
import static org.hamcrest.MatcherAssert.assertThat;
import static org.hamcrest.Matchers.instanceOf;
import static org.hamcrest.collection.IsIterableContainingInOrder.contains;
import static org.junit.Assert.*;

public class FileTransactionLogTest {

    @Rule
    public final TemporaryDirectory temporaryDirectory = new TemporaryDirectory();

    @Test
    public void testAppendAndStreamInlineTransaction() {
        RpcTransaction transaction = repoCreationTransaction();

        FileTransactionLog txLog = new FileTransactionLog(temporaryDirectory.getRoot().resolve("deadbeef.dat"));
        txLog.append(transaction);

        assertThat(toList(txLog), contains(tx(transaction)));
        assertThat(toList("deadbeef.dat"), contains(tx(transaction)));
    }

    @Test
    public void testAppendAndStreamLargeTransaction() {
        RpcTransaction transaction = refUpdateTransaction(1024);

        FileTransactionLog txLog = new FileTransactionLog(temporaryDirectory.getRoot().resolve("deadbeef.dat"));
        txLog.append(transaction);

        assertThat(toList(txLog), contains(tx(transaction)));
        assertThat(toList("deadbeef.dat"), contains(tx(transaction)));
    }

    @Test
    public void testAppendCreatesParentDirectory() {
        Path path = temporaryDirectory.getRoot().resolve("sub-dir").resolve("tx.dat");
        RpcTransaction transaction = repoCreationTransaction();

        FileTransactionLog txLog = new FileTransactionLog(path);
        txLog.append(transaction);

        assertThat(toList(txLog), contains(tx(transaction)));
        assertThat(toList("sub-dir/tx.dat"), contains(tx(transaction)));
    }

    @Test
    public void testAppendCreatesMultipleSegmentFiles() {
        RpcTransaction transaction1 = repoCreationTransaction();
        RpcTransaction transaction2 = refUpdateTransaction(2);
        RpcTransaction transaction3 = refUpdateTransaction(2);
        RpcTransaction transaction4 = refUpdateTransaction(2);
        RpcTransaction transaction5 = refUpdateTransaction(2);
        RpcTransaction transaction6 = refUpdateTransaction(2);
        RpcTransaction transaction7 = refUpdateTransaction(2);
        RpcTransaction transaction8 = refUpdateTransaction(2);

        Path logFile = temporaryDirectory.getRoot().resolve("deadbeef.dat");
        FileTransactionLog txLog = new FileTransactionLog(logFile, 3);

        //Append 3 transactions, which should fill exactly 1 segment
        txLog.append(transaction1);
        txLog.append(transaction2);
        txLog.append(transaction3);
        assertThat(toList(txLog), contains(tx(transaction1), tx(transaction2), tx(transaction3)));
        assertFileCount(1, logFile.getParent());

        //Append 3 more transactions, which should fill another segment
        txLog.append(transaction4);
        txLog.append(transaction5);
        txLog.append(transaction6);
        assertThat(toList(txLog), contains(tx(transaction1), tx(transaction2), tx(transaction3),
                tx(transaction4), tx(transaction5), tx(transaction6)));
        assertFileCount(2, logFile.getParent());

        //Append the final 2 transactions, which should created another segment
        txLog.append(transaction7);
        txLog.append(transaction8);
        assertThat(toList(txLog), contains(tx(transaction1), tx(transaction2), tx(transaction3),
                tx(transaction4), tx(transaction5), tx(transaction6), tx(transaction7), tx(transaction8)));
        assertFileCount(3, logFile.getParent());

        //Now consume 1 transaction from the log
        assertThat(take(txLog, 1), contains(tx(transaction1)));
        assertThat(toList(txLog), contains(tx(transaction2), tx(transaction3), tx(transaction4), tx(transaction5),
                tx(transaction6), tx(transaction7), tx(transaction8)));
        assertFileCount(3, logFile.getParent());

        //Consume 2 more transaction; this should trigger removal of the oldest segment
        assertThat(take(txLog, 2), contains(tx(transaction2), tx(transaction3)));
        assertThat(toList(txLog), contains(tx(transaction4), tx(transaction5), tx(transaction6),
                tx(transaction7), tx(transaction8)));
        assertFileCount(2, logFile.getParent());

        //Consume 4 more transaction; this should trigger removal of another segment
        assertThat(take(txLog, 4), contains(tx(transaction4), tx(transaction5), tx(transaction6), tx(transaction7)));
        assertThat(toList(txLog), contains(tx(transaction8)));
        assertFileCount(1, logFile.getParent());
    }

    @Test
    public void testAppendAndStreamMultipleTransactionsInSingleAppend() {
        Path path = temporaryDirectory.getRoot().resolve("tx.dat");
        RpcTransaction txRepoCreation = repoCreationTransaction();
        RpcTransaction txRefUpdates1 = refUpdateTransaction(1);
        RpcTransaction txRefUpdates2 = refUpdateTransaction(3);
        RpcTransaction txRefUpdates3 = refUpdateTransaction(512);
        RpcTransaction txRepoDeletion = repoDeletionTransaction();

        FileTransactionLog txLog = new FileTransactionLog(path);
        txLog.append(txRepoCreation);
        txLog.append(txRefUpdates1);
        txLog.append(txRefUpdates2);
        txLog.append(txRefUpdates3);
        txLog.append(txRepoDeletion);

        List<RpcTransaction> transactions = toList(txLog);
        assertThat(transactions, contains(
                tx(txRepoCreation), tx(txRefUpdates1), tx(txRefUpdates2), tx(txRefUpdates3), tx(txRepoDeletion)));
    }

    @Test
    public void testAppendAndStreamMultipleTransactionsInMultipleAppends() {
        Path path = temporaryDirectory.getRoot().resolve("tx.dat");
        RpcTransaction txRepoCreation = repoCreationTransaction();
        RpcTransaction txRefUpdates1 = refUpdateTransaction(1);
        RpcTransaction txRefUpdates2 = refUpdateTransaction(3);
        RpcTransaction txRefUpdates3 = refUpdateTransaction(512);
        RpcTransaction txRepoDeletion = repoDeletionTransaction();

        FileTransactionLog txLog = new FileTransactionLog(path);
        txLog.append(txRepoCreation);
        txLog.append(txRefUpdates1);
        txLog.append(txRefUpdates2);
        txLog.append(txRefUpdates3);
        txLog.append(txRepoDeletion);

        List<RpcTransaction> transactions = toList(txLog);
        assertThat(transactions, contains(
                tx(txRepoCreation), tx(txRefUpdates1), tx(txRefUpdates2), tx(txRefUpdates3), tx(txRepoDeletion)));
    }

    @Test
    public void testProcessWholeBatchAndAbort() throws IOException {
        Path path = temporaryDirectory.getRoot().resolve("tx.dat");
        RpcTransaction txRepoCreation = repoCreationTransaction();
        RpcTransaction txRefUpdates1 = refUpdateTransaction(512); // large tx stored in separate details file
        RpcTransaction txRefUpdates2 = refUpdateTransaction(3);
        RpcTransaction txRefUpdates3 = refUpdateTransaction(512); // large tx stored in separate details file
        RpcTransaction txRepoDeletion = repoDeletionTransaction();

        FileTransactionLog txLog = new FileTransactionLog(path);
        txLog.append(txRepoCreation);
        txLog.append(txRefUpdates1);
        txLog.append(txRefUpdates2);
        txLog.append(txRefUpdates3);
        txLog.append(txRepoDeletion);

        //Verify that the details file for txRefUpdates1 and txRefUpdates3 have been created
        try (Stream<Path> stream = Files.list(temporaryDirectory.getRoot())) {
            List<Path> files = stream.collect(Collectors.toList());
            assertEquals(3, files.size());
        }

        //Process the transaction log in batches of 3, consuming the first batch
        txLog.process(transactions -> {
            assertEquals(3, transactions.size());
            assertThat(transactions, contains(tx(txRepoCreation), tx(txRefUpdates1), tx(txRefUpdates2)));
            return Result.abort();
        }, 3);

        List<RpcTransaction> remaining = toList(txLog);
        assertEquals(2, remaining.size());
        assertThat(remaining, contains(tx(txRefUpdates3), tx(txRepoDeletion)));

        //Verify that the details file for txRefUpdates1 has been deleted when the item was consumed
        try (Stream<Path> files = Files.list(temporaryDirectory.getRoot())) {
            assertEquals(2, files.count());
        }
    }

    @Test
    public void testProcessPartialBatchAndAbort() {
        Path path = temporaryDirectory.getRoot().resolve("tx.dat");
        RpcTransaction txRepoCreation = repoCreationTransaction();
        RpcTransaction txRefUpdates1 = refUpdateTransaction(512); // large tx stored in separate details file
        RpcTransaction txRefUpdates2 = refUpdateTransaction(3);
        RpcTransaction txRefUpdates3 = refUpdateTransaction(512); // large tx stored in separate details file
        RpcTransaction txRepoDeletion = repoDeletionTransaction();

        FileTransactionLog txLog = new FileTransactionLog(path);
        txLog.append(txRepoCreation);
        txLog.append(txRefUpdates1);
        txLog.append(txRefUpdates2);
        txLog.append(txRefUpdates3);
        txLog.append(txRepoDeletion);

        //Process the transaction log in batches of 3, consuming the first 2 transactions
        txLog.process(transactions -> {
            assertEquals(3, transactions.size());
            assertThat(transactions, contains(tx(txRepoCreation), tx(txRefUpdates1), tx(txRefUpdates2)));
            return Result.abortAndReturn(1);
        }, 3);

        List<RpcTransaction> remaining = toList(txLog);
        assertEquals(3, remaining.size());
        assertThat(remaining, contains(tx(txRefUpdates2), tx(txRefUpdates3), tx(txRepoDeletion)));
    }

    @Test
    public void testProcessConsumingAll() {
        Path path = temporaryDirectory.getRoot().resolve("tx.dat");
        RpcTransaction txRepoCreation = repoCreationTransaction();
        RpcTransaction txRefUpdates = refUpdateTransaction(512);
        RpcTransaction txRepoDeletion = repoDeletionTransaction();

        FileTransactionLog txLog = new FileTransactionLog(path);
        txLog.append(txRepoCreation);
        txLog.append(txRefUpdates);
        txLog.append(txRepoDeletion);

        //Process the transaction log, consuming all
        ImmutableList.Builder<RpcTransaction> builder = ImmutableList.builder();
        txLog.process(transactions -> {
            builder.addAll(transactions);
            return Result.proceed();
        }, 2);

        List<RpcTransaction> processed = builder.build();
        assertEquals(3, processed.size());
        assertThat(processed, contains(tx(txRepoCreation), tx(txRefUpdates), tx(txRepoDeletion)));

        List<RpcTransaction> remaining = toList(txLog);
        assertEquals(0, remaining.size());

        //Verify that the details file for txRefUpdates has been deleted when the item was consumed,
        //and that the transaction log file has been deleted after being fully consumed.
        assertFileCount(0, path.getParent());

        //Append another tx and verify that it can be read
        txRefUpdates = refUpdateTransaction(128);
        txLog.append(txRefUpdates);
        remaining = toList(txLog);
        assertEquals(1, remaining.size());
        assertThat(remaining.get(0), tx(txRefUpdates));
    }

    @Test
    public void testProcessRetainingAll() {
        Path path = temporaryDirectory.getRoot().resolve("tx.dat");
        RpcTransaction txRepoCreation = repoCreationTransaction();
        RpcTransaction txRefUpdates1 = refUpdateTransaction(1);
        RpcTransaction txRefUpdates2 = refUpdateTransaction(3);
        RpcTransaction txRefUpdates3 = refUpdateTransaction(512);
        RpcTransaction txRepoDeletion = repoDeletionTransaction();

        FileTransactionLog txLog = new FileTransactionLog(path);
        txLog.append(txRepoCreation);
        txLog.append(txRefUpdates1);
        txLog.append(txRefUpdates2);
        txLog.append(txRefUpdates3);
        txLog.append(txRepoDeletion);

        //Process the transaction log, keeping all the transactions offered
        txLog.process(transactions -> Result.abortAndReturn(transactions.size()), 2);

        assertThat(toList(txLog), contains(
                tx(txRepoCreation), tx(txRefUpdates1), tx(txRefUpdates2), tx(txRefUpdates3), tx(txRepoDeletion)));
    }

    @Test
    public void testProcessRetainsBatchWhenProcessorReturnsErrorResult() {
        Path path = temporaryDirectory.getRoot().resolve("tx.dat");
        RpcTransaction txRepoCreation = repoCreationTransaction();
        RpcTransaction txRefUpdates1 = refUpdateTransaction(1);
        RpcTransaction txRefUpdates2 = refUpdateTransaction(3);
        RpcTransaction txRefUpdates3 = refUpdateTransaction(512);
        RpcTransaction txRepoDeletion = repoDeletionTransaction();

        FileTransactionLog txLog = new FileTransactionLog(path);
        txLog.append(txRepoCreation);
        txLog.append(txRefUpdates1);
        txLog.append(txRefUpdates2);
        txLog.append(txRefUpdates3);
        txLog.append(txRepoDeletion);

        TransactionLog.Processor processor = new TransactionLog.Processor() {

            int batchesProcessed;

            @Nonnull
            @Override
            public Result process(@Nonnull List<RpcTransaction> transactions) {
                if (++batchesProcessed == 2) {
                    return Result.error(1, new IOException("testing exception handling"));
                }
                return Result.proceed();
            }
        };

        TransactionLogException exception = assertThrows(TransactionLogException.class,
                () -> txLog.process(processor, 2));
        assertThat(exception.getCause(), instanceOf(IOException.class));

        //The first batch should be consumed, the second partially consumed, and third should be retained
        assertThat(toList(txLog), contains(tx(txRefUpdates3), tx(txRepoDeletion)));
    }

    @Test
    public void testProcessRetainsBatchWhenProcessorThrows() {
        Path path = temporaryDirectory.getRoot().resolve("tx.dat");
        RpcTransaction txRepoCreation = repoCreationTransaction();
        RpcTransaction txRefUpdates1 = refUpdateTransaction(1);
        RpcTransaction txRefUpdates2 = refUpdateTransaction(3);
        RpcTransaction txRefUpdates3 = refUpdateTransaction(512);
        RpcTransaction txRepoDeletion = repoDeletionTransaction();

        FileTransactionLog txLog = new FileTransactionLog(path);
        txLog.append(txRepoCreation);
        txLog.append(txRefUpdates1);
        txLog.append(txRefUpdates2);
        txLog.append(txRefUpdates3);
        txLog.append(txRepoDeletion);

        TransactionLog.Processor processor = new TransactionLog.Processor() {

            int batchesProcessed;

            @Nonnull
            @Override
            public Result process(@Nonnull List<RpcTransaction> transactions) {
                if (++batchesProcessed == 2) {
                    //Throw on the 2nd batch
                    throw new UnsupportedOperationException("testing exception handling");
                }
                return Result.proceed();
            }
        };

        assertThrows(UnsupportedOperationException.class, () -> txLog.process(processor, 2));
        //The first batch should be consumed, but the second and third batches should be retained
        assertThat(toList(txLog), contains(tx(txRefUpdates2), tx(txRefUpdates3), tx(txRepoDeletion)));
    }

    @Test
    public void testStreamNonExistingFile() {
        assertEquals(ImmutableList.of(), toList("does-not-exist.dat"));
    }

    private static void assertFileCount(int expectedCount, Path directory) {
        try (Stream<Path> files = Files.list(directory)) {
            assertEquals(expectedCount, files.count());
        } catch (IOException e) {
            e.printStackTrace();
            fail(e.getMessage());
        }
    }

    private static List<RpcTransaction> take(FileTransactionLog log, int count) {
        ConsumingProcessor processor = new ConsumingProcessor(count);
        log.process(processor, 3);
        return processor.getTransactions();
    }

    private static RpcTransaction refUpdateTransaction(int numRefUpdates) {
        RpcUpdateRefsRequest.Builder builder = RpcUpdateRefsRequest.newBuilder()
                .setRepository("p/0001/h/deadbeef/r/2020")
                .setTransactional(true);
        IntStream.rangeClosed(1, numRefUpdates)
                .forEach(num -> builder.addRefUpdates(RpcProposedRefUpdate.newBuilder()
                        .setRefId(ByteString.copyFromUtf8("refs/heads/branch-" + num))
                        .setOldHash("0000000000000000000000000000000000000000")
                        .setNewHash("6ea70303824cdd755a1db6326d3b14d5ba5" + leftPad(Integer.toString(num), 5, '0'))));
        return RpcTransaction.newBuilder()
                .setRepository("p/0001/h/deadbeef/r/2020")
                .setTimestamp(System.currentTimeMillis())
                .setUpdateRefs(builder)
                .build();
    }

    private static RpcTransaction repoCreationTransaction() {
        return RpcTransaction.newBuilder()
                .setRepository("p/0001/h/deadbeef/r/2020")
                .setTimestamp(System.currentTimeMillis())
                .build();
    }

    private static RpcTransaction repoDeletionTransaction() {
        return RpcTransaction.newBuilder()
                .setRepository("p/0001/h/deadbeef/r/2020")
                .setTimestamp(System.currentTimeMillis())
                .build();
    }

    private static List<RpcTransaction> toList(FileTransactionLog txLog) {
        return txLog.stream().collect(Collectors.toList());
    }

    private static List<RpcTransaction> toList(Path path) {
        return toList(new FileTransactionLog(path));
    }

    private List<RpcTransaction> toList(String relativePath) {
        return toList(temporaryDirectory.getRoot().resolve(relativePath));
    }

    private static class ConsumingProcessor implements TransactionLog.Processor {

        private final ImmutableList.Builder<RpcTransaction> builder;
        private int remaining;

        ConsumingProcessor(int count) {
            builder = ImmutableList.builder();
            remaining = count;
        }

        public List<RpcTransaction> getTransactions() {
            return builder.build();
        }

        @Nonnull
        @Override
        public Result process(@Nonnull List<RpcTransaction> transactions) {
            int batchSize = transactions.size();
            if (remaining < batchSize) {
                builder.addAll(transactions.subList(0, remaining));
                Result result = Result.abortAndReturn(batchSize - remaining);
                remaining = 0;
                return result;
            }
            builder.addAll(transactions);
            remaining -= batchSize;
            return Result.proceed();
        }
    }
}