package freenet.client.async;

import static org.junit.Assert.*;

import java.io.DataInputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.util.Arrays;
import java.util.Collections;
import java.util.HashSet;
import java.util.List;

import org.junit.Assert;
import org.junit.Before;
import org.junit.Test;

import freenet.client.ClientMetadata;
import freenet.client.FetchContext;
import freenet.client.FetchException;
import freenet.client.HighLevelSimpleClientImpl;
import freenet.client.InsertContext;
import freenet.client.InsertContext.CompatibilityMode;
import freenet.client.Metadata;
import freenet.client.Metadata.SplitfileAlgorithm;
import freenet.client.MetadataParseException;
import freenet.client.MetadataUnresolvedException;
import freenet.client.OnionFECCodec;
import freenet.client.events.SimpleEventProducer;
import freenet.crypt.CRCChecksumChecker;
import freenet.crypt.ChecksumFailedException;
import freenet.crypt.DummyRandomSource;
import freenet.keys.CHKBlock;
import freenet.keys.CHKEncodeException;
import freenet.keys.ClientCHK;
import freenet.keys.ClientCHKBlock;
import freenet.keys.FreenetURI;
import freenet.keys.Key;
import freenet.keys.NodeCHK;
import freenet.node.BaseSendableGet;
import freenet.node.KeysFetchingLocally;
import freenet.node.SendableRequestItemKey;
import freenet.support.CheatingTicker;
import freenet.support.DummyJobRunner;
import freenet.support.MemoryLimitedJobRunner;
import freenet.support.PooledExecutor;
import freenet.support.Ticker;
import freenet.support.WaitableExecutor;
import freenet.support.api.Bucket;
import freenet.support.api.BucketFactory;
import freenet.support.api.LockableRandomAccessBuffer;
import freenet.support.api.LockableRandomAccessBufferFactory;
import freenet.support.compress.Compressor.COMPRESSOR_TYPE;
import freenet.support.io.ArrayBucketFactory;
import freenet.support.io.BucketTools;
import freenet.support.io.ByteArrayRandomAccessBufferFactory;
import freenet.support.io.NativeThread;
import freenet.support.io.StorageFormatException;

public class SplitFileFetcherStorageTest {

    // Setup code is considerable. See below for actual tests ...

    static DummyRandomSource random;
    static final KeySalter salt = Key::getRoutingKey;
    static BucketFactory bf = new ArrayBucketFactory();
    static LockableRandomAccessBufferFactory rafFactory = new ByteArrayRandomAccessBufferFactory();
    static final WaitableExecutor exec = new WaitableExecutor(new PooledExecutor());
    static final PersistentJobRunner jobRunner = new DummyJobRunner(exec, null);
    static final Ticker ticker = new CheatingTicker(exec);
    static MemoryLimitedJobRunner memoryLimitedJobRunner = new MemoryLimitedJobRunner(9 * 1024 * 1024L, 20, exec, NativeThread.JAVA_PRIORITY_RANGE);
    static final int BLOCK_SIZE = CHKBlock.DATA_LENGTH;
    private static final OnionFECCodec codec = new OnionFECCodec();
    private static final int MAX_SEGMENT_SIZE = 256;
    static final int KEY_LENGTH = 32;
    static final CompatibilityMode COMPATIBILITY_MODE = InsertContext.CompatibilityMode.COMPAT_1416;
    static FreenetURI URI;
    private static final List<COMPRESSOR_TYPE> NO_DECOMPRESSORS = Collections.emptyList();

    @Before
    public void setUp() {
        // Reset RNG for each test (it's static so not reset by junit).
        random = new DummyRandomSource(1234);
        URI = FreenetURI.generateRandomCHK(random);
    }

    private static class TestSplitfile {
        final Bucket originalData;
        final Metadata metadata;
        final byte[][] dataBlocks;
        final byte[][] checkBlocks;
        final ClientCHK[] dataKeys;
        final ClientCHK[] checkKeys;
        private final byte[] cryptoKey;
        private final byte cryptoAlgorithm;
        private final int[] segmentDataBlockCount;
        private final int[] segmentCheckBlockCount;
        private final boolean persistent;
        final MyKeysFetchingLocally fetchingKeys;

        private TestSplitfile(Bucket data, Metadata m, byte[][] originalDataBlocks,
                              byte[][] originalCheckBlocks, ClientCHK[] dataKeys, ClientCHK[] checkKeys,
                              byte[] cryptoKey, byte cryptoAlgorithm, int[] segmentDataBlockCount,
                              int[] segmentCheckBlockCount, boolean persistent) {
            this.originalData = data;
            this.metadata = m;
            this.dataBlocks = originalDataBlocks;
            this.checkBlocks = originalCheckBlocks;
            this.dataKeys = dataKeys;
            this.checkKeys = checkKeys;
            this.cryptoKey = cryptoKey;
            this.cryptoAlgorithm = cryptoAlgorithm;
            this.segmentDataBlockCount = segmentDataBlockCount;
            this.segmentCheckBlockCount = segmentCheckBlockCount;
            this.persistent = persistent;
            this.fetchingKeys = new MyKeysFetchingLocally();
        }

        void free() {
            originalData.free();
        }

