/*
 * Copyright 2022 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.docs;

import static com.google.common.collect.ImmutableList.toImmutableList;

import java.util.List;
import java.util.Map;
import java.util.stream.Stream;

import com.google.common.collect.ImmutableListMultimap;
import com.google.common.collect.Iterables;
import com.google.common.collect.ListMultimap;
import com.google.common.collect.Streams;

import com.linecorp.armeria.common.HttpHeaders;

final class ExampleSupport {

    private final Map<String, ListMultimap<String, HttpHeaders>> exampleHeaders;
    private final Map<String, ListMultimap<String, String>> exampleRequests;
    private final Map<String, ListMultimap<String, String>> examplePaths;
    private final Map<String, ListMultimap<String, String>> exampleQueries;

    ExampleSupport(Map<String, ListMultimap<String, HttpHeaders>> exampleHeaders,
                   Map<String, ListMultimap<String, String>> exampleRequests,
                   Map<String, ListMultimap<String, String>> examplePaths,
                   Map<String, ListMultimap<String, String>> exampleQueries) {
        this.exampleHeaders = exampleHeaders;
        this.exampleRequests = exampleRequests;
        this.examplePaths = examplePaths;
        this.exampleQueries = exampleQueries;
    }

    /**
     * Adds examples to the {@link ServiceSpecification} and {@link ServiceInfo} objects.
     * @param spec Spec itself.
     * @return new {@link ServiceSpecification} with examples.
     */
    ServiceSpecification addExamples(ServiceSpecification spec) {
        final List<ServiceInfo> serviceWithExample =
                spec.services().stream().map(this::addServiceExamples).collect(toImmutableList());
        final Iterable<HttpHeaders> exampleHeaders =
                Iterables.concat(spec.exampleHeaders(),
                                 this.exampleHeaders.getOrDefault("", ImmutableListMultimap.of()).get(""));

        return new ServiceSpecification(serviceWithExample, spec.enums(), spec.structs(), spec.exceptions(),
                                        exampleHeaders, spec.docServiceRoute());
    }

    private ServiceInfo addServiceExamples(ServiceInfo service) {
        final ListMultimap<String, HttpHeaders> exampleHeaders =
                this.exampleHeaders.getOrDefault(service.name(), ImmutableListMultimap.of());
        final ListMultimap<String, String> exampleRequests =
                this.exampleRequests.getOrDefault(service.name(), ImmutableListMultimap.of());
        final ListMultimap<String, String> examplePaths =
                this.examplePaths.getOrDefault(service.name(), ImmutableListMultimap.of());
        final ListMultimap<String, String> exampleQueries =
                this.exampleQueries.getOrDefault(service.name(), ImmutableListMultimap.of());

        final List<MethodInfo> methodsWithExamples =
                service.methods().stream()
                       .map(m -> new MethodInfo(
                               m.name(), m.returnTypeSignature(), m.parameters(), m.useParameterAsRoot(),
                               m.exceptionTypeSignatures(),
                               m.endpoints(), concatAndDedup(exampleHeaders.get(m.name()), m.exampleHeaders()),
                               // Show the examples added via `DocServiceBuilder` before the examples
                               // generated by the plugin.
                               concatAndDedup(exampleRequests.get(m.name()), m.exampleRequests()),
                               concatAndDedup(examplePaths.get(m.name()), m.examplePaths()),
                               concatAndDedup(exampleQueries.get(m.name()), m.exampleQueries()), m.httpMethod(),
                               m.descriptionInfo(), m.id()
                       ))
                       .collect(toImmutableList());
        final Iterable<HttpHeaders> serviceExampleHeaders =
                Iterables.concat(service.exampleHeaders(), exampleHeaders.get(""));

        return service.withMethods(methodsWithExamples)
                      .withExampleHeaders(serviceExampleHeaders);
    }

    private static <T> Iterable<T> concatAndDedup(Iterable<T> first, Iterable<T> second) {
        return Stream.concat(Streams.stream(first), Streams.stream(second)).distinct()
                     .collect(toImmutableList());
    }
}
