/**
 * Copyright 2015 Palantir Technologies
 *
 * Licensed under the BSD-3 License (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://opensource.org/licenses/BSD-3-Clause
 *
 * 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.palantir.atlasdb.keyvalue.remoting.proxy;

import java.lang.reflect.InvocationHandler;
import java.lang.reflect.InvocationTargetException;
import java.lang.reflect.Method;
import java.lang.reflect.Proxy;
import java.util.Set;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import com.google.common.base.Function;
import com.google.common.base.Preconditions;
import com.google.common.base.Supplier;
import com.google.common.collect.ImmutableSet;
import com.palantir.atlasdb.keyvalue.api.KeyValueService;
import com.palantir.atlasdb.keyvalue.partition.api.DynamicPartitionMap;
import com.palantir.atlasdb.keyvalue.partition.exception.ClientVersionTooOldException;
import com.palantir.atlasdb.keyvalue.partition.exception.EndpointVersionTooOldException;
import com.palantir.atlasdb.keyvalue.remoting.RemotingKeyValueService.LONG_HOLDER;
import com.palantir.atlasdb.keyvalue.remoting.RemotingKeyValueService.STRING_HOLDER;
import com.palantir.atlasdb.keyvalue.remoting.RemotingPartitionMapService;
import com.palantir.common.supplier.RemoteContextHolder;
import com.palantir.common.supplier.ServiceContext;

/**
 * This is for the endpoint to ensure that client and server partition map versions are compatible.
 * It will throw <code>VersionTooOldException</code> if the client is out of date.
 *
 * The server version supplier is passed as an argument. Client version supplier is taken
 * from <code>RemoteContextHolder.INBOX.getProviderForKey(HOLDER.PM_VERSION)</code>.
 *
 * @see ClientVersionTooOldException
 * @see RemoteContextHolder
 * @see LONG_HOLDER
 * @see STRING_HOLDER
 *
 * @author htarasiuk
 *
 */
public class VersionCheckProxy implements InvocationHandler {
    private static final Logger log = LoggerFactory.getLogger(VersionCheckProxy.class);
    private final Supplier<Long> serverVersionProvider;
    private final Function<? super DynamicPartitionMap, Void> serverPartitionMapUpdater;
    private final KeyValueService delegate;

    private VersionCheckProxy(Supplier<Long> serverVersionProvider, Function<? super DynamicPartitionMap, Void> serverPartitionMapUpdater, KeyValueService delegate) {
        this.serverVersionProvider = serverVersionProvider;
        this.serverPartitionMapUpdater = serverPartitionMapUpdater;
        this.delegate = delegate;
    }


    /*
     * Table metadata methods are version check-exempt because the metadata
     * is stored on all enpoints anyway.
     *
     * This is necessary because the metadata skips transaction manager and
     * does not retry automatically on retryable exception.
     *
     * Some other "technical" methods are also exempt.
     */
    private static final boolean isMethodVersionExempt(Method method) {
        return method.getDeclaringClass() != KeyValueService.class || EXEMPT_METHODS.contains(method.getName());
    }

    private static final Set<String> EXEMPT_METHODS = ImmutableSet.<String>builder()
            // Table and table metadata methods
            .add("dropTable")
            .add("createTable")
            .add("createTables")
            .add("truncateTable")
            .add("truncateTables")
            .add("getAllTableNames")
            .add("getAllTableNames")
            .add("getMetadataForTable")
            .add("getMetadataForTables")
            .add("putMetadataForTable")
            .add("putMetadataForTables")
            // Technical methods
            .add("close")
            .add("teardown")
            .add("initializeFromFreshInstance")
            .add("compactInternally")
            .build();

    @Override
    public Object invoke(Object proxy, Method method, Object[] args) throws Throwable {

        ServiceContext<Long> remoteVersionClientCtx = RemoteContextHolder.INBOX.getProviderForKey(LONG_HOLDER.PM_VERSION);
        ServiceContext<String> newPmsUriCtx = RemoteContextHolder.INBOX.getProviderForKey(STRING_HOLDER.PMS_URI);

        // Only check the version for appropriate methods
        if (!isMethodVersionExempt(method)) {
            Long clientVersion = remoteVersionClientCtx.get();
            Long serverVersion = Preconditions.checkNotNull(serverVersionProvider.get());
            if (serverVersion < 0L) {
                // In this case the version check is simply disabled.
                assert clientVersion == null || clientVersion < 0;
            } else {
                if (!serverVersion.equals(clientVersion)) {
                    log.warn("Server map version: " + serverVersion + ", client map version: " + clientVersion);
                }
                if (clientVersion < serverVersion) {
                    throw new ClientVersionTooOldException();
                }
                if (clientVersion > serverVersion) {
                    String newPmsUri = newPmsUriCtx.get();
                    if (newPmsUri != null) {
                        try {
                            serverPartitionMapUpdater.apply(RemotingPartitionMapService.createClientSide(newPmsUri).getMap());
                        } catch (RuntimeException e) {
                            log.warn("Could not update server partition map");
                            throw new EndpointVersionTooOldException("Server partition map out-of-date and update failed.", e);
                        }
                    } else {
                        throw new EndpointVersionTooOldException("Server partition map out-of-date.");
                    }
                }
            }
        }

        try {
            return method.invoke(delegate, args);
        } catch (InvocationTargetException e) {
            throw e.getCause();
        }
    }

    /**
     *
     * @param delegate
     * @param serverVersionProvider Use <code>Suppliers.<Long>ofInstance(-1L)</code> to disable version check. In
     * such case this proxy is just a no-op.
     * @return
     */
    public static KeyValueService newProxyInstance(KeyValueService delegate, Supplier<Long> serverVersionProvider,
                                                   Function<? super DynamicPartitionMap, Void> serverPartitionMapUpdater) {
        VersionCheckProxy vcp = new VersionCheckProxy(serverVersionProvider, serverPartitionMapUpdater, delegate);
        return (KeyValueService) Proxy.newProxyInstance(
                KeyValueService.class.getClassLoader(),
                new Class<?>[] { KeyValueService.class }, vcp);
    }

}
@@ -801,30 +801,8 @@
 oxy;
-%0Aimport java.util.Set;
 %0A%0Aim
@@ -986,55 +986,8 @@
 er;%0A
-import com.google.common.collect.ImmutableSet;%0A
 impo
@@ -3303,790 +3303,15 @@
 lass
- %7C%7C EXEMPT_METHODS.contains(method.getName());%0A    %7D%0A%0A    private static final Set%3CString%3E EXEMPT_METHODS = ImmutableSet.%3CString%3Ebuilder()%0A            // Table and table metadata methods%0A            .add(%22dropTable%22)%0A            .add(%22createTable%22)%0A            .add(%22createTables%22)%0A            .add(%22truncateTable%22)%0A            .add(%22truncateTables%22)%0A            .add(%22getAllTableNames%22)%0A            .add(%22getAllTableNames%22)%0A            .add(%22getMetadataForTable%22)%0A            .add(%22getMetadataForTables%22)%0A            .add(%22putMetadataForTable%22)%0A            .add(%22putMetadataForTables%22)%0A            // Technical methods%0A            .add(%22close%22)%0A            .add(%22teardown%22)%0A            .add(%22initializeFromFreshInstance%22)%0A            .add(%22compactInternally%22)%0A            .build();
+;%0A    %7D
 %0A%0A  

Remove method exemption. All methods must check partition version!