        static TestSplitfile constructSingleSegment(long size, int checkBlocks, boolean persistent) throws IOException, CHKEncodeException, MetadataUnresolvedException, MetadataParseException {
            assertTrue(checkBlocks <= MAX_SEGMENT_SIZE);
            assertTrue(size < MAX_SEGMENT_SIZE * (long) BLOCK_SIZE);
            Bucket data = makeRandomBucket(size);
            byte[][] originalDataBlocks = splitAndPadBlocks(data, size);
            int dataBlocks = originalDataBlocks.length;
            assertTrue(dataBlocks <= MAX_SEGMENT_SIZE);
            assertTrue(dataBlocks + checkBlocks <= MAX_SEGMENT_SIZE);
            byte[][] originalCheckBlocks = constructBlocks(checkBlocks);
            codec.encode(originalDataBlocks, originalCheckBlocks, falseArray(checkBlocks), BLOCK_SIZE);
            ClientMetadata cm = new ClientMetadata(null);
            // FIXME no hashes for now.
            // FIXME no compression for now.
            byte[] cryptoKey = randomKey();
            byte cryptoAlgorithm = Key.ALGO_AES_CTR_256_SHA256;
            ClientCHK[] dataKeys = makeKeys(originalDataBlocks, cryptoKey, cryptoAlgorithm);
            ClientCHK[] checkKeys = makeKeys(originalCheckBlocks, cryptoKey, cryptoAlgorithm);
            Metadata m = new Metadata(SplitfileAlgorithm.ONION_STANDARD, dataKeys, checkKeys, dataBlocks,
                checkBlocks, 0, cm, size, null, null, size, false, null, null, size, size, dataBlocks,
                dataBlocks + checkBlocks, false, COMPATIBILITY_MODE,
                cryptoAlgorithm, cryptoKey, true, 0
            );
            // Make sure the metadata is reusable.
            // FIXME also necessary as the above constructor doesn't set segments.
            Bucket metaBucket = m.toBucket(bf);
            Metadata m1 = Metadata.construct(metaBucket);
            Bucket copyBucket = m1.toBucket(bf);
            assertTrue(BucketTools.equalBuckets(metaBucket, copyBucket));
            metaBucket.free();
            copyBucket.free();
            return new TestSplitfile(data, m1, originalDataBlocks, originalCheckBlocks, dataKeys, checkKeys,
                cryptoKey, cryptoAlgorithm, null, null, persistent
            );
        }

        /**
         * Create a multi-segment test splitfile. The main complication with multi-segment is that we can't
         * choose the number of blocks in each segment arbitrarily; that depends on the metadata format; the
         * caller must ensure that the number are consistent.
         *
         * @param size
         * @param segmentDataBlockCount    The actual number of data blocks in each segment. Must be consistent
         *                                 with the other parameters; this cannot be chosen freely due to the metadata format.
         * @param segmentCheckBlockCount   The actual number of check blocks in each segment. Must be consistent
         *                                 with the other parameters; this cannot be chosen freely due to the metadata format.
         * @param segmentSize              The "typical" number of data blocks in a segment.
         * @param checkSegmentSize         The "typical" number of check blocks in a segment.
         * @param deductBlocksFromSegments The number of segments from which a single block has been deducted.
         *                                 This is used when the number of data blocks isn't an exact multiple of the number of segments.
         * @param topCompatibilityMode     The "short" value of the definitive compatibility mode used to create
         *                                 the splitfile. This must again be consistent with the rest, as it is sometimes used in decoding.
         * @return
         * @throws IOException
         * @throws CHKEncodeException
         * @throws MetadataUnresolvedException
         * @throws MetadataParseException
         */
        static TestSplitfile constructMultipleSegments(long size, int[] segmentDataBlockCount,
                                                       int[] segmentCheckBlockCount, int segmentSize, int checkSegmentSize,
                                                       int deductBlocksFromSegments, CompatibilityMode topCompatibilityMode)
            throws IOException, CHKEncodeException, MetadataUnresolvedException, MetadataParseException {
            int dataBlocks = sum(segmentDataBlockCount);
            int checkBlocks = sum(segmentCheckBlockCount);
            int segments = segmentDataBlockCount.length;
            assertEquals((size + BLOCK_SIZE - 1) / BLOCK_SIZE, dataBlocks);
            assertEquals(segments, segmentCheckBlockCount.length);
            Bucket data = makeRandomBucket(size);
            byte[][] originalDataBlocks = splitAndPadBlocks(data, size);
            byte[][] originalCheckBlocks = constructBlocks(checkBlocks);
            int startDataBlock = 0;
            int startCheckBlock = 0;
            for (int seg = 0; seg < segments; seg++) {
                byte[][] segmentDataBlocks = Arrays.copyOfRange(originalDataBlocks, startDataBlock, startDataBlock + segmentDataBlockCount[seg]);
                byte[][] segmentCheckBlocks = Arrays.copyOfRange(originalCheckBlocks, startCheckBlock, startCheckBlock + segmentCheckBlockCount[seg]);
                codec.encode(segmentDataBlocks, segmentCheckBlocks, falseArray(segmentCheckBlocks.length), BLOCK_SIZE);
                startDataBlock += segmentDataBlockCount[seg];
                startCheckBlock += segmentCheckBlockCount[seg];
            }
            ClientMetadata cm = new ClientMetadata(null);
            // FIXME no hashes for now.
            // FIXME no compression for now.
            byte[] cryptoKey = randomKey();
            byte cryptoAlgorithm = Key.ALGO_AES_CTR_256_SHA256;
            ClientCHK[] dataKeys = makeKeys(originalDataBlocks, cryptoKey, cryptoAlgorithm);
            ClientCHK[] checkKeys = makeKeys(originalCheckBlocks, cryptoKey, cryptoAlgorithm);
            Metadata m = new Metadata(SplitfileAlgorithm.ONION_STANDARD, dataKeys, checkKeys, segmentSize,
                checkSegmentSize, deductBlocksFromSegments, cm, size, null, null, size, false, null, null,
                size, size, dataBlocks, dataBlocks + checkBlocks, false, topCompatibilityMode,
                cryptoAlgorithm, cryptoKey, true /* FIXME try older splitfiles pre-single-key?? */,
                0
            );
            // Make sure the metadata is reusable.
            // FIXME also necessary as the above constructor doesn't set segments.
            Bucket metaBucket = m.toBucket(bf);
            Metadata m1 = Metadata.construct(metaBucket);
            Bucket copyBucket = m1.toBucket(bf);
            assertTrue(BucketTools.equalBuckets(metaBucket, copyBucket));
            metaBucket.free();
            copyBucket.free();
            return new TestSplitfile(data, m1, originalDataBlocks, originalCheckBlocks, dataKeys, checkKeys,
                cryptoKey, cryptoAlgorithm, segmentDataBlockCount, segmentCheckBlockCount, false
            );
        }

        public CHKBlock encodeDataBlock(int i) throws CHKEncodeException {
            return ClientCHKBlock.encodeSplitfileBlock(dataBlocks[i], cryptoKey, cryptoAlgorithm).getBlock();
        }

