package com.proofpoint.http.client.testing;

import com.google.common.base.Function;
import com.google.common.util.concurrent.ForwardingListenableFuture;
import com.google.common.util.concurrent.ListenableFuture;
import com.google.common.util.concurrent.ListeningExecutorService;
import com.google.common.util.concurrent.MoreExecutors;
import com.proofpoint.http.client.AsyncHttpClient;
import com.proofpoint.http.client.Request;
import com.proofpoint.http.client.RequestStats;
import com.proofpoint.http.client.Response;
import com.proofpoint.http.client.ResponseHandler;
import com.proofpoint.units.Duration;

import java.util.concurrent.Callable;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicReference;

import static com.google.common.base.Preconditions.checkNotNull;
import static com.google.common.base.Preconditions.checkState;

public class TestingHttpClient
        implements AsyncHttpClient
{
    private final Function<Request, Response> processor;
    private final ListeningExecutorService executor;

    private final RequestStats stats = new RequestStats();
    private final AtomicBoolean closed = new AtomicBoolean();

    public TestingHttpClient(Function<Request, Response> processor)
    {
        this(processor, MoreExecutors.sameThreadExecutor());
    }

    public TestingHttpClient(Function<Request, Response> processor, ExecutorService executor)
    {
        this.processor = processor;
        this.executor = MoreExecutors.listeningDecorator(executor);
    }

    @Override
    public <T, E extends Exception> AsyncHttpResponseFuture<T> executeAsync(final Request request, final ResponseHandler<T, E> responseHandler)
    {
        checkNotNull(request, "request is null");
        checkNotNull(responseHandler, "responseHandler is null");
        checkState(!closed.get(), "client is closed");

        final AtomicReference<String> state = new AtomicReference<>("SENDING_REQUEST");
        ListenableFuture<T> future = executor.submit(new Callable<T>()
        {
            @Override
            public T call()
                    throws Exception
            {
                return execute(request, responseHandler, state);
            }
        });

        return new TestingAsyncHttpResponseFuture<>(future, state);
    }

    @Override
    public <T, E extends Exception> T execute(Request request, ResponseHandler<T, E> responseHandler)
            throws E
    {
        checkNotNull(request, "request is null");
        checkNotNull(responseHandler, "responseHandler is null");
        checkState(!closed.get(), "client is closed");
        return execute(request, responseHandler, new AtomicReference<>("SENDING_REQUEST"));
    }

    private <T, E extends Exception> T execute(Request request, ResponseHandler<T, E> responseHandler, AtomicReference<String> state)
            throws E
    {
        state.set("PROCESSING_REQUEST");
        Response response;
        Duration requestProcessingTime = null;
        try {
            long requestStart = System.nanoTime();
            response = processor.apply(request);
            requestProcessingTime = Duration.nanosSince(requestStart);
        }
        catch (Throwable e) {
            state.set("FAILED");
            stats.record(request.getMethod(),
                    0,
                    0,
                    0,
                    requestProcessingTime,
                    null);
            if (e instanceof Exception) {
                return responseHandler.handleException(request, (Exception) e);
            }
            else {
                throw e;
            }
        }
        checkState(response != null, "response is null");

        // notify handler
        state.set("PROCESSING_RESPONSE");
        long responseStart = System.nanoTime();
        try {
            return responseHandler.handle(request, response);
        }
        finally {
            state.set("DONE");
            Duration responseProcessingTime = Duration.nanosSince(responseStart);
            stats.record(request.getMethod(),
                    response.getStatusCode(),
                    response.getBytesRead(),
                    response.getBytesRead(),
                    requestProcessingTime,
                    responseProcessingTime);
        }
    }

    @Override
    public RequestStats getStats()
    {
        return stats;
    }

    @Override
    public void close()
    {
        closed.set(true);
    }

    private class TestingAsyncHttpResponseFuture<T>
            extends ForwardingListenableFuture<T>
            implements AsyncHttpResponseFuture<T>
    {
        private final AtomicReference<String> state;
        private final ListenableFuture<T> future;

        private TestingAsyncHttpResponseFuture(ListenableFuture<T> future, AtomicReference<String> state)
        {
            this.future = future;
            this.state = state;
        }

        @Override
        public ListenableFuture<T> delegate()
        {
            return future;
        }

        @Override
        public String getState()
        {
            return state.get();
        }
    }
}

@@ -605,16 +605,61 @@
 ation;%0A%0A
+import javax.validation.constraints.NotNull;%0A
 import j
@@ -688,16 +688,16 @@
 llable;%0A
-
 import j
@@ -1055,35 +1055,17 @@
 nal 
-Function%3CRequest, Response%3E
+Processor
  pro
@@ -1423,35 +1423,466 @@
 ent(
-Function%3CRequest, Response%3E
+Processor processor)%0A    %7B%0A        this(processor, MoreExecutors.sameThreadExecutor());%0A    %7D%0A%0A    public TestingHttpClient(final Function%3CRequest, Response%3E processor, ExecutorService executor)%0A    %7B%0A        this(new Processor()%0A        %7B%0A            @Override%0A            public Response handle(Request request)%0A            %7B%0A                return processor.apply(request);%0A            %7D%0A        %7D, executor);%0A    %7D%0A%0A    public TestingHttpClient(Processor
  pro
@@ -3584,37 +3584,38 @@
 nse = processor.
-apply
+handle
 (request);%0A     
@@ -4114,16 +4114,24 @@
    throw
+ (Error)
  e;%0A    
@@ -4949,32 +4949,32 @@
 d close()%0A    %7B%0A
-
         closed.s
@@ -4982,32 +4982,168 @@
 t(true);%0A    %7D%0A%0A
+    public interface Processor%0A    %7B%0A        @NotNull%0A        Response handle(Request request)%0A                throws Exception;%0A    %7D%0A%0A
     private clas

Add ability for TestingHttpClient to send checked exceptions