/*
 * Copyright 2019 LINE Corporation
 *
 * LINE Corporation licenses this file to you 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:
 *
 *   https://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.linecorp.armeria.server.thrift;

import static com.google.common.base.Preconditions.checkArgument;
import static java.util.Objects.requireNonNull;

import java.util.LinkedHashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.function.BiFunction;
import java.util.function.Function;

import org.apache.thrift.protocol.TBinaryProtocol;
import org.apache.thrift.protocol.TCompactProtocol;

import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableListMultimap;
import com.google.common.collect.ImmutableSet;
import com.google.common.collect.Multimaps;

import com.linecorp.armeria.common.RpcResponse;
import com.linecorp.armeria.common.SerializationFormat;
import com.linecorp.armeria.common.annotation.Nullable;
import com.linecorp.armeria.common.thrift.ThriftProtocolFactoryProvider;
import com.linecorp.armeria.common.thrift.ThriftSerializationFormats;
import com.linecorp.armeria.server.RpcService;
import com.linecorp.armeria.server.ServerBuilder;
import com.linecorp.armeria.server.ServiceRequestContext;

/**
 * A fluent builder to build an instance of {@link THttpService}. This builder allows to bind multiple thrift
 * service implementations along with mixing TMultiplexed protocol to a route.
 * <h2>Example</h2>
 * <pre>{@code
 * Server server =
 *     Server.builder()
 *           .http(8080)
 *           .service("/", THttpService.builder()
 *                                     .addService(new FooService())              // foo() method
 *                                     .addService(new BarService())              // bar() method
 *                                     .addService("foobar", new FooBarService()) // TMultiplexed service
 *                                     .build())
 *           .build();
 * }</pre>
 *
 * <p>When the thrift request has a method {@code foo()} then {@code FooService.foo()} handles the request and
 * similarly when thrift request has a method {@code bar()} then {@code BarService.bar()} handles the request.
 * And when the service name is "foobar" then FooBarService</p>
 *
 * @see THttpService
 * @see ThriftCallService
 */
public final class THttpServiceBuilder {

    private static final BiFunction<? super ServiceRequestContext, ? super Throwable, ? extends RpcResponse>
            defaultExceptionHandler = (ctx, cause) -> RpcResponse.ofFailure(cause);

    private final ImmutableListMultimap.Builder<String, Object> implementationsBuilder =
            ImmutableListMultimap.builder();
    private SerializationFormat defaultSerializationFormat = ThriftSerializationFormats.BINARY;
    private Set<SerializationFormat> otherSerializationFormats = ThriftSerializationFormats.values();
    private boolean createOtherSerializations = true;
    @Nullable
    private Function<? super RpcService, ? extends RpcService> decoratorFunction;
    private BiFunction<? super ServiceRequestContext, ? super Throwable, ? extends RpcResponse>
            exceptionHandler = defaultExceptionHandler;

    // -1 means to use the default request length of the Server.
    private int maxRequestStringLength = -1;
    private int maxRequestContainerLength = -1;

    THttpServiceBuilder() {}

    /**
     * Adds a new {@code TMultiplexed} service to the builder.
     *
     * @param name name of the service.
     * @param implementation an implementation of {@code *.Iface} or {@code *.AsyncIface} service interface
     *                       generated by the Apache Thrift compiler.
     */
    public THttpServiceBuilder addService(String name, Object implementation) {
        implementationsBuilder.put(name, implementation);
        return this;
    }

    /**
     * Adds a new service implementation to the builder.
     *
     * @param implementation an implementation of {@code *.Iface} or {@code *.AsyncIface} service interface
     *                       generated by the Apache Thrift compiler
     */
    public THttpServiceBuilder addService(Object implementation) {
        return addService("", implementation);
    }

    /**
     * Adds other {@link SerializationFormat} to the builder. By default, all {@link SerializationFormat}s in
     * {@link ThriftSerializationFormats} are supported. If nothing is specified then they are added.
     * To add a new custom Thrift serialization format,
     * define a new SPI {@link ThriftProtocolFactoryProvider}.
     *
     * <p>Currently, the only way to specify a serialization format at request time is by using the HTTP session
     * protocol and setting the {@code "Content-Type"} header to the appropriate
     * {@link SerializationFormat#mediaType()}.
     */
    public THttpServiceBuilder otherSerializationFormats(SerializationFormat otherSerializationFormat) {
        requireNonNull(otherSerializationFormat, "otherSerializationFormat");
        return otherSerializationFormats(ImmutableList.of(otherSerializationFormat));
    }