        public CHKBlock encodeCheckBlock(int i) throws CHKEncodeException {
            return ClientCHKBlock.encodeSplitfileBlock(checkBlocks[i], cryptoKey, cryptoAlgorithm).getBlock();
        }

        public CHKBlock encodeBlock(int block) throws CHKEncodeException {
            if (block < dataBlocks.length) {
                return encodeDataBlock(block);
            } else {
                return encodeCheckBlock(block - dataBlocks.length);
            }
        }

        public int findCheckBlock(byte[] data, int start) {
            start++;
            for (int i = start; i < checkBlocks.length; i++) {
                if (checkBlocks[i] == data) return i;
            }
            for (int i = start; i < checkBlocks.length; i++) {
                if (Arrays.equals(checkBlocks[i], data)) return i;
            }
            return -1;
        }

        public int findDataBlock(byte[] data, int start) {
            start++;
            for (int i = start; i < dataBlocks.length; i++) {
                if (dataBlocks[i] == data) return i;
            }
            for (int i = start; i < dataBlocks.length; i++) {
                if (Arrays.equals(dataBlocks[i], data)) return i;
            }
            return -1;
        }

        public StorageCallback createStorageCallback() {
            return new StorageCallback(this);
        }

        public SplitFileFetcherStorage createStorage(StorageCallback cb) throws FetchException, MetadataParseException, IOException {
            return createStorage(cb, makeFetchContext());
        }

        public SplitFileFetcherStorage createStorage(final StorageCallback cb, FetchContext ctx) throws FetchException, MetadataParseException, IOException {
            LockableRandomAccessBufferFactory f = new LockableRandomAccessBufferFactory() {

                @Override
                public LockableRandomAccessBuffer makeRAF(long size) throws IOException {
                    LockableRandomAccessBuffer t = rafFactory.makeRAF(size);
                    cb.snoopRAF(t);
                    return t;
                }

                @Override
                public LockableRandomAccessBuffer makeRAF(
                    byte[] initialContents,
                    int offset,
                    int size,
                    boolean readOnly
                ) throws IOException {
                    LockableRandomAccessBuffer t = rafFactory.makeRAF(initialContents, offset, size, readOnly);
                    cb.snoopRAF(t);
                    return t;
                }

            };
            return new SplitFileFetcherStorage(metadata, cb, NO_DECOMPRESSORS, metadata.getClientMetadata(), false,
                (short) COMPATIBILITY_MODE.ordinal(), ctx, false, salt, URI, URI, true, new byte[0], random, bf,
                f, jobRunner, ticker, memoryLimitedJobRunner, new CRCChecksumChecker(), persistent, null, null, fetchingKeys
            );
        }

        /**
         * Restore a splitfile fetcher from a file.
         *
         * @throws StorageFormatException
         * @throws IOException
         * @throws FetchException
         */
        public SplitFileFetcherStorage createStorage(
            StorageCallback cb,
            FetchContext ctx,
            LockableRandomAccessBuffer raf
        ) throws IOException, StorageFormatException, FetchException {
            assertTrue(persistent);
            return new SplitFileFetcherStorage(raf, false, cb, ctx, random, jobRunner, fetchingKeys, ticker, memoryLimitedJobRunner, new CRCChecksumChecker(), false, null, false, false);
        }

        public FetchContext makeFetchContext() {
            return HighLevelSimpleClientImpl.makeDefaultFetchContext(
                Long.MAX_VALUE,
                Long.MAX_VALUE,
                bf,
                new SimpleEventProducer()
            );
        }

        public void verifyOutput(SplitFileFetcherStorage storage) throws IOException {
            StreamGenerator g = storage.streamGenerator();
            Bucket out = bf.makeBucket(-1);
            try (OutputStream os = out.getOutputStream()) {
                g.writeTo(os, null);
            }
            assertTrue(BucketTools.equalBuckets(originalData, out));
            out.free();
        }

        public NodeCHK getCHK(int block) {
            if (block < dataBlocks.length) {
                return dataKeys[block].getNodeCHK();
            } else {
                return checkKeys[block - dataBlocks.length].getNodeCHK();
            }
        }

        public int segmentFor(int block) {
            int total = 0;
            // Must be consistent with the getCHK() counting etc.
            // Count data blocks first then check blocks.
            for (int i = 0; i < segmentDataBlockCount.length; i++) {
                total += segmentDataBlockCount[i];
                if (block < total) return i;
            }
            for (int i = 0; i < segmentCheckBlockCount.length; i++) {
                total += segmentCheckBlockCount[i];
                if (block < total) return i;
            }

            return -1;
        }

    }

    public static ClientCHK[] makeKeys(byte[][] blocks, byte[] cryptoKey, byte cryptoAlgorithm) throws CHKEncodeException {
        ClientCHK[] keys = new ClientCHK[blocks.length];
        for (int i = 0; i < blocks.length; i++) {
            keys[i] = ClientCHKBlock.encodeSplitfileBlock(blocks[i], cryptoKey, cryptoAlgorithm).getClientKey();
        }
        return keys;
    }

    public static int sum(int[] values) {
        int total = 0;
        for (int x : values) {
            total += x;
        }
        return total;
    }

    private static class StorageCallback implements SplitFileFetcherStorageCallback {

        final TestSplitfile splitfile;
        final boolean[] encodedBlocks;
        private boolean succeeded;
        private boolean closed;
        private boolean failed;
        private LockableRandomAccessBuffer raf;

        public StorageCallback(TestSplitfile splitfile) {
            this.splitfile = splitfile;
            encodedBlocks = new boolean[splitfile.dataBlocks.length + splitfile.checkBlocks.length];
        }

        synchronized void snoopRAF(LockableRandomAccessBuffer t) {
            this.raf = t;
        }

        synchronized LockableRandomAccessBuffer getRAF() {
            return raf;
        }

        @Override
        public synchronized void onSuccess() {
            succeeded = true;
            notifyAll();
        }

        @Override
        public synchronized void onClosed() {
            closed = true;
            notifyAll();
        }

