/*
 * Copyright (C) 2014 Indeed Inc.
 *
 * Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except
 * in compliance with the License. You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software distributed under the
 * License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
 * express or implied. See the License for the specific language governing permissions and
 * limitations under the License.
 */
 package com.indeed.imhotep;

import com.google.common.base.Function;
import com.google.common.base.Throwables;
import com.google.common.collect.Iterables;
import com.google.common.collect.Lists;
import com.google.common.io.ByteStreams;
import com.google.common.primitives.Doubles;
import com.google.common.primitives.Ints;
import com.google.common.primitives.Longs;
import com.indeed.imhotep.api.DocIterator;
import com.indeed.imhotep.api.FTGSIterator;
import com.indeed.imhotep.api.ImhotepOutOfMemoryException;
import com.indeed.imhotep.api.RawFTGSIterator;
import com.indeed.imhotep.io.ImhotepProtobufShipping;
import com.indeed.imhotep.io.LimitedBufferedOutputStream;
import com.indeed.imhotep.io.Streams;
import com.indeed.imhotep.io.TempFileSizeLimitExceededException;
import com.indeed.imhotep.io.WriteLimitExceededException;
import com.indeed.imhotep.marshal.ImhotepClientMarshaller;
import com.indeed.imhotep.protobuf.DatasetInfoMessage;
import com.indeed.imhotep.protobuf.GroupMultiRemapMessage;
import com.indeed.imhotep.protobuf.GroupRemapMessage;
import com.indeed.imhotep.protobuf.HostAndPort;
import com.indeed.imhotep.protobuf.ImhotepRequest;
import com.indeed.imhotep.protobuf.ImhotepResponse;
import com.indeed.imhotep.protobuf.IntFieldAndTerms;
import com.indeed.imhotep.protobuf.QueryRemapMessage;
import com.indeed.imhotep.protobuf.RegroupConditionMessage;
import com.indeed.imhotep.protobuf.ShardInfoMessage;
import com.indeed.imhotep.protobuf.StringFieldAndTerms;
import com.indeed.imhotep.service.InputStreamDocIterator;

import com.indeed.util.core.Throwables2;
import org.apache.log4j.Logger;

import javax.annotation.Nullable;
import java.io.BufferedInputStream;
import java.io.File;
import java.io.FileInputStream;
import java.io.FileOutputStream;
import java.io.FilterInputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.net.InetSocketAddress;
import java.net.Socket;
import java.net.SocketTimeoutException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicLong;

/**
 * @author jsgroth
 *
 * an ImhotepSession for talking to a remote ImhotepDaemon over a Socket using protobufs
 */
public class ImhotepRemoteSession extends AbstractImhotepSession {
    private static final Logger log = Logger.getLogger(ImhotepRemoteSession.class);

    public static final int DEFAULT_MERGE_THREAD_LIMIT = ImhotepRequest.getDefaultInstance().getMergeThreadLimit();

    private static final int DEFAULT_SOCKET_TIMEOUT = (int)TimeUnit.MINUTES.toMillis(30);

    private static final int CURRENT_CLIENT_VERSION = 2; // id to be incremented as changes to the client are done

    private final String host;
    private final int port;
    private final String sessionId;
    private final int socketTimeout;
    private final AtomicLong tempFileSizeBytesLeft;

    private int numStats = 0;

    public ImhotepRemoteSession(String host, int port, String sessionId, AtomicLong tempFileSizeBytesLeft) {
        this(host, port, sessionId, tempFileSizeBytesLeft, DEFAULT_SOCKET_TIMEOUT);
    }
    
    public ImhotepRemoteSession(String host, int port, String sessionId, @Nullable AtomicLong tempFileSizeBytesLeft, int socketTimeout) {
        this.host = host;
        this.port = port;
        this.sessionId = sessionId;
        this.socketTimeout = socketTimeout;
        this.tempFileSizeBytesLeft = tempFileSizeBytesLeft;
    }

    @Deprecated
    public static List<ShardInfo> getShardList(final String host, final int port) throws IOException {
        log.trace("sending get shard request to "+host+":"+port);
        final ImhotepRequest request = getBuilderForType(ImhotepRequest.RequestType.GET_SHARD_LIST)
                .build();

        final ImhotepResponse response = sendRequest(request, host, port);

        final List<ShardInfoMessage> protoShardInfo = response.getShardInfoList();
        final List<ShardInfo> ret = new ArrayList<ShardInfo>(protoShardInfo.size());
        for (final ShardInfoMessage shardInfo : protoShardInfo) {
            ret.add(ShardInfo.fromProto(shardInfo));
        }
        return ret;
    }