    /**
     * Adds other {@link SerializationFormat}s to the builder. If nothing is specified then all
     * {@link SerializationFormat}s returned by {@link ThriftSerializationFormats#values()}
     * are added.
     *
     * <p>Currently, the only way to specify a serialization format at request time is by using the HTTP session
     * protocol and setting the {@code "Content-Type"} header to the appropriate
     * {@link SerializationFormat#mediaType()}.
     */
    public THttpServiceBuilder otherSerializationFormats(
            Iterable<SerializationFormat> otherSerializationFormats) {
        requireNonNull(otherSerializationFormats, "otherSerializationFormats");
        if (createOtherSerializations) {
            this.otherSerializationFormats = new LinkedHashSet<>();
            createOtherSerializations = false;
        }
        otherSerializationFormats.forEach(this.otherSerializationFormats::add);
        return this;
    }

    /**
     * Adds the default serialization format which will be used when the client does not specify one in
     * request.
     *
     * <p>Currently, the only way to specify a serialization format is by using the HTTP session
     * protocol and setting the {@code "Content-Type"} header to the appropriate
     * {@link SerializationFormat#mediaType()}.
     */
    public THttpServiceBuilder defaultSerializationFormat(SerializationFormat defaultSerializationFormat) {
        requireNonNull(defaultSerializationFormat, "defaultSerializationFormat");
        this.defaultSerializationFormat = defaultSerializationFormat;
        return this;
    }

    /**
     * Sets the maximum allowed number of bytes to read from the transport for
     * variable-length fields (such as strings or binary).
     * If unspecified, the default value of {@link ServerBuilder#maxRequestLength(long)} will be used.
     *
     * <p>Note that this options is only valid for {@link TBinaryProtocol} and {@link TCompactProtocol}.
     *
     * @param maxRequestStringLength the maximum allowed string length. {@code 0} disables the length limit.
     * @see ServerBuilder#maxRequestLength(long)
     */
    public THttpServiceBuilder maxRequestStringLength(int maxRequestStringLength) {
        checkArgument(maxRequestStringLength >= 0, "maxRequestStringLength: %s (expected: >= 0)",
                      maxRequestStringLength);
        this.maxRequestStringLength = maxRequestStringLength;
        return this;
    }

    /**
     * Sets the maximum allowed number of containers to read from the transport for maps, sets and lists.
     * If unspecified, the default value of {@link ServerBuilder#maxRequestLength(long)} will be used.
     *
     * <p>Note that this options is only valid for {@link TBinaryProtocol} and {@link TCompactProtocol}.
     *
     * @param maxRequestContainerLength the maximum allowed string length. {@code 0} disables the length limit.
     */
    public THttpServiceBuilder maxRequestContainerLength(int maxRequestContainerLength) {
        checkArgument(maxRequestContainerLength >= 0, "maxRequestContainerLength: %s (expected: >= 0)",
                      maxRequestContainerLength);
        this.maxRequestContainerLength = maxRequestContainerLength;
        return this;
    }

    /**
     * Sets the {@link BiFunction} that returns an {@link RpcResponse} using the given {@link Throwable}
     * and {@link ServiceRequestContext}.
     */
    public THttpServiceBuilder exceptionHandler(
            BiFunction<? super ServiceRequestContext, ? super Throwable, ? extends RpcResponse>
                    exceptionHandler) {
        this.exceptionHandler = requireNonNull(exceptionHandler, "exceptionHandler");
        return this;
    }

    /**
     * A {@code Function<? super RpcService, ? extends RpcService>} to decorate the {@link RpcService}.
     */
    public THttpServiceBuilder decorate(Function<? super RpcService, ? extends RpcService> decoratorFunction) {
        requireNonNull(decoratorFunction, "decoratorFunction");
        if (this.decoratorFunction == null) {
            this.decoratorFunction = decoratorFunction;
        } else {
            this.decoratorFunction = this.decoratorFunction.andThen(decoratorFunction);
        }
        return this;
    }

    private RpcService decorate(RpcService service) {
        if (decoratorFunction != null) {
            return service.decorate(decoratorFunction);
        }
        return service;
    }

    /**
     * Builds a new instance of {@link THttpService}.
     */
    public THttpService build() {
        @SuppressWarnings("UnstableApiUsage")
        final Map<String, List<Object>> implementations = Multimaps.asMap(implementationsBuilder.build());

        final ThriftCallService tcs = ThriftCallService.of(implementations);
        return build0(tcs);
    }

    /**
     * Returns a newly-created {@link THttpService} decorator with the properties of this builder.
     */
    public Function<? super RpcService, THttpService> newDecorator() {
        return this::build0;
    }

    private THttpService build0(RpcService tcs) {
        final ImmutableSet.Builder<SerializationFormat> builder = ImmutableSet.builder();
        builder.add(defaultSerializationFormat);
        builder.addAll(otherSerializationFormats);

        return new THttpService(decorate(tcs), defaultSerializationFormat, builder.build(),
                                maxRequestStringLength, maxRequestContainerLength, exceptionHandler);
    }
}