        @Override
        public short getPriorityClass() {
            return 0;
        }

        @Override
        public synchronized void failOnDiskError(IOException e) {
            failed = true;
            notifyAll();
            System.err.println("Failed on disk error: " + e);
            e.printStackTrace();
        }

        @Override
        public void setSplitfileBlocks(int requiredBlocks, int remainingBlocks) {
            assertEquals(requiredBlocks, splitfile.dataBlocks.length);
            assertEquals(remainingBlocks, splitfile.checkBlocks.length);
        }

        @Override
        public void onSplitfileCompatibilityMode(
            CompatibilityMode min,
            CompatibilityMode max,
            byte[] customSplitfileKey,
            boolean compressed,
            boolean bottomLayer,
            boolean definitiveAnyway
        ) {
            // Ignore. FIXME?
        }

        @Override
        public void queueHeal(byte[] data, byte[] cryptoKey, byte cryptoAlgorithm) {
            assertArrayEquals(cryptoKey, splitfile.cryptoKey);
            assertEquals(cryptoAlgorithm, splitfile.cryptoAlgorithm);
            int x = -1;
            boolean progress = false;
            while ((x = splitfile.findCheckBlock(data, x)) != -1) {
                synchronized (this) {
                    encodedBlocks[x + splitfile.dataBlocks.length] = true;
                }
                progress = true;
            }
            if (!progress) {
                // Data block?
                while ((x = splitfile.findDataBlock(data, x)) != -1) {
                    synchronized (this) {
                        encodedBlocks[x] = true;
                    }
                    progress = true;
                }
            }
            if (!progress) {
                Assert.fail("Queued healing block not in the original block list");
            }
            assertTrue(progress);
        }

        synchronized void markDownloadedBlock(int block) {
            encodedBlocks[block] = true;
        }

        public synchronized void checkFailed() {
            assertFalse(failed);
        }

        public synchronized void waitForFinished() {
            while (!(succeeded || failed)) {
                try {
                    wait();
                } catch (InterruptedException e) {
                    // Ignore.
                }
            }
        }

        public synchronized void waitForFailed() {
            while (!(succeeded || failed)) {
                try {
                    wait();
                } catch (InterruptedException e) {
                    // Ignore.
                }
            }
            assertTrue(failed);
        }

        public void waitForFree() {
            synchronized (this) {
                while (!closed) {
                    try {
                        wait();
                    } catch (InterruptedException e) {
                        // Ignore.
                    }
                }
                assertTrue(succeeded);
            }
            synchronized (this) {
                for (int i = 0; i < encodedBlocks.length; i++) {
                    assertTrue("Block " + i + " not found or decoded", encodedBlocks[i]);
                }
            }
        }

        @Override
        public void onFetchedBlock() {
            // Ignore.
        }

        @Override
        public void fail(FetchException fe) {
            synchronized (this) {
                failed = true;
                notifyAll();
            }
        }

        @Override
        public void onFailedBlock() {
            // Ignore.
        }

        @Override
        public void maybeAddToBinaryBlob(ClientCHKBlock decodedBlock) {
            // Ignore.
        }

        @Override
        public boolean wantBinaryBlob() {
            return false;
        }

        @Override
        public BaseSendableGet getSendableGet() {
            return null;
        }

        @Override
        public void restartedAfterDataCorruption() {
            // Ignore
        }

        @Override
        public void clearCooldown() {
            // Ignore.
        }

        @Override
        public void reduceCooldown(long wakeupTime) {
            // Ignore.
        }

        @Override
        public HasKeyListener getHasKeyListener() {
            return null;
        }

        @Override
        public void failOnDiskError(ChecksumFailedException e) {
            synchronized (this) {
                failed = true;
                notifyAll();
            }
        }

        @Override
        public KeySalter getSalter() {
            return salt;
        }

        @Override
        public void onResume(int succeeded, int failed, ClientMetadata mimeType, long finalSize) {
            // Ignore.
        }

    }

    public static Bucket makeRandomBucket(long size) throws IOException {
        Bucket b = bf.makeBucket(size);
        BucketTools.fill(b, random, size);
        return b;
    }

    public static byte[][] splitAndPadBlocks(Bucket data, long size) throws IOException {
        int n = (int) ((size + BLOCK_SIZE - 1) / BLOCK_SIZE);
        byte[][] blocks = new byte[n][];
        InputStream is = data.getInputStream();
        DataInputStream dis = new DataInputStream(is);
        for (int i = 0; i < n; i++) {
            blocks[i] = new byte[BLOCK_SIZE];
            if (i < n - 1) {
                dis.readFully(blocks[i]);
            } else {
                int length = (int) (size - i * BLOCK_SIZE);
                dis.readFully(blocks[i], 0, length);
                // Now pad it ...
                blocks[i] = BucketTools.pad(blocks[i], BLOCK_SIZE, length);
            }
        }
        return blocks;
    }

    public static byte[] randomKey() {
        byte[] buf = new byte[KEY_LENGTH];
        random.nextBytes(buf);
        return buf;
    }

    public static boolean[] falseArray(int checkBlocks) {
        return new boolean[checkBlocks];
    }

    public static byte[][] constructBlocks(int n) {
        byte[][] blocks = new byte[n][];
        for (int i = 0; i < n; i++) {
            blocks[i] = new byte[BLOCK_SIZE];
        }
        return blocks;
    }

    // Actual tests ...

    @Test
    public void testSingleSegment() throws CHKEncodeException, IOException, FetchException, MetadataParseException, MetadataUnresolvedException {
        // 2 data blocks.
        //testSingleSegment(1, 2, BLOCK_SIZE);
        // We don't test this case because it just copies the data block to the check blocks.
        // Which breaks some of the scripts here.
        testSingleSegment(2, 1, BLOCK_SIZE * 2);
        testSingleSegment(2, 1, BLOCK_SIZE + 1);
        testSingleSegment(2, 2, BLOCK_SIZE * 2);
        testSingleSegment(2, 2, BLOCK_SIZE + 1);
        testSingleSegment(2, 3, BLOCK_SIZE * 2);
        testSingleSegment(2, 3, BLOCK_SIZE + 1);
        testSingleSegment(128, 128, BLOCK_SIZE * 128);
        testSingleSegment(128, 128, BLOCK_SIZE * 128 - 1);
        testSingleSegment(129, 127, BLOCK_SIZE * 129);
        testSingleSegment(129, 127, BLOCK_SIZE * 129 - 1);
        testSingleSegment(127, 129, BLOCK_SIZE * 127);
        testSingleSegment(127, 129, BLOCK_SIZE * 127 - 1);
    }

