package peergos.server.storage;

import peergos.server.storage.auth.*;
import peergos.shared.cbor.*;
import peergos.shared.crypto.hash.*;
import peergos.shared.io.ipfs.cid.*;
import peergos.shared.io.ipfs.multihash.*;
import peergos.shared.storage.*;
import peergos.shared.storage.auth.*;
import peergos.shared.util.*;

import java.util.*;
import java.util.concurrent.*;

public class AuthedCachingStorage extends DelegatingStorage {
    private final ContentAddressedStorage target;
    private final LRUCache<Multihash, byte[]> cache;
    private final LRUCache<Multihash, CompletableFuture<Optional<CborObject>>> pending;
    private final LRUCache<Multihash, CompletableFuture<Optional<byte[]>>> pendingRaw;
    private final BlockRequestAuthoriser authoriser;
    private final Hasher h;
    private final Cid ourNodeId;
    private final int maxValueSize, cacheSize;

    public AuthedCachingStorage(ContentAddressedStorage target,
                                BlockRequestAuthoriser authoriser,
                                Hasher h,
                                int cacheSize,
                                int maxValueSize) {
        super(target);
        this.target = target;
        this.ourNodeId = target.id().join();
        this.authoriser = authoriser;
        this.h = h;
        this.cache = new LRUCache<>(cacheSize);
        this.maxValueSize = maxValueSize;
        this.cacheSize = cacheSize;
        this.pending = new LRUCache<>(100);
        this.pendingRaw = new LRUCache<>(100);
    }

    public Collection<byte[]> getCached() {
        return cache.values();
    }

    @Override
    public CompletableFuture<BlockStoreProperties> blockStoreProperties() {
        return target.blockStoreProperties();
    }

    @Override
    public ContentAddressedStorage directToOrigin() {
        return new AuthedCachingStorage(target.directToOrigin(), authoriser, h, cacheSize, maxValueSize);
    }

    @Override
    public void clearBlockCache() {
        cache.clear();
        target.clearBlockCache();
    }

    @Override
    public CompletableFuture<List<Cid>> put(PublicKeyHash owner,
                                            PublicKeyHash writer,
                                            List<byte[]> signedHashes,
                                            List<byte[]> blocks,
                                            TransactionId tid) {
        return target.put(owner, writer, signedHashes, blocks, tid)
                .thenApply(res -> {
                    for (int i=0; i < blocks.size(); i++) {
                        byte[] block = blocks.get(i);
                        if (block.length < maxValueSize)
                            cache.put(res.get(i), block);
                    }
                    return res;
                });
    }

    private CompletableFuture<byte[]> authoriseGet(Cid key, byte[] block, Optional<BatWithId> bat) {
        return bat.map(b -> b.bat.generateAuth(key, ourNodeId, 300, S3Request.currentDatetime(), bat.get().id, h)
                .thenApply(BlockAuth::encode)).orElse(Futures.of(""))
                .thenCompose(auth -> authoriser.allowRead(key, block, ourNodeId, auth))
                .thenCompose(allow -> allow ? Futures.of(block) : Futures.errored(new Throwable("Unauthorised!")));
    }

    @Override
    public CompletableFuture<Optional<CborObject>> get(Cid key, Optional<BatWithId> bat) {
        if (cache.containsKey(key))
            return authoriseGet(key, cache.get(key), bat)
                    .thenApply(res -> Optional.of(CborObject.fromByteArray(res)));

        if (pending.containsKey(key))
            return pending.get(key)
                    .thenCompose(copt -> copt.isEmpty() ?
                            Futures.of(Optional.empty()) :
                            authoriseGet(key, copt.get().serialize(), bat)
                                    .thenApply(b -> copt));

        CompletableFuture<Optional<CborObject>> pipe = new CompletableFuture<>();
        pending.put(key, pipe);

        CompletableFuture<Optional<CborObject>> result = new CompletableFuture<>();
        target.get(key, bat).thenAccept(cborOpt -> {
            if (cborOpt.isPresent()) {
                byte[] value = cborOpt.get().toByteArray();
                if (value.length > 0 && value.length < maxValueSize)
                    cache.put(key, value);
            }
            pending.remove(key);
            pipe.complete(cborOpt);
            result.complete(cborOpt);
        }).exceptionally(t -> {
            pending.remove(key);
            pipe.completeExceptionally(t);
            result.completeExceptionally(t);
            return null;
        });
        return result;
    }

    @Override
    public CompletableFuture<List<Cid>> putRaw(PublicKeyHash owner,
                                               PublicKeyHash writer,
                                               List<byte[]> signatures,
                                               List<byte[]> blocks,
                                               TransactionId tid,
                                               ProgressConsumer<Long> progressConsumer) {
        return target.putRaw(owner, writer, signatures, blocks, tid, progressConsumer)
                .thenApply(res -> {
                    for (int i=0; i < blocks.size(); i++) {
                        byte[] block = blocks.get(i);
                        if (block.length < maxValueSize)
                            cache.put(res.get(i), block);
                    }
                    return res;
                });
    }

    @Override
    public CompletableFuture<Optional<byte[]>> getRaw(Cid key, Optional<BatWithId> bat) {
        if (cache.containsKey(key))
            return authoriseGet(key, cache.get(key), bat)
                    .thenApply(res -> Optional.of(res));

        if (pendingRaw.containsKey(key))
            return pendingRaw.get(key)
                    .thenCompose(opt -> opt.isEmpty() ?
                            Futures.of(Optional.empty()) :
                            authoriseGet(key, opt.get(), bat)
                                    .thenApply(b -> opt));

        CompletableFuture<Optional<byte[]>> pipe = new CompletableFuture<>();
        pendingRaw.put(key, pipe);
        return target.getRaw(key, bat).thenApply(rawOpt -> {
            if (rawOpt.isPresent()) {
                byte[] value = rawOpt.get();
                if (value.length > 0 && value.length < maxValueSize)
                    cache.put(key, value);
            }
            pendingRaw.remove(key);
            pipe.complete(rawOpt);
            return rawOpt;
        }).exceptionally(t -> {
            pending.remove(key);
            pipe.completeExceptionally(t);
            return Optional.empty();
        });
    }
}

@@ -530,16 +530,77 @@
  cache;%0A
+    private final LRUCache%3CMultihash, Boolean%3E legacyBlocks;%0A
     priv
@@ -1400,24 +1400,79 @@
 cacheSize);%0A
+        this.legacyBlocks = new LRUCache%3C%3E(cacheSize);%0A
         this
@@ -3008,32 +3008,145 @@
 tWithId%3E bat) %7B%0A
+        if (key.isRaw() && bat.isEmpty() && legacyBlocks.containsKey(key))%0A            return Futures.of(block);%0A
         return b
@@ -6708,32 +6708,34 @@
  %3C maxValueSize)
+ %7B
 %0A               
@@ -6754,32 +6754,157 @@
 ut(key, value);%0A
+                    if (bat.isEmpty() && key.isRaw())%0A                        legacyBlocks.put(key, true);%0A                %7D%0A
             %7D%0A  

Make sure the AuthedCache is consistent with new legacy raw blocks additions