package edu.uw.zookeeper.safari.backend;

import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;

import com.google.common.base.Function;
import com.google.common.eventbus.Subscribe;
import com.google.common.util.concurrent.ListenableFuture;

import edu.uw.zookeeper.common.Promise;
import edu.uw.zookeeper.common.PromiseTask;
import edu.uw.zookeeper.common.TimeValue;
import edu.uw.zookeeper.data.ZNodeLabel;
import edu.uw.zookeeper.protocol.ConnectMessage;
import edu.uw.zookeeper.protocol.Message;
import edu.uw.zookeeper.protocol.Operation;
import edu.uw.zookeeper.protocol.ProtocolCodec;
import edu.uw.zookeeper.protocol.ProtocolCodecConnection;
import edu.uw.zookeeper.protocol.ProtocolResponseMessage;
import edu.uw.zookeeper.protocol.ProtocolRequestMessage;
import edu.uw.zookeeper.protocol.client.AssignXidProcessor;
import edu.uw.zookeeper.protocol.client.ClientConnectionExecutor;
import edu.uw.zookeeper.protocol.client.ConnectTask;
import edu.uw.zookeeper.protocol.proto.IMultiRequest;
import edu.uw.zookeeper.protocol.proto.OpCodeXid;
import edu.uw.zookeeper.protocol.proto.Records;
import edu.uw.zookeeper.safari.Identifier;
import edu.uw.zookeeper.safari.peer.protocol.ShardedOperation;
import edu.uw.zookeeper.safari.peer.protocol.ShardedRequest;
import edu.uw.zookeeper.safari.peer.protocol.ShardedRequestMessage;
import edu.uw.zookeeper.safari.peer.protocol.ShardedResponseMessage;

public class ShardedClientConnectionExecutor<C extends ProtocolCodecConnection<? super Message.ClientSession, ? extends ProtocolCodec<?,?>, ?>> extends ClientConnectionExecutor<C> {

    public static <C extends ProtocolCodecConnection<? super Message.ClientSession, ? extends ProtocolCodec<?,?>, ?>> ShardedClientConnectionExecutor<C> newInstance(
            ShardedOperationTranslators translator,
            Function<ZNodeLabel.Path, Identifier> lookup,
            ConnectMessage.Request request,
            C connection,
            ScheduledExecutorService executor) {
        return newInstance(
                translator,
                lookup,
                request,
                AssignXidProcessor.newInstance(),
                connection,
                executor);
    }

    public static <C extends ProtocolCodecConnection<? super Message.ClientSession, ? extends ProtocolCodec<?,?>, ?>> ShardedClientConnectionExecutor<C> newInstance(
            ShardedOperationTranslators translator,
            Function<ZNodeLabel.Path, Identifier> lookup,
            ConnectMessage.Request request,
            AssignXidProcessor xids,
            C connection,
            ScheduledExecutorService executor) {
        return newInstance(
                translator,
                lookup,
                ConnectTask.create(connection, request),
                xids,
                connection,
                TimeValue.create(Long.valueOf(request.getTimeOut()), TimeUnit.MILLISECONDS),
                executor);
    }

    public static <C extends ProtocolCodecConnection<? super Message.ClientSession, ? extends ProtocolCodec<?,?>, ?>> ShardedClientConnectionExecutor<C> newInstance(
            ShardedOperationTranslators translator,
            Function<ZNodeLabel.Path, Identifier> lookup,
            ListenableFuture<ConnectMessage.Response> session,
            AssignXidProcessor xids,
            C connection,
            TimeValue timeOut,
            ScheduledExecutorService executor) {
        return new ShardedClientConnectionExecutor<C>(
                lookup,
                translator,
                session, 
                xids,
                connection,
                timeOut,
                executor);
    }

    protected final AtomicInteger lastXid = new AtomicInteger(-1);
    protected final Function<ZNodeLabel.Path, Identifier> lookup;
    protected final ShardedOperationTranslators translator;
    
    public ShardedClientConnectionExecutor(
            Function<ZNodeLabel.Path, Identifier> lookup,
            ShardedOperationTranslators translator,
            ListenableFuture<ConnectMessage.Response> session,
            AssignXidProcessor xids,
            C connection,
            TimeValue timeOut,
            ScheduledExecutorService executor) {
        super(session, xids, connection, timeOut, executor);
        this.lookup = lookup;
        this.translator = translator;
    }

