/*
 * Copyright (c) 2016-2018 Flux Capacitor.
 *
 * 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 io.fluxcapacitor.javaclient.publishing;

import io.fluxcapacitor.common.Registration;
import io.fluxcapacitor.common.api.SerializedMessage;
import io.fluxcapacitor.javaclient.common.Message;
import io.fluxcapacitor.javaclient.common.serialization.MessageSerializer;
import io.fluxcapacitor.javaclient.publishing.client.GatewayClient;
import io.fluxcapacitor.javaclient.tracking.handling.HandlerRegistry;
import lombok.AllArgsConstructor;

import java.util.Optional;
import java.util.concurrent.CompletableFuture;

import static java.lang.String.format;

@AllArgsConstructor
public class DefaultGenericGateway implements RequestGateway {
    private final GatewayClient gatewayClient;
    private final RequestHandler requestHandler;
    private final MessageSerializer serializer;
    private final HandlerRegistry localHandlerRegistry;

    @Override
    public void sendAndForget(Message message) {
        SerializedMessage serializedMessage = serializer.serialize(message);
        Optional<CompletableFuture<Message>> localResult 
                = localHandlerRegistry.handle(message.getPayload(), serializedMessage);
        if (!localResult.isPresent()) {
            try {
                gatewayClient.send(serializedMessage);
            } catch (Exception e) {
                throw new GatewayException(format("Failed to send and forget %s", message.getPayload().toString()), e);
            }
        }
    }

    @Override
    public CompletableFuture<Message> sendForMessage(Message message) {
        SerializedMessage serializedMessage = serializer.serialize(message);
        Optional<CompletableFuture<Message>> localResult 
                = localHandlerRegistry.handle(message.getPayload(), serializedMessage);
        if (localResult.isPresent()) {
            return localResult.get();
        } else {
            try {
                return requestHandler.sendRequest(serializedMessage, gatewayClient::send);
            } catch (Exception e) {
                throw new GatewayException(format("Failed to send %s", message.getPayload().toString()), e);
            }
        }
    }

    @Override
    public Registration registerLocalHandler(Object target) {
        return localHandlerRegistry.registerHandler(target);
    }

}

@@ -1041,16 +1041,44 @@
 tructor;
+%0Aimport lombok.SneakyThrows;
 %0A%0Aimport
@@ -1144,16 +1144,65 @@
 eFuture;
+%0Aimport java.util.concurrent.CompletionException;
 %0A%0Aimport
@@ -1525,32 +1525,50 @@
 %0A%0A    @Override%0A
+    @SneakyThrows%0A
     public void 
@@ -1964,32 +1964,32 @@
 (Exception e) %7B%0A
-
                 
@@ -2098,32 +2098,261 @@
 ;%0A            %7D%0A
+        %7D else if (localResult.get().isCompletedExceptionally()) %7B%0A            try %7B%0A                localResult.get().getNow(null);%0A            %7D catch (CompletionException e) %7B%0A                throw e.getCause();%0A            %7D%0A
         %7D%0A    %7D%0A

Throw local handler exception on send and forget