    public static List<DatasetInfo> getShardInfoList(final String host, final int port) throws IOException {
        final ImhotepRequest request = getBuilderForType(ImhotepRequest.RequestType.GET_SHARD_INFO_LIST)
                .build();

        final ImhotepResponse response = sendRequest(request, host, port);

        final List<DatasetInfoMessage> protoShardInfo = response.getDatasetInfoList();
        final List<DatasetInfo> ret = Lists.newArrayListWithCapacity(protoShardInfo.size());
        for (final DatasetInfoMessage datasetInfo : protoShardInfo) {
            ret.add(DatasetInfo.fromProto(datasetInfo));
        }
        return ret;
    }

    public static ImhotepStatusDump getStatusDump(final String host, final int port) throws IOException {
        final ImhotepRequest request = getBuilderForType(ImhotepRequest.RequestType.GET_STATUS_DUMP)
                .build();

        final ImhotepResponse response = sendRequest(request, host, port);

        return ImhotepStatusDump.fromProto(response.getStatusDump());
    }

    public static ImhotepRemoteSession openSession(final String host, final int port, final String dataset, final List<String> shards, @Nullable String sessionId) throws ImhotepOutOfMemoryException, IOException {
        return openSession(host, port, dataset, shards, DEFAULT_MERGE_THREAD_LIMIT, getUsername(), false, -1, sessionId, -1, null);
    }

    public static ImhotepRemoteSession openSession(final String host, final int port, final String dataset, final List<String> shards,
                                                   final int mergeThreadLimit, @Nullable String sessionId) throws ImhotepOutOfMemoryException, IOException {
        return openSession(host, port, dataset, shards, mergeThreadLimit, getUsername(), false, -1, sessionId, -1, null);
    }

    public static ImhotepRemoteSession openSession(final String host, final int port, final String dataset, final List<String> shards,
                                                   final int mergeThreadLimit, final String username,
                                                   final boolean optimizeGroupZeroLookups, final int socketTimeout, @Nullable String sessionId, final long tempFileSizeLimit, @Nullable final AtomicLong tempFileSizeBytesLeft) throws ImhotepOutOfMemoryException, IOException {
        final Socket socket = newSocket(host, port, socketTimeout);
        final OutputStream os = Streams.newBufferedOutputStream(socket.getOutputStream());
        final InputStream is = Streams.newBufferedInputStream(socket.getInputStream());

        try {
            log.trace("sending open request to "+host+":"+port+" for shards "+shards);
            final ImhotepRequest openSessionRequest = getBuilderForType(ImhotepRequest.RequestType.OPEN_SESSION)
                    .setUsername(username)
                    .setDataset(dataset)
                    .setMergeThreadLimit(mergeThreadLimit)
                    .addAllShardRequest(shards)
                    .setOptimizeGroupZeroLookups(optimizeGroupZeroLookups)
                    .setClientVersion(CURRENT_CLIENT_VERSION)
                    .setSessionId(sessionId == null ? "" : sessionId)
                    .setTempFileSizeLimit(tempFileSizeLimit)
                    .build();
            try {
                ImhotepProtobufShipping.sendProtobuf(openSessionRequest, os);
    
                log.trace("waiting for confirmation from "+host+":"+port);
                final ImhotepResponse response = ImhotepProtobufShipping.readResponse(is);
                if (response.getResponseCode() == ImhotepResponse.ResponseCode.OUT_OF_MEMORY) {
                    throw new ImhotepOutOfMemoryException();
                } else if (response.getResponseCode() == ImhotepResponse.ResponseCode.OTHER_ERROR) {
                    throw buildExceptionFromResponse(response, host, port);
                }
                if (sessionId == null) sessionId = response.getSessionId();
    
                log.trace("session created, id "+sessionId);
                return new ImhotepRemoteSession(host, port, sessionId, tempFileSizeBytesLeft, socketTimeout);
            } catch (SocketTimeoutException e) {
                throw buildExceptionAfterSocketTimeout(e, host, port);
            }
        } finally {
            closeSocket(socket, is, os);
        }
    }

    public static String getUsername() {
        return System.getProperty("user.name");
    }

