package peergos.shared.storage;

import peergos.shared.cbor.CborObject;
import peergos.shared.crypto.hash.PublicKeyHash;
import peergos.shared.io.ipfs.multihash.Multihash;
import peergos.shared.user.fs.FragmentWithHash;
import peergos.shared.util.ProgressConsumer;

import java.util.List;
import java.util.Optional;
import java.util.Random;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ScheduledThreadPoolExecutor;
import java.util.concurrent.TimeUnit;
import java.util.function.Function;
import java.util.function.Supplier;

public class RetryStorage implements ContentAddressedStorage {

    private static final int RANDOM_SEED = 987447;
    private static Random random = new Random(RANDOM_SEED);
    private final ContentAddressedStorage target;
    private final ScheduledThreadPoolExecutor executor = new ScheduledThreadPoolExecutor(1);
    private final int maxAttempts;


    public RetryStorage(ContentAddressedStorage target, int maxAttempts) {
        this.target = target;
        this.maxAttempts = maxAttempts;
    }

    private <V> void retryAfter(Supplier<CompletableFuture<V>> method, int milliseconds) {
        executor.schedule(method::get, milliseconds, TimeUnit.MILLISECONDS);
    }

    private int jitter(int minMilliseconds, int rangeMilliseconds) {
        return minMilliseconds + random.nextInt(rangeMilliseconds);
    }

    private <V> CompletableFuture<V> runWithRetry(Supplier<CompletableFuture<V>> f) {
        return recurse(maxAttempts, f);
    }

    private <V> CompletableFuture<V> recurse(int retriesLeft, Supplier<CompletableFuture<V>> f) {
        Function<Integer, CompletableFuture<V>> compose =
                i-> {
                    CompletableFuture<V> res = new CompletableFuture<>();
                    f.get()
                            .thenAccept(res::complete)
                            .exceptionally(e ->  {
                                if (i == 1) {
                                    res.completeExceptionally(e);
                                } else {
                                    retryAfter(() -> recurse(i - 1, f)
                                            .thenAccept(res::complete)
                                            .exceptionally(t -> {
                                                res.completeExceptionally(t);
                                                return null;
                                            }),
                                            jitter((maxAttempts + 1 - i) * 1000, 500));
                                }
                                return null;
                            });
                    return res;
                };
        return compose.apply(retriesLeft);
    }
    @Override
    public CompletableFuture<BlockStoreProperties> blockStoreProperties() {
        return runWithRetry(() -> target.blockStoreProperties());
    }
    @Override
    public CompletableFuture<Multihash> id() {
        return runWithRetry(() -> target.id());
    }

    @Override
    public CompletableFuture<TransactionId> startTransaction(PublicKeyHash owner) {
        return runWithRetry(() -> target.startTransaction(owner));
    }

    @Override
    public CompletableFuture<Boolean> closeTransaction(PublicKeyHash owner, TransactionId tid) {
        return runWithRetry(() -> target.closeTransaction(owner, tid));
    }

    @Override
    public CompletableFuture<List<Multihash>> put(PublicKeyHash owner, PublicKeyHash writer, List<byte[]> signatures, List<byte[]> blocks, TransactionId tid) {
        return runWithRetry(() -> target.put(owner, writer, signatures, blocks, tid));
    }

    @Override
    public CompletableFuture<Optional<CborObject>> get(Multihash hash) {
        return runWithRetry(() -> target.get(hash));
    }

    @Override
    public CompletableFuture<List<Multihash>> putRaw(PublicKeyHash owner,
                                                     PublicKeyHash writer,
                                                     List<byte[]> signatures,
                                                     List<byte[]> blocks,
                                                     TransactionId tid,
                                                     ProgressConsumer<Long> progressCounter) {
        return runWithRetry(() -> target.putRaw(owner, writer, signatures, blocks, tid, progressCounter));
    }

    @Override
    public CompletableFuture<Optional<byte[]>> getRaw(Multihash hash) {
        return runWithRetry(() -> target.getRaw(hash));
    }

