package com.tradeshift.reaktive.actors;

import java.util.UUID;
import java.util.function.Function;

import akka.actor.ActorRef;
import akka.actor.ActorSystem;
import akka.actor.Props;
import akka.cluster.sharding.ClusterSharding;
import akka.cluster.sharding.ClusterShardingSettings;
import akka.cluster.sharding.ShardRegion.MessageExtractor;
import akka.persistence.AbstractPersistentActor;

/**
 * Base class for setting up sharding of persistent actors that:
 * 
 * - Have their persistenceId conform to [prefix] + "_" + [uuid], e.g. "doc_249e1098-1cd9-4ffe-a494-73a430983590"
 * - Respond to a specific base class of commands
 * - Send the commands unchanged from the SharedRegion router onto the persistent actors themselves
 */
public class PersistentActorSharding<A extends AbstractPersistentActor, C> {
    private final Class<A> actorType;
    private final String entityIdPrefix;
    private final Function<C, UUID> entityIdPostfix;
    
    public static <A extends AbstractPersistentActor, C> PersistentActorSharding<A,C> of(Class<A> actorType, String entityIdPrefix, Function<C, UUID> entityIdPostfix) {
        return new PersistentActorSharding<>(actorType, entityIdPrefix, entityIdPostfix);
    }
    
    protected PersistentActorSharding(Class<A> actorType, String entityIdPrefix, Function<C, UUID> entityIdPostfix) {
        this.actorType = actorType;
        this.entityIdPrefix = entityIdPrefix;
        this.entityIdPostfix = entityIdPostfix;
    }

    private final MessageExtractor messageExtractor = new MessageExtractor() {
        private final int numberOfShards = 256;
        
        @Override
        public String entityId(Object command) {
            return getEntityId(command);
        }

        @Override
        public String shardId(Object command) {
            return String.valueOf(entityId(command).hashCode() % numberOfShards);
        }
        
        @Override
        public Object entityMessage(Object command) {
            // we don't need need to unwrap messages sent to the router, they can just be
            // forwarded to the target persistent actor directly.
            return command;
        }
    };
    
    /**
     * Starts the cluster router (ShardRegion) for this persistent actor type on the given actor system,
     * and returns its ActorRef. If it's already running, just returns the ActorRef.
     */
    public ActorRef shardRegion(ActorSystem system) {
        return ClusterSharding.get(system).start(
            entityIdPrefix, 
            Props.create(actorType), 
            ClusterShardingSettings.create(system), 
            messageExtractor);
    }
    
    /** 
     * Returns the UUID from a generated persistence ID.
     * @param persistenceId A persistenceId of an actor that was spawned by sending a command to it through the 
     * ActorRef returned by {@link #shardRegion}. 
     */
    public UUID getUUIDFromPersistenceId(String persistenceId) {
        return UUID.fromString(persistenceId.substring(entityIdPrefix.length() + 1));
    };
    
    @SuppressWarnings("unchecked")
    protected String getEntityId(Object command) {
        return entityIdPrefix + "_" + entityIdPostfix.apply((C) command);
    }
    
}

@@ -769,43 +769,8 @@
 ing%3C
-A extends AbstractPersistentActor, 
 C%3E %7B
@@ -788,34 +788,27 @@
 e final 
-Class%3CA%3E actorType
+Props props
 ;%0A    pr
@@ -896,24 +896,125 @@
 stfix;%0A    %0A
+    /**%0A     * Creates a PersistentActorSharding for an actor that has a no-args constructor%0A     */%0A
     public s
@@ -1086,10 +1086,8 @@
 ing%3C
-A,
 C%3E o
@@ -1213,16 +1213,29 @@
 rding%3C%3E(
+Props.create(
 actorTyp
@@ -1235,16 +1235,17 @@
 ctorType
+)
 , entity
@@ -1291,60 +1291,448 @@
     
-protected PersistentActorSharding(Class%3CA%3E actorType
+/**%0A     * Creates a PersistentActorSharding for an actor that is created according to %5Bprops%5D. The actor must be a subclass of %7B@link AbstractPersistentActor%7D.%0A     */%0A    public static %3CC%3E PersistentActorSharding%3CC%3E of(Props props, String entityIdPrefix, Function%3CC, UUID%3E entityIdPostfix) %7B%0A        return new PersistentActorSharding%3C%3E(props, entityIdPrefix, entityIdPostfix);%0A    %7D%0A    %0A    protected PersistentActorSharding(Props props
 , St
@@ -1806,29 +1806,21 @@
 his.
-actorType = actorType
+props = props
 ;%0A  
@@ -2969,39 +2969,21 @@
         
-Props.create(actorType)
+props
 , %0A     

Allow PersistentActorSharding for actors that have non-empty constructors