    @Override
    public long getTotalDocFreq(String[] intFields, String[] stringFields) {
        final ImhotepRequest request = getBuilderForType(ImhotepRequest.RequestType.GET_TOTAL_DOC_FREQ)
                .setSessionId(sessionId)
                .addAllIntFields(Arrays.asList(intFields))
                .addAllStringFields(Arrays.asList(stringFields))
                .build();

        try {
            final ImhotepResponse response = sendRequest(request, host, port, socketTimeout);
            return response.getTotalDocFreq();
        } catch (IOException e) {
            throw new RuntimeException(e);
        }
    }

    @Override
    public long[] getGroupStats(int stat) {
        final ImhotepRequest request = getBuilderForType(ImhotepRequest.RequestType.GET_GROUP_STATS)
                .setSessionId(sessionId)
                .setStat(stat)
                .build();
        final ImhotepResponse response;
        try {
            response = sendRequest(request, host, port, socketTimeout);
        } catch (IOException e) {
            throw new RuntimeException(e);
        }

        final List<Long> groupStats = response.getGroupStatList();
        final long[] ret = new long[groupStats.size()];
        for (int i = 0; i < ret.length; ++i) {
            ret[i] = groupStats.get(i);
        }
        return ret;
    }

    @Override
    public FTGSIterator getFTGSIterator(String[] intFields, String[] stringFields) {
        final ImhotepRequest request = getBuilderForType(ImhotepRequest.RequestType.GET_FTGS_ITERATOR)
                .setSessionId(sessionId)
                .addAllIntFields(Arrays.asList(intFields))
                .addAllStringFields(Arrays.asList(stringFields))
                .build();

        return fileBufferedFTGSRequest(request);
    }

    @Override
    public FTGSIterator getSubsetFTGSIterator(Map<String, long[]> intFields, Map<String, String[]> stringFields) {
        final ImhotepRequest.Builder requestBuilder = getBuilderForType(ImhotepRequest.RequestType.GET_SUBSET_FTGS_ITERATOR)
                .setSessionId(sessionId);
        addSubsetFieldsAndTermsToBuilder(intFields, stringFields, requestBuilder);
        return fileBufferedFTGSRequest(requestBuilder.build());
    }

    private void addSubsetFieldsAndTermsToBuilder(Map<String, long[]> intFields, Map<String, String[]> stringFields, ImhotepRequest.Builder requestBuilder) {
        for (Map.Entry<String, long[]> entry : intFields.entrySet()) {
            final IntFieldAndTerms.Builder builder = IntFieldAndTerms.newBuilder().setField(entry.getKey());
            for (long term : entry.getValue()) {
                builder.addTerms(term);
            }
            requestBuilder.addIntFieldsToTerms(builder);
        }
        for (Map.Entry<String, String[]> entry : stringFields.entrySet()) {
            requestBuilder.addStringFieldsToTerms(
                    StringFieldAndTerms.newBuilder()
                            .setField(entry.getKey())
                            .addAllTerms(Arrays.asList(entry.getValue()))
            );
        }
    }

    public DocIterator getDocIterator(final String[] intFields, final String[] stringFields) throws ImhotepOutOfMemoryException {
        final ImhotepRequest request = getBuilderForType(ImhotepRequest.RequestType.GET_DOC_ITERATOR)
                .setSessionId(sessionId)
                .addAllIntFields(Arrays.asList(intFields))
                .addAllStringFields(Arrays.asList(stringFields))
                .build();
        try {
            final Socket socket = newSocket(host, port, socketTimeout);
            final InputStream is = Streams.newBufferedInputStream(socket.getInputStream());
            final OutputStream os = Streams.newBufferedOutputStream(socket.getOutputStream());
            try {
                sendRequest(request, is, os, host, port);
            } catch (IOException e) {
                closeSocket(socket, is, os);
                throw e;
            }
            return new InputStreamDocIterator(is, intFields.length, stringFields.length);
        } catch (IOException e) {
            throw Throwables.propagate(e);
        }
    }

    public RawFTGSIterator[] getFTGSIteratorSplits(final String[] intFields, final String[] stringFields) {
        throw new UnsupportedOperationException();
    }

