/**
 *
 * Copyright (c) Microsoft Corporation. All rights reserved.
 * Licensed under the MIT License. See License.txt in the project root for license information.
 *
 */

package com.microsoft.rest;

import com.squareup.okhttp.Headers;
import com.squareup.okhttp.Interceptor;
import com.squareup.okhttp.Request;
import com.squareup.okhttp.Response;

import java.io.IOException;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * An instance of this class enables adding custom headers in client requests
 * when added to the {@link com.squareup.okhttp.OkHttpClient} interceptors.
 */
public class CustomHeaderInterceptor implements Interceptor {
    private Map<String, List<String>> headers;

    /**
     * Initialize an instance of {@link CustomHeaderInterceptor} class.
     */
    public CustomHeaderInterceptor() {
        headers = new HashMap<String, List<String>>();
    }

    /**
     * Add a single header key-value pair.
     *
     * @param name the name of the header.
     * @param value the value of the header.
     * @return the interceptor instance itself.
     */
    public CustomHeaderInterceptor addHeader(String name, String value) {
        this.headers.put(name, Collections.singletonList(value));
        return this;
    }

    /**
     * Add all headers in a {@link Headers} object.
     *
     * @param headers an OkHttp {@link Headers} object.
     * @return the interceptor instance itself.
     */
    public CustomHeaderInterceptor addHeaders(Headers headers) {
        this.headers.putAll(headers.toMultimap());
        return this;
    }

    /**
     * Add all headers in a header map.
     *
     * @param headers a map of headers.
     * @return the interceptor instance itself.
     */
    public CustomHeaderInterceptor addHeaderMap(Map<String, String> headers) {
        for (Map.Entry<String, String> header : headers.entrySet()) {
            this.headers.put(header.getKey(), Collections.singletonList(header.getValue()));
        }
        return this;
    }

    /**
     * Add all headers in a header multimap.
     *
     * @param headers a multimap of headers.
     * @return the interceptor instance itself.
     */
    public CustomHeaderInterceptor addHeaderMultimap(Map<String, List<String>> headers) {
        this.headers.putAll(headers);
        return this;
    }

    /**
     * Remove a header.
     *
     * @param name the name of the header to remove.
     * @return the interceptor instance itself.
     */
    public CustomHeaderInterceptor removeHeader(String name) {
        this.headers.remove(name);
        return this;
    }

    @Override
    public Response intercept(Chain chain) throws IOException {
        Request.Builder builder = chain.request().newBuilder();
        for (Map.Entry<String, List<String>> header : headers.entrySet()) {
            for (String value : header.getValue()) {
                builder = builder.addHeader(header.getKey(), value);
            }
        }
        return chain.proceed(builder.build());
    }
}

@@ -372,16 +372,41 @@
 eption;%0A
+import java.util.Arrays;%0A
 import j
@@ -952,32 +952,329 @@
 ing%3E%3E();%0A    %7D%0A%0A
+    /**%0A     * Initialize an instance of %7B@link CustomHeaderInterceptor%7D class.%0A     *%0A     * @param key the key for the header%0A     * @param value the value of the header%0A     */%0A    public CustomHeaderInterceptor(String key, String value) %7B%0A        this();%0A        addHeader(key, value);%0A    %7D%0A%0A
     /**%0A     * A

Support default headers in azure clients