    private void testSingleSegment(int dataBlocks, int checkBlocks, long size) throws CHKEncodeException, IOException, FetchException, MetadataParseException, MetadataUnresolvedException {
        assertTrue(dataBlocks * (long) BLOCK_SIZE >= size);
        TestSplitfile test = TestSplitfile.constructSingleSegment(size, checkBlocks, false);
        testDataBlocksOnly(test);
        if (checkBlocks >= dataBlocks) {
            testCheckBlocksOnly(test);
        }
        testRandomMixture(test);
        test.free();
    }

    private void testDataBlocksOnly(TestSplitfile test) throws IOException, CHKEncodeException, FetchException, MetadataParseException {
        StorageCallback cb = test.createStorageCallback();
        SplitFileFetcherStorage storage = test.createStorage(cb);
        SplitFileFetcherSegmentStorage segment = storage.segments[0];
        for (int i = 0; i < test.checkBlocks.length; i++) {
            segment.onNonFatalFailure(test.dataBlocks.length + i);
        }
        for (int i = 0; i < test.dataBlocks.length; i++) {
            assertFalse(segment.hasStartedDecode());
            assertTrue(segment.onGotKey(test.dataKeys[i].getNodeCHK(), test.encodeDataBlock(i)));
            cb.markDownloadedBlock(i);
        }
        cb.checkFailed();
        assertTrue(segment.hasStartedDecode());
        cb.checkFailed();
        waitForDecode(segment);
        cb.checkFailed();
        cb.waitForFinished();
        cb.checkFailed();
        test.verifyOutput(storage);
        cb.checkFailed();
        storage.finishedFetcher();
        cb.checkFailed();
        waitForFinished(segment);
        cb.checkFailed();
        cb.waitForFree();
        cb.checkFailed();
    }

    private void testCheckBlocksOnly(TestSplitfile test) throws IOException, CHKEncodeException, FetchException, MetadataParseException {
        StorageCallback cb = test.createStorageCallback();
        SplitFileFetcherStorage storage = test.createStorage(cb);
        SplitFileFetcherSegmentStorage segment = storage.segments[0];
        for (int i = 0; i < test.dataBlocks.length; i++) {
            segment.onNonFatalFailure(i);
        }
        for (int i = test.dataBlocks.length; i < test.checkBlocks.length; i++) {
            segment.onNonFatalFailure(i + test.dataBlocks.length);
        }
        for (int i = 0; i < test.dataBlocks.length /* only need that many to decode */; i++) {
            assertFalse(segment.hasStartedDecode());
            assertTrue(segment.onGotKey(test.checkKeys[i].getNodeCHK(), test.encodeCheckBlock(i)));
            cb.markDownloadedBlock(i + test.dataBlocks.length);
        }
        cb.checkFailed();
        assertTrue(segment.hasStartedDecode());
        cb.checkFailed();
        waitForDecode(segment);
        cb.checkFailed();
        cb.waitForFinished();
        cb.checkFailed();
        test.verifyOutput(storage);
        cb.checkFailed();
        storage.finishedFetcher();
        cb.checkFailed();
        waitForFinished(segment);
        cb.checkFailed();
        cb.waitForFree();
        cb.checkFailed();
    }

    private void testRandomMixture(TestSplitfile test) throws FetchException, MetadataParseException, IOException, CHKEncodeException {
        StorageCallback cb = test.createStorageCallback();
        SplitFileFetcherStorage storage = test.createStorage(cb);
        SplitFileFetcherSegmentStorage segment = storage.segments[0];
        int total = test.dataBlocks.length + test.checkBlocks.length;
        for (int i = 0; i < total; i++)
            segment.onNonFatalFailure(i); // We want healing on all blocks that aren't found.
        boolean[] hits = new boolean[total];
        for (int i = 0; i < test.dataBlocks.length; i++) {
            int block;
            do {
                block = random.nextInt(total);
            } while (hits[block]);
            hits[block] = true;
            assertFalse(segment.hasStartedDecode());
            assertTrue(segment.onGotKey(test.getCHK(block), test.encodeBlock(block)));
            cb.markDownloadedBlock(block);
        }
        //printChosenBlocks(hits);
        cb.checkFailed();
        assertTrue(segment.hasStartedDecode());
        cb.checkFailed();
        waitForDecode(segment);
        cb.checkFailed();
        cb.waitForFinished();
        cb.checkFailed();
        test.verifyOutput(storage);
        cb.checkFailed();
        storage.finishedFetcher();
        cb.checkFailed();
        waitForFinished(segment);
        cb.checkFailed();
        cb.waitForFree();
        cb.checkFailed();
    }

    // FIXME LATER Test cross-segment.

    @Test
    public void testMultiSegment() throws CHKEncodeException, IOException, MetadataUnresolvedException, MetadataParseException, FetchException {
        // We have to be consistent with the format, but we can in fact play with the segment sizes 
        // to some degree.

        // Simplest case: Same number of blocks in each segment.
        // 2 blocks in each of 2 segments.
        testMultiSegment(32768 * 4, new int[]{2, 2}, new int[]{3, 3}, 2, 3, 0,
            InsertContext.CompatibilityMode.COMPAT_1416);
        testMultiSegment(32768 * 4 - 1, new int[]{2, 2}, new int[]{3, 3}, 2, 3, 0,
            InsertContext.CompatibilityMode.COMPAT_1416);

        // 3 blocks in 3 segments
        testMultiSegment(32768 * 9 - 1, new int[]{3, 3, 3}, new int[]{4, 4, 4}, 3, 4, 0,
            InsertContext.CompatibilityMode.COMPAT_1416);

        // Deduct blocks. This is how we handle this situation in modern splitfiles.
        testMultiSegment(32768 * 7 - 1, new int[]{3, 2, 2}, new int[]{4, 4, 4}, 3, 4, 2,
            InsertContext.CompatibilityMode.COMPAT_1416);

        // Sharp truncation. This is how we used to handle non-divisible numbers...
        testMultiSegment(32768 * 9 - 1, new int[]{7, 2}, new int[]{7, 2}, 7, 7, 0,
            InsertContext.CompatibilityMode.COMPAT_1416);
        // Still COMPAT_1416 because has crypto key etc.

        // FIXME test old splitfiles.
        // FIXME test really old splitfiles: last data block not padded
        // FIXME test really old splitfiles: non-redundant support
    }