    @Override
    public CompletableFuture<List<Multihash>> pinUpdate(PublicKeyHash owner, Multihash existing, Multihash updated) {
        return runWithRetry(() -> target.pinUpdate(owner, existing, updated));
    }

    @Override
    public CompletableFuture<List<Multihash>> recursivePin(PublicKeyHash owner, Multihash hash) {
        return runWithRetry(() -> target.recursivePin(owner, hash));
    }

    @Override
    public CompletableFuture<List<Multihash>> recursiveUnpin(PublicKeyHash owner, Multihash hash) {
        return runWithRetry(() -> target.recursiveUnpin(owner, hash));
    }

    @Override
    public CompletableFuture<Boolean> gc() {
        return runWithRetry(() -> target.gc());
    }

    @Override
    public CompletableFuture<List<Multihash>> getLinks(Multihash root) {
        return runWithRetry(() -> target.getLinks(root));
    }

    @Override
    public CompletableFuture<Optional<Integer>> getSize(Multihash block) {
        return runWithRetry(() -> target.getSize(block));
    }

    @Override
    public CompletableFuture<List<FragmentWithHash>> downloadFragments(List<Multihash> hashes,
                                                                       ProgressConsumer<Long> monitor,
                                                                       double spaceIncreaseFactor) {
        return runWithRetry(() -> target.downloadFragments(hashes, monitor, spaceIncreaseFactor));
    }

    @Override
    public CompletableFuture<List<PresignedUrl>> authReads(List<Multihash> blocks) {
        return runWithRetry(() -> target.authReads(blocks));
    }

    @Override
    public CompletableFuture<List<PresignedUrl>> authWrites(PublicKeyHash owner,
                                                            PublicKeyHash writer,
                                                            List<byte[]> signedHashes,
                                                            List<Integer> blockSizes,
                                                            boolean isRaw,
                                                            TransactionId tid) {
        return runWithRetry(() -> target.authWrites(owner, writer, signedHashes, blockSizes, isRaw, tid));
    }
}

@@ -1610,100 +1610,8 @@
 ) %7B%0A
-        Function%3CInteger, CompletableFuture%3CV%3E%3E compose =%0A                i-%3E %7B%0A            
     
@@ -1668,32 +1668,38 @@
 re%3C%3E();%0A        
+try %7B%0A
             f.ge
@@ -1718,32 +1718,24 @@
             
-       
- 
 .thenAccept(
@@ -1765,32 +1765,24 @@
             
-        
 .exceptional
@@ -1788,17 +1788,16 @@
 lly(e -%3E
- 
  %7B%0A     
@@ -1819,21 +1819,23 @@
     
- 
-       if (i
+if (retriesLeft
  == 
@@ -1839,24 +1839,16 @@
 == 1) %7B%0A
-        
         
@@ -1921,24 +1921,16 @@
         
-      
-  
 %7D else %7B
@@ -1930,24 +1930,16 @@
  else %7B%0A
-        
         
@@ -1983,17 +1983,27 @@
 recurse(
-i
+retriesLeft
  - 1, f)
@@ -2363,24 +2363,16 @@
         
-        
 jitter((
@@ -2389,17 +2389,27 @@
 s + 1 - 
-i
+retriesLeft
 ) * 1000
@@ -2445,24 +2445,10 @@
     
-        %7D%0A      
+%7D%0A
     
@@ -2459,34 +2459,32 @@
                 
-  
 return null;%0A   
@@ -2492,32 +2492,36 @@
                 
+%7D);%0A
         %7D);%0A    
@@ -2517,41 +2517,71 @@
    %7D
-);%0A                    return res
+ catch (Throwable t) %7B%0A            res.completeExceptionally(t)
 ;%0A  
@@ -2590,18 +2590,9 @@
     
- 
-       %7D;
+%7D
 %0A   
@@ -2607,34 +2607,11 @@
 urn 
-compose.apply(retriesLeft)
+res
 ;%0A  

Simplify RetryStorage