    public RawFTGSIterator getFTGSIteratorSplit(final String[] intFields, final String[] stringFields, final int splitIndex, final int numSplits) {
        final ImhotepRequest request = getBuilderForType(ImhotepRequest.RequestType.GET_FTGS_SPLIT)
                .setSessionId(sessionId)
                .addAllIntFields(Arrays.asList(intFields))
                .addAllStringFields(Arrays.asList(stringFields))
                .setSplitIndex(splitIndex)
                .setNumSplits(numSplits)
                .build();

        return sendGetFTGSIteratorSplit(request);
    }

    @Override
    public RawFTGSIterator[] getSubsetFTGSIteratorSplits(Map<String, long[]> intFields, Map<String, String[]> stringFields) {
        throw new UnsupportedOperationException();
    }

    @Override
    public RawFTGSIterator getSubsetFTGSIteratorSplit(Map<String, long[]> intFields, Map<String, String[]> stringFields, int splitIndex, int numSplits) {
        final ImhotepRequest.Builder requestBuilder = getBuilderForType(ImhotepRequest.RequestType.GET_SUBSET_FTGS_SPLIT)
                .setSessionId(sessionId)
                .setSplitIndex(splitIndex)
                .setNumSplits(numSplits);
        addSubsetFieldsAndTermsToBuilder(intFields, stringFields, requestBuilder);
        return sendGetFTGSIteratorSplit(requestBuilder.build());
    }

    private RawFTGSIterator sendGetFTGSIteratorSplit(ImhotepRequest request) {
        try {
            final Socket socket = newSocket(host, port, socketTimeout);
            final InputStream is = Streams.newBufferedInputStream(socket.getInputStream());
            final OutputStream os = Streams.newBufferedOutputStream(socket.getOutputStream());
            try {
                sendRequest(request, is, os, host, port);
            } catch (IOException e) {
                closeSocket(socket, is, os);
                throw e;
            }
            return new ClosingInputStreamFTGSIterator(socket, is, os, numStats);
        } catch (IOException e) {
            throw new RuntimeException(e); // TODO
        }
    }

    public RawFTGSIterator mergeFTGSSplit(final String[] intFields, final String[] stringFields, final String sessionId, final InetSocketAddress[] nodes, final int splitIndex) {
        final ImhotepRequest request = getBuilderForType(ImhotepRequest.RequestType.MERGE_FTGS_SPLIT)
                .setSessionId(sessionId)
                .addAllIntFields(Arrays.asList(intFields))
                .addAllStringFields(Arrays.asList(stringFields))
                .setSplitIndex(splitIndex)
                .addAllNodes(Iterables.transform(Arrays.asList(nodes), new Function<InetSocketAddress, HostAndPort>() {
                    public HostAndPort apply(final InetSocketAddress input) {
                        return HostAndPort.newBuilder().setHost(input.getHostName()).setPort(input.getPort()).build();
                    }
                }))
                .build();

        return fileBufferedFTGSRequest(request);
    }

    @Override
    public RawFTGSIterator mergeSubsetFTGSSplit(Map<String, long[]> intFields, Map<String, String[]> stringFields, String sessionId, InetSocketAddress[] nodes, int splitIndex) {
        final ImhotepRequest.Builder requestBuilder = getBuilderForType(ImhotepRequest.RequestType.MERGE_SUBSET_FTGS_SPLIT)
                .setSessionId(sessionId)
                .setSplitIndex(splitIndex)
                .addAllNodes(Iterables.transform(Arrays.asList(nodes), new Function<InetSocketAddress, HostAndPort>() {
                    public HostAndPort apply(final InetSocketAddress input) {
                        return HostAndPort.newBuilder().setHost(input.getHostName()).setPort(input.getPort()).build();
                    }
                }));
        addSubsetFieldsAndTermsToBuilder(intFields, stringFields, requestBuilder);
        return fileBufferedFTGSRequest(requestBuilder.build());
    }