    private void testMultiSegment(long size, int[] segmentDataBlockCount,
                                  int[] segmentCheckBlockCount, int segmentSize, int checkSegmentSize,
                                  int deductBlocksFromSegments, CompatibilityMode topCompatibilityMode) throws CHKEncodeException, IOException, MetadataUnresolvedException, MetadataParseException, FetchException {
        TestSplitfile test = TestSplitfile.constructMultipleSegments(size, segmentDataBlockCount,
            segmentCheckBlockCount, segmentSize, checkSegmentSize, deductBlocksFromSegments,
            topCompatibilityMode);
        testRandomMixtureMultiSegment(test);
        test.free();

    }

    private void testRandomMixtureMultiSegment(TestSplitfile test) throws CHKEncodeException, IOException, FetchException, MetadataParseException {
        StorageCallback cb = test.createStorageCallback();
        SplitFileFetcherStorage storage = test.createStorage(cb);
        int total = test.dataBlocks.length + test.checkBlocks.length;
        for (SplitFileFetcherSegmentStorage segment : storage.segments) {
            for (int i = 0; i < segment.totalBlocks(); i++)
                segment.onNonFatalFailure(i); // We want healing on all blocks that aren't found.
        }
        boolean[] hits = new boolean[total];
        for (int i = 0; i < test.dataBlocks.length; i++) {
            int block;
            do {
                block = random.nextInt(total);
            } while (hits[block]);
            hits[block] = true;
            SplitFileFetcherSegmentStorage segment = storage.segments[test.segmentFor(block)];
            if (segment.hasStartedDecode()) {
                i--;
                continue;
            }
            assertTrue(segment.onGotKey(test.getCHK(block), test.encodeBlock(block)));
            cb.markDownloadedBlock(block);
        }
        printChosenBlocks(hits);
        cb.checkFailed();
        for (SplitFileFetcherSegmentStorage segment : storage.segments) {
            assertTrue(segment.hasStartedDecode()); // All segments have started decoding.
        }
        cb.checkFailed();
        for (SplitFileFetcherSegmentStorage segment : storage.segments) {
            waitForDecode(segment);
        }
        cb.checkFailed();
        cb.waitForFinished();
        cb.checkFailed();
        test.verifyOutput(storage);
        cb.checkFailed();
        storage.finishedFetcher();
        cb.checkFailed();
        for (SplitFileFetcherSegmentStorage segment : storage.segments) {
            waitForFinished(segment);
        }
        cb.checkFailed();
        cb.waitForFree();
        cb.checkFailed();
    }

    private void printChosenBlocks(boolean[] hits) {
        StringBuilder sb = new StringBuilder();
        sb.append("Blocks: ");
        for (int i = 0; i < hits.length; i++) {
            if (hits[i]) {
                sb.append(i);
                sb.append(" ");
            }
        }
        sb.setLength(sb.length() - 1);
        System.out.println(sb.toString());
    }

    private void waitForFinished(SplitFileFetcherSegmentStorage segment) {
        while (!segment.isFinished()) {
            assertFalse(segment.hasFailed());
            try {
                Thread.sleep(1);
            } catch (InterruptedException e) {
                // Ignore.
            }
        }
    }

    private void waitForDecode(SplitFileFetcherSegmentStorage segment) {
        while (!segment.hasSucceeded()) {
            assertFalse(segment.hasFailed());
            try {
                Thread.sleep(1);
            } catch (InterruptedException e) {
                // Ignore.
            }
        }
    }

    private static class MyKeysFetchingLocally implements KeysFetchingLocally {
        private final HashSet<Key> keys = new HashSet<Key>();

        @Override
        public long checkRecentlyFailed(Key key, boolean realTime) {
            return 0;
        }

        @Override
        public boolean hasKey(Key key, BaseSendableGet getterWaiting) {
            return keys.contains(key);
        }

        @Override
        public boolean hasInsert(SendableRequestItemKey token) {
            return false;
        }

        public void add(Key k) {
            keys.add(k);
        }

        public void clear() {
            keys.clear();
        }

    }

    @Test
    public void testChooseKeyOneTry() throws CHKEncodeException, IOException, MetadataUnresolvedException, MetadataParseException, FetchException {
        int dataBlocks = 3, checkBlocks = 3;
        TestSplitfile test = TestSplitfile.constructSingleSegment(dataBlocks * BLOCK_SIZE, checkBlocks, false);
        StorageCallback cb = test.createStorageCallback();
        FetchContext ctx = test.makeFetchContext();
        ctx.maxSplitfileBlockRetries = 0;
        SplitFileFetcherStorage storage = test.createStorage(cb, ctx);
        boolean[] tried = new boolean[dataBlocks + checkBlocks];
        innerChooseKeyTest(dataBlocks, checkBlocks, storage.segments[0], tried, test, false);
        assertNull(storage.chooseRandomKey());
        cb.waitForFailed();
    }

