/*
 * Copyright (C) 2012 Facebook, 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 io.airlift.drift.server;

import com.google.common.collect.ImmutableMap;
import com.google.common.util.concurrent.Futures;
import com.google.common.util.concurrent.ListenableFuture;
import io.airlift.drift.TApplicationException;
import io.airlift.drift.TApplicationException.Type;
import io.airlift.drift.codec.ThriftCodecManager;
import io.airlift.drift.codec.metadata.ThriftMethodMetadata;
import io.airlift.drift.codec.metadata.ThriftServiceMetadata;
import io.airlift.drift.server.stats.MethodInvocationStat;
import io.airlift.drift.server.stats.MethodInvocationStatsFactory;
import io.airlift.drift.transport.MethodMetadata;
import io.airlift.drift.transport.server.ServerInvokeRequest;
import io.airlift.drift.transport.server.ServerMethodInvoker;

import java.util.Collection;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Optional;

import static java.lang.String.format;

class DriftServerMethodInvoker
        implements ServerMethodInvoker
{
    private final Map<String, ServiceMethod> methods;
    private final Map<String, MethodInvocationStat> stats;

    public DriftServerMethodInvoker(
            ThriftCodecManager codecManager,
            Collection<DriftService> services,
            List<MethodInvocationFilter> filters,
            MethodInvocationStatsFactory methodInvocationStatsFactory)
    {
        Map<String, ServiceMethod> processorMap = new HashMap<>();
        ImmutableMap.Builder<String, MethodInvocationStat> stats = ImmutableMap.builder();
        for (DriftService service : services) {
            ThriftServiceMetadata serviceMetadata = new ThriftServiceMetadata(service.getService().getClass(), codecManager.getCatalog());
            for (ThriftMethodMetadata thriftMethodMetadata : serviceMetadata.getMethods().values()) {
                if (processorMap.containsKey(thriftMethodMetadata.getName())) {
                    throw new IllegalArgumentException(format("Multiple methods named '%s' are annotated with @ThriftMethod in the given services", thriftMethodMetadata.getName()));
                }
                ServiceMethod serviceMethod = new ServiceMethod(codecManager, service.getService(), thriftMethodMetadata, filters);
                processorMap.put(thriftMethodMetadata.getName(), serviceMethod);
                if (service.isStatsEnabled()) {
                    stats.put(thriftMethodMetadata.getName(), methodInvocationStatsFactory.getStat(serviceMetadata, service.getQualifier(), serviceMethod.getMethodMetadata()));
                }
            }
        }
        methods = ImmutableMap.copyOf(processorMap);
        this.stats = stats.build();
    }

    @Override
    public Optional<MethodMetadata> getMethodMetadata(String name)
    {
        ServiceMethod method = methods.get(name);
        if (method == null) {
            return Optional.empty();
        }
        return Optional.of(method.getMethodMetadata());
    }

    @Override
    public ListenableFuture<Object> invoke(ServerInvokeRequest request)
    {
        ServiceMethod method = methods.get(request.getMethod().getName());
        if (method == null) {
            return Futures.immediateFailedFuture(new TApplicationException(Type.UNKNOWN_METHOD, "Invalid method name: '" + request.getMethod().getName() + "'"));
        }

        return method.invokeMethod(request);
    }

    @Override
    public void recordResult(String methodName, long startTime, ListenableFuture<Object> result)
    {
        MethodInvocationStat stat = stats.get(methodName);
        if (stat == null) {
            return;
        }
        stat.recordResult(startTime, result);
    }
}