    private RawFTGSIterator fileBufferedFTGSRequest(ImhotepRequest request) {
        try {
            final Socket socket = newSocket(host, port, socketTimeout);
            final InputStream is = Streams.newBufferedInputStream(socket.getInputStream());
            final OutputStream os = Streams.newBufferedOutputStream(socket.getOutputStream());
            try {
                sendRequest(request, is, os, host, port);
            } catch (IOException e) {
                closeSocket(socket, is, os);
                throw e;
            }
            File tmp = null;
            try {
                tmp = File.createTempFile("ftgs", ".tmp");
                OutputStream out = null;
                try {
                    final long start = System.currentTimeMillis();
                    out = new LimitedBufferedOutputStream(new FileOutputStream(tmp), tempFileSizeBytesLeft);
                    ByteStreams.copy(is, out);
                    if(log.isDebugEnabled()) {
                        log.debug("time to copy split data to file: " + (System.currentTimeMillis() - start) + " ms, file length: " + tmp.length());
                    }
                } catch (Throwable t) {
                    tmp.delete();
                    if(t instanceof WriteLimitExceededException) {
                        throw new TempFileSizeLimitExceededException(t);
                    }
                    throw Throwables2.propagate(t, IOException.class);
                } finally {
                    if (out != null) {
                        out.close();
                    }
                }
                final BufferedInputStream bufferedInputStream = new BufferedInputStream(new FileInputStream(tmp));
                final InputStream in = new FilterInputStream(bufferedInputStream) {
                    public void close() throws IOException {
                        bufferedInputStream.close();
                    }
                };
                return new InputStreamFTGSIterator(in, numStats);
            } finally {
                if (tmp != null) {
                    tmp.delete();
                }
                closeSocket(socket, is, os);
            }
        } catch (IOException e) {
            throw new RuntimeException(e); // TODO
        }
    }

    @Override
    public int regroup(GroupMultiRemapRule[] rawRules, boolean errorOnCollisions) throws ImhotepOutOfMemoryException {
        try {
            final ImhotepResponse response = sendMultisplitRegroupRequest(rawRules, sessionId, errorOnCollisions);
            return response.getNumGroups();
        } catch (IOException e) {
            throw new RuntimeException(e);
        }
    }

    @Override
    public int regroup(int numRawRules, Iterator<GroupMultiRemapRule> rawRules, boolean errorOnCollisions) throws ImhotepOutOfMemoryException {
        try {
            final ImhotepResponse response = sendMultisplitRegroupRequest(numRawRules, rawRules, sessionId, errorOnCollisions);
            return response.getNumGroups();
        } catch (IOException e) {
            throw new RuntimeException(e);
        }
    }

    @Override
    public int regroup(GroupRemapRule[] rawRules) throws ImhotepOutOfMemoryException {
        final List<GroupRemapMessage> protoRules = ImhotepClientMarshaller.marshal(rawRules);

        final ImhotepRequest request = getBuilderForType(ImhotepRequest.RequestType.REGROUP)
                .setSessionId(sessionId)
                .addAllRemapRules(protoRules)
                .build();

        try {
            final ImhotepResponse response = sendRequestWithMemoryException(request, host, port, socketTimeout);
            return response.getNumGroups();
        } catch (IOException e) {
            throw new RuntimeException(e);
        }
    }

    @Override
    public int regroup(QueryRemapRule rule) throws ImhotepOutOfMemoryException {
        final QueryRemapMessage protoRule = ImhotepClientMarshaller.marshal(rule);

        final ImhotepRequest request = getBuilderForType(ImhotepRequest.RequestType.QUERY_REGROUP)
                .setSessionId(sessionId)
                .setQueryRemapRule(protoRule)
                .build();

        try {
            final ImhotepResponse response = sendRequestWithMemoryException(request, host, port, socketTimeout);
            return response.getNumGroups();
        } catch (IOException e) {
            throw new RuntimeException(e);
        }
    }

    @Override
    public void intOrRegroup(String field, long[] terms, int targetGroup, int negativeGroup, int positiveGroup) throws ImhotepOutOfMemoryException {
        final ImhotepRequest request = getBuilderForType(ImhotepRequest.RequestType.INT_OR_REGROUP)
                .setSessionId(sessionId)
                .setField(field)
                .addAllIntTerm(Longs.asList(terms))
                .setTargetGroup(targetGroup)
                .setNegativeGroup(negativeGroup)
                .setPositiveGroup(positiveGroup)
                .build();

        try {
            sendRequestWithMemoryException(request, host, port, socketTimeout);
        } catch (IOException e) {
            throw new RuntimeException(e);
        }
    }

    @Override
    public void stringOrRegroup(String field, String[] terms, int targetGroup, int negativeGroup, int positiveGroup) throws ImhotepOutOfMemoryException {
        final ImhotepRequest request = getBuilderForType(ImhotepRequest.RequestType.STRING_OR_REGROUP)
                .setSessionId(sessionId)
                .setField(field)
                .addAllStringTerm(Arrays.asList(terms))
                .setTargetGroup(targetGroup)
                .setNegativeGroup(negativeGroup)
                .setPositiveGroup(positiveGroup)
                .build();

        try {
            sendRequestWithMemoryException(request, host, port, socketTimeout);
        } catch (IOException e) {
            throw new RuntimeException(e);
        }
    }