    private void innerChooseKeyTest(int dataBlocks, int checkBlocks, SplitFileFetcherSegmentStorage storage, boolean[] tried, TestSplitfile test, boolean cooldown) {
        final MyKeysFetchingLocally keys = test.fetchingKeys;
        keys.clear();
        for (int i = 0; i < dataBlocks + checkBlocks; i++) {
            int chosen = storage.chooseRandomKey();
            assertTrue(chosen != -1);
            assertFalse(tried[chosen]);
            tried[chosen] = true;
            Key k = test.getCHK(chosen);
            keys.add(k);
        }
        // Every block has been tried.
        for (boolean b : tried) {
            assertTrue(b);
        }
        if (cooldown) {
            assertEquals(storage.chooseRandomKey(), -1);
            // In infinite cooldown, waiting for all requests to complete.
            assertEquals(storage.getOverallCooldownTime(), Long.MAX_VALUE);
        }
        // Every request is running.
        // When we complete a request, we remove it from KeysFetchingLocally *and* call onNonFatalFailure.
        // This will reset the cooldown.
        for (int i = 0; i < dataBlocks + checkBlocks; i++) {
            storage.onNonFatalFailure(i);
        }
        if (!cooldown) {
            // Cleared cooldown, if any.
            assertEquals(storage.getOverallCooldownTime(), 0);
        }
        assertTrue(storage.getOverallCooldownTime() != Long.MAX_VALUE || storage.hasFailed());
        // Now all the requests have completed, the keys are no longer being fetched.
        keys.clear();
        // Will be able to fetch keys immediately, unless in cooldown.
    }

    @Test
    public void testChooseKeyThreeTries() throws CHKEncodeException, IOException, MetadataUnresolvedException, MetadataParseException, FetchException {
        int dataBlocks = 3, checkBlocks = 3;
        TestSplitfile test = TestSplitfile.constructSingleSegment(dataBlocks * BLOCK_SIZE, checkBlocks, false);
        StorageCallback cb = test.createStorageCallback();
        FetchContext ctx = test.makeFetchContext();
        ctx.maxSplitfileBlockRetries = 2;
        SplitFileFetcherStorage storage = test.createStorage(cb, ctx);
        for (int i = 0; i < 3; i++) {
            boolean[] tried = new boolean[dataBlocks + checkBlocks];
            innerChooseKeyTest(dataBlocks, checkBlocks, storage.segments[0], tried, test, false);
        }
        assertNull(storage.chooseRandomKey());
        cb.waitForFailed();
    }

    @Test
    public void testChooseKeyCooldown() throws CHKEncodeException, IOException, MetadataUnresolvedException, MetadataParseException, FetchException, InterruptedException {
        int dataBlocks = 3, checkBlocks = 3;
        int COOLDOWN_TIME = 200;
        TestSplitfile test = TestSplitfile.constructSingleSegment(dataBlocks * BLOCK_SIZE, checkBlocks, false);
        StorageCallback cb = test.createStorageCallback();
        FetchContext ctx = test.makeFetchContext();
        ctx.maxSplitfileBlockRetries = 5;
        ctx.setCooldownRetries(3);
        ctx.setCooldownTime(COOLDOWN_TIME, true);
        SplitFileFetcherStorage storage = test.createStorage(cb, ctx);
        // 3 tries for each block.
        long now = System.currentTimeMillis();
        for (int i = 0; i < 3; i++) {
            boolean[] tried = new boolean[dataBlocks + checkBlocks];
            innerChooseKeyTest(dataBlocks, checkBlocks, storage.segments[0], tried, test, true);
        }
        assertTrue(storage.segments[0].getOverallCooldownTime() > now);
        assertNotEquals(Long.MAX_VALUE, storage.segments[0].getOverallCooldownTime());
        // Now in cooldown.
        test.fetchingKeys.clear();
        assertNull(storage.chooseRandomKey());
        Thread.sleep(COOLDOWN_TIME + COOLDOWN_TIME / 2 + 1);
        cb.checkFailed();
        // Should be out of cooldown now.
        for (int i = 0; i < 3; i++) {
            boolean[] tried = new boolean[dataBlocks + checkBlocks];
            innerChooseKeyTest(dataBlocks, checkBlocks, storage.segments[0], tried, test, true);
        }
        // Now it should fail.
        cb.waitForFailed();
    }

    @Test
    public void testWriteReadSegmentKeys() throws FetchException, MetadataParseException, IOException, CHKEncodeException, MetadataUnresolvedException, ChecksumFailedException {
        int dataBlocks = 3, checkBlocks = 3;
        TestSplitfile test = TestSplitfile.constructSingleSegment(dataBlocks*BLOCK_SIZE, checkBlocks, true);
        StorageCallback cb = test.createStorageCallback();
        SplitFileFetcherStorage storage = test.createStorage(cb);
        SplitFileFetcherSegmentStorage segment = storage.segments[0];
        SplitFileSegmentKeys keys = segment.getSegmentKeys();
        SplitFileSegmentKeys moreKeys = segment.readSegmentKeys();
        assertEquals(keys, moreKeys);
        storage.close();
    }

    /**
     * Test persistence: Create and then reload. Don't do anything.
     */
    @Test
    public void testPersistenceReload() throws CHKEncodeException, IOException, MetadataUnresolvedException, MetadataParseException, FetchException, StorageFormatException {
        int dataBlocks = 2;
        int checkBlocks = 3;
        long size = 32768 * 2 - 1;
        assertTrue(dataBlocks * (long) BLOCK_SIZE >= size);
        TestSplitfile test = TestSplitfile.constructSingleSegment(size, checkBlocks, true);
        StorageCallback cb = test.createStorageCallback();
        SplitFileFetcherStorage storage = createSplitFileFetcherStorageTwice(test, cb);
        storage.close();
    }