    public ShardedOperation.Request<?> shard(Operation.Request request) {
        Records.Request record = 
                (Records.Request) ((request instanceof Records.Request) ?
                 request :
                ((Operation.RecordHolder<?>) request).record());
        Identifier shard = Identifier.zero();
        if (request instanceof ShardedOperation) {
            shard = ((ShardedOperation) request).getIdentifier();
        } else {
            if (record instanceof IMultiRequest) {
                throw new UnsupportedOperationException();
            } else if (record instanceof Records.PathGetter) {
                shard = lookup.apply(ZNodeLabel.Path.of(((Records.PathGetter) record).getPath()));
            } 
        }
        if (shard == null) {
            // TODO
            throw new UnsupportedOperationException();
        }
        record = translator.get(shard).apply(record);
        ShardedOperation.Request<?> sharded;
        if (request instanceof Records.Request) {
            sharded = ShardedRequest.of(
                    shard, record);
        } else {
            int xid = ((Operation.ProtocolRequest<?>) request).xid();
            if (lastXid.get() > 0 && lastXid.get() != xid - 1) {
                throw new AssertionError(String.format("%d %d", lastXid.get(), xid));
            }
            lastXid.set(xid);
            sharded = ShardedRequestMessage.of(
                    shard,
                    ProtocolRequestMessage.of(xid, record));
        }
        return sharded;
    }

    @Override
    @Subscribe
    public void handleResponse(Message.ServerResponse<?> message) {
        if ((state() != State.TERMINATED) && !(message instanceof ShardedOperation)) {
            timeOut.send(message);
            int xid = message.xid();
            if (xid != OpCodeXid.PING.xid()) {
                ShardedResponseMessage<?> unshardedResponse;
                PendingResponseTask next = pending.peek();
                if ((next != null) && (next.getXid() == xid)) {
                    pending.remove(next);
                    Identifier id = ((ShardedRequestTask) next.delegate()).getIdentifier();
                    Records.Response record = message.record();
                    Records.Response translated = translator.get(id).apply(record);
                    if (translated == record) {
                        unshardedResponse = ShardedResponseMessage.of(id, message);
                    } else {
                        unshardedResponse = ShardedResponseMessage.of(
                                id,
                                ProtocolResponseMessage.of(
                                        xid, message.zxid(), translated));
                    }
                    
                    next.set(unshardedResponse);
                } else {
                    if (xid != OpCodeXid.NOTIFICATION.xid()) {
                        // FIXME is this an error?
                        logger.warn("{} != {} ({})", next, message, this);
                    }
                    
                    Identifier id = Identifier.zero();
                    Records.Response record = message.record();
                    if (record instanceof Records.PathGetter) {
                        id = lookup.apply(ZNodeLabel.Path.of(((Records.PathGetter) record).getPath()));
                    }
                    if (! Identifier.zero().equals(id)) {
                        Records.Response translated = translator.get(id).apply(record);
                        if (translated == record) {
                            unshardedResponse = ShardedResponseMessage.of(id, message);
                        } else {
                            unshardedResponse = ShardedResponseMessage.of(
                                    id,
                                    ProtocolResponseMessage.of(
                                            xid, message.zxid(), translated));
                        }
                    } else {
                        unshardedResponse = ShardedResponseMessage.of(id, message);
                    }
                }
    
                post(unshardedResponse);
            }
        }
    }
    
    @Override
    protected boolean apply(PromiseTask<Operation.Request, Message.ServerResponse<?>> input) {
        return super.apply(new ShardedRequestTask(shard(input.task()), input));
    }
    
    protected static class ShardedRequestTask extends PromiseTask<Operation.Request, Message.ServerResponse<?>> implements ShardedOperation {

        protected final ShardedOperation.Request<?> sharded;
        
        public ShardedRequestTask(
                ShardedOperation.Request<?> sharded,
                Promise<Message.ServerResponse<?>> delegate) {
            super(sharded.getRequest(), delegate);
            this.sharded = sharded;
        }

        @Override
        public Identifier getIdentifier() {
            return sharded.getIdentifier();
        }
    }
}

@@ -130,58 +130,8 @@
 nit;
-%0Aimport java.util.concurrent.atomic.AtomicInteger;
 %0A%0Aim
@@ -3701,75 +3701,8 @@
  %7D%0A%0A
-    protected final AtomicInteger lastXid = new AtomicInteger(-1);%0A
     
@@ -5506,203 +5506,8 @@
 ();%0A
-            if (lastXid.get() %3E 0 && lastXid.get() != xid - 1) %7B%0A                throw new AssertionError(String.format(%22%25d %25d%22, lastXid.get(), xid));%0A            %7D%0A            lastXid.set(xid);%0A
     

remove debugging code