    @Override
    public void regexRegroup(String field, String regex, int targetGroup, int negativeGroup, int positiveGroup) throws ImhotepOutOfMemoryException {
        final ImhotepRequest request = getBuilderForType(ImhotepRequest.RequestType.REGEX_REGROUP)
                .setSessionId(sessionId)
                .setField(field)
                .setRegex(regex)
                .setTargetGroup(targetGroup)
                .setNegativeGroup(negativeGroup)
                .setPositiveGroup(positiveGroup)
                .build();

        try {
            sendRequestWithMemoryException(request, host, port, socketTimeout);
        } catch (IOException e) {
            throw new RuntimeException(e);
        }
    }

    @Override
    public void randomRegroup(String field, boolean isIntField, String salt, double p, int targetGroup, int negativeGroup,
                              int positiveGroup) throws ImhotepOutOfMemoryException {
        final ImhotepRequest request = getBuilderForType(ImhotepRequest.RequestType.RANDOM_REGROUP)
                .setSessionId(sessionId)
                .setField(field)
                .setIsIntField(isIntField)
                .setSalt(salt)
                .setP(p)
                .setTargetGroup(targetGroup)
                .setNegativeGroup(negativeGroup)
                .setPositiveGroup(positiveGroup)
                .build();

        try {
            sendRequestWithMemoryException(request, host, port, socketTimeout);
        } catch (IOException e) {
            throw new RuntimeException(e);
        }
    }

    @Override
    public void randomMultiRegroup(String field, boolean isIntField, String salt, int targetGroup, double[] percentages,
                                   int[] resultGroups) throws ImhotepOutOfMemoryException {
        final ImhotepRequest request = getBuilderForType(ImhotepRequest.RequestType.RANDOM_MULTI_REGROUP)
                .setSessionId(sessionId)
                .setField(field)
                .setIsIntField(isIntField)
                .setSalt(salt)
                .setTargetGroup(targetGroup)
                .addAllPercentages(Doubles.asList(percentages))
                .addAllResultGroups(Ints.asList(resultGroups))
                .build();

        try {
            sendRequestWithMemoryException(request, host, port, socketTimeout);
        } catch (IOException e) {
            throw new RuntimeException(e);
        }
    }

    @Override
    public int metricRegroup(int stat, long min, long max, long intervalSize, boolean noGutters) throws ImhotepOutOfMemoryException {
        final ImhotepRequest request = getBuilderForType(ImhotepRequest.RequestType.METRIC_REGROUP)
                .setSessionId(sessionId)
                .setXStat(stat)
                .setXMin(min)
                .setXMax(max)
                .setXIntervalSize(intervalSize)
                .setNoGutters(noGutters)
                .build();

        try {
            final ImhotepResponse response = sendRequestWithMemoryException(request, host, port, socketTimeout);
            return response.getNumGroups();
        } catch (IOException e) {
            throw new RuntimeException(e);
        }
    }

    @Override
    public int metricRegroup2D(int xStat, long xMin, long xMax, long xIntervalSize, int yStat, long yMin, long yMax, long yIntervalSize) throws ImhotepOutOfMemoryException {
        final ImhotepRequest request = getBuilderForType(ImhotepRequest.RequestType.METRIC_REGROUP_2D)
                .setSessionId(sessionId)
                .setXStat(xStat)
                .setXMin(xMin)
                .setXMax(xMax)
                .setXIntervalSize(xIntervalSize)
                .setYStat(yStat)
                .setYMin(yMin)
                .setYMax(yMax)
                .setYIntervalSize(yIntervalSize)
                .build();

        try {
            final ImhotepResponse response = sendRequestWithMemoryException(request, host, port, socketTimeout);
            return response.getNumGroups();
        } catch (IOException e) {
            throw new RuntimeException(e);
        }
    }