    @Test
    public void testPersistenceReloadThenFetch() throws IOException, StorageFormatException, CHKEncodeException, MetadataUnresolvedException, MetadataParseException, FetchException {
        int dataBlocks = 2;
        int checkBlocks = 3;
        long size = 32768 * 2 - 1;
        assertTrue(dataBlocks * (long) BLOCK_SIZE >= size);
        TestSplitfile test = TestSplitfile.constructSingleSegment(size, checkBlocks, true);
        StorageCallback cb = test.createStorageCallback();
        SplitFileFetcherStorage storage = createSplitFileFetcherStorageTwice(test, cb);
        SplitFileFetcherSegmentStorage segment = storage.segments[0];
        assertFalse(segment.corruptMetadata());
        int total = test.dataBlocks.length + test.checkBlocks.length;
        for (int i = 0; i < total; i++) {
            segment.onNonFatalFailure(i); // We want healing on all blocks that aren't found.
        }
        boolean[] hits = new boolean[total];
        for (int i = 0; i < test.dataBlocks.length; i++) {
            int block;
            do {
                block = random.nextInt(total);
            } while (hits[block]);
            hits[block] = true;
            assertFalse(segment.hasStartedDecode());
            assertTrue(segment.onGotKey(test.getCHK(block), test.encodeBlock(block)));
            cb.markDownloadedBlock(block);
        }
        //printChosenBlocks(hits);
        cb.checkFailed();
        assertTrue(segment.hasStartedDecode());
        cb.checkFailed();
        waitForDecode(segment);
        cb.checkFailed();
        cb.waitForFinished();
        cb.checkFailed();
        test.verifyOutput(storage);
        cb.checkFailed();
        storage.finishedFetcher();
        cb.checkFailed();
        waitForFinished(segment);
        cb.checkFailed();
        cb.waitForFree();
        cb.checkFailed();
    }

    @Test
    public void testPersistenceReloadThenChooseKey() throws IOException, StorageFormatException, CHKEncodeException, MetadataUnresolvedException, MetadataParseException, FetchException {
        int dataBlocks = 2;
        int checkBlocks = 3;
        long size = 32768 * 2 - 1;
        assertTrue(dataBlocks * (long) BLOCK_SIZE >= size);
        TestSplitfile test = TestSplitfile.constructSingleSegment(size, checkBlocks, true);
        StorageCallback cb = test.createStorageCallback();
        FetchContext ctx = test.makeFetchContext();
        ctx.maxSplitfileBlockRetries = 2;
        test.createStorage(cb, ctx);
        // No need to shutdown the old storage.
        SplitFileFetcherStorage storage = test.createStorage(cb, ctx, cb.getRAF());
        for (int i = 0; i < 3; i++) {
            boolean[] tried = new boolean[dataBlocks + checkBlocks];
            innerChooseKeyTest(dataBlocks, checkBlocks, storage.segments[0], tried, test, false);
        }
        test.fetchingKeys.clear();
        assertNull(storage.chooseRandomKey());
        cb.waitForFailed();
    }

    @Test
    public void testPersistenceReloadBetweenChooseKey() throws IOException, StorageFormatException, CHKEncodeException, MetadataUnresolvedException, MetadataParseException, FetchException {
        int dataBlocks = 2;
        int checkBlocks = 3;
        long size = 32768 * 2 - 1;
        assertTrue(dataBlocks * (long) BLOCK_SIZE >= size);
        TestSplitfile test = TestSplitfile.constructSingleSegment(size, checkBlocks, true);
        StorageCallback cb = test.createStorageCallback();
        FetchContext ctx = test.makeFetchContext();
        ctx.maxSplitfileBlockRetries = 2;
        test.createStorage(cb, ctx);
        // No need to shutdown the old storage.
        SplitFileFetcherStorage storage = test.createStorage(cb, ctx, cb.getRAF());
        for (int i = 0; i < 3; i++) {
            boolean[] tried = new boolean[dataBlocks + checkBlocks];
            innerChooseKeyTest(dataBlocks, checkBlocks, storage.segments[0], tried, test, false);
            // Reload.
            exec.waitForIdle();
            try {
                storage = test.createStorage(cb, ctx, cb.getRAF());
                assertTrue(i != 2);
                storage.start(false);
            } catch (FetchException e) {
                if (i != 2) throw e; // Already failed on the final iteration, otherwise is an error.
                return;
            }
        }
        test.fetchingKeys.clear();
        assertNull(storage.chooseRandomKey());
        cb.waitForFailed();
    }

    @Test
    public void testPersistenceReloadBetweenFetches() throws IOException, StorageFormatException, CHKEncodeException, MetadataUnresolvedException, MetadataParseException, FetchException {
        int dataBlocks = 2;
        int checkBlocks = 3;
        long size = 32768 * 2 - 1;
        assertTrue(dataBlocks * (long) BLOCK_SIZE >= size);
        TestSplitfile test = TestSplitfile.constructSingleSegment(size, checkBlocks, true);
        StorageCallback cb = test.createStorageCallback();
        SplitFileFetcherStorage storage = createSplitFileFetcherStorageTwice(test, cb);
        SplitFileFetcherSegmentStorage segment = storage.segments[0];
        assertFalse(segment.corruptMetadata());
        int total = test.dataBlocks.length + test.checkBlocks.length;
        for (int i = 0; i < total; i++) {
            segment.onNonFatalFailure(i); // We want healing on all blocks that aren't found.
        }
        boolean[] hits = new boolean[total];
        for (int i = 0; i < test.dataBlocks.length; i++) {
            int block;
            do {
                block = random.nextInt(total);
            } while (hits[block]);
            hits[block] = true;
            assertFalse(segment.hasStartedDecode());
            assertTrue(segment.onGotKey(test.getCHK(block), test.encodeBlock(block)));
            cb.markDownloadedBlock(block);
            if (i != test.dataBlocks.length - 1) {
                // Reload.
                exec.waitForIdle();
                storage = test.createStorage(cb, test.makeFetchContext(), cb.getRAF());
                segment = storage.segments[0];
                storage.start(false);
            }
        }
        //printChosenBlocks(hits);
        cb.checkFailed();
        assertTrue(segment.hasStartedDecode());
        cb.checkFailed();
        waitForDecode(segment);
        cb.checkFailed();
        cb.waitForFinished();
        cb.checkFailed();
        test.verifyOutput(storage);
        cb.checkFailed();
        storage.finishedFetcher();
        cb.checkFailed();
        waitForFinished(segment);
        cb.checkFailed();
        cb.waitForFree();
        cb.checkFailed();
    }

    private SplitFileFetcherStorage createSplitFileFetcherStorageTwice(TestSplitfile test, StorageCallback cb) throws FetchException, MetadataParseException, IOException, StorageFormatException {
        test.createStorage(cb);
        // No need to shutdown the old storage.
        return test.createStorage(cb, test.makeFetchContext(), cb.getRAF());
    }
}