    public int metricFilter(int stat, long min, long max, boolean negate) throws ImhotepOutOfMemoryException {
        final ImhotepRequest request = getBuilderForType(ImhotepRequest.RequestType.METRIC_FILTER)
                .setSessionId(sessionId)
                .setXStat(stat)
                .setXMin(min)
                .setXMax(max)
                .setNegate(negate)
                .build();
        try {
            final ImhotepResponse response = sendRequestWithMemoryException(request, host, port, socketTimeout);
            return response.getNumGroups();
        } catch (IOException e) {
            throw Throwables.propagate(e);
        }
    }

    @Override
    public List<TermCount> approximateTopTerms(String field, boolean isIntField, int k) {
        final ImhotepRequest request = getBuilderForType(ImhotepRequest.RequestType.APPROXIMATE_TOP_TERMS)
                .setSessionId(sessionId)
                .setField(field)
                .setIsIntField(isIntField)
                .setK(k)
                .build();

        try {
            final ImhotepResponse response = sendRequest(request, host, port, socketTimeout);
            return ImhotepClientMarshaller.marshal(response.getTopTermsList());
        } catch (IOException e) {
            throw new RuntimeException(e);
        }
    }

    @Override
    public int pushStat(String statName) throws ImhotepOutOfMemoryException {
        final ImhotepRequest request = getBuilderForType(ImhotepRequest.RequestType.PUSH_STAT)
                .setSessionId(sessionId)
                .setMetric(statName)
                .build();

        try {
            final ImhotepResponse response = sendRequestWithMemoryException(request, host, port, socketTimeout);
            numStats = response.getNumStats();
            return numStats;
        } catch (IOException e) {
            throw new RuntimeException(e);
        }
    }

    @Override
    public int pushStats(final List<String> statNames) throws ImhotepOutOfMemoryException {
        for (String statName : statNames) {
            this.pushStat(statName);
        }

        return numStats;
    }

    @Override
    public int popStat() {
        final ImhotepRequest request = getBuilderForType(ImhotepRequest.RequestType.POP_STAT)
                .setSessionId(sessionId)
                .build();

        try {
            final ImhotepResponse response = sendRequest(request, host, port, socketTimeout);
            numStats = response.getNumStats();
            return numStats;
        } catch (IOException e) {
            throw new RuntimeException(e);
        }
    }

    @Override
    public int getNumStats() {
        // TODO: really should ask the remote session just to be sure.
        return numStats;
    }

    @Override
    public int getNumGroups() {
        final ImhotepRequest request = getBuilderForType(ImhotepRequest.RequestType.GET_NUM_GROUPS)
                .setSessionId(sessionId)
                .build();
        try {
            final ImhotepResponse response = sendRequest(request, host, port, socketTimeout);
            return response.getNumGroups();
        } catch (IOException e) {
            throw Throwables.propagate(e);
        }
    }

    @Override
    public void createDynamicMetric(String name) throws ImhotepOutOfMemoryException {
        final ImhotepRequest request = getBuilderForType(ImhotepRequest.RequestType.CREATE_DYNAMIC_METRIC)
                .setSessionId(sessionId)
                .setDynamicMetricName(name)
                .build();

        try {
            sendRequestWithMemoryException(request, host, port, socketTimeout);
        } catch (IOException e) {
            throw new RuntimeException(e);
        }
    }

    @Override
    public void updateDynamicMetric(String name, int[] deltas) {
        final ImhotepRequest request = getBuilderForType(ImhotepRequest.RequestType.UPDATE_DYNAMIC_METRIC)
                .setSessionId(sessionId)
                .setDynamicMetricName(name)
                .addAllDynamicMetricDeltas(Ints.asList(deltas))
                .build();

        try {
            sendRequest(request, host, port);
        } catch (SocketTimeoutException e) {
            throw new RuntimeException(buildExceptionAfterSocketTimeout(e, host, port));
        } catch (IOException e) {
            throw new RuntimeException(e);
        }
    }

    @Override
    public void conditionalUpdateDynamicMetric(String name, RegroupCondition[] conditions, int[] deltas) {
        List<RegroupConditionMessage> conditionMessages = ImhotepClientMarshaller.marshal(conditions);

        final ImhotepRequest request = getBuilderForType(ImhotepRequest.RequestType.CONDITIONAL_UPDATE_DYNAMIC_METRIC)
                .setSessionId(sessionId)
                .setDynamicMetricName(name)
                .addAllConditions(conditionMessages)
                .addAllDynamicMetricDeltas(Ints.asList(deltas))
                .build();
        try {
            sendRequest(request, host, port, socketTimeout);
        } catch (IOException e) {
