/*
 * Copyright Red Hat, Inc. and/or its affiliates
 * and other contributors as indicated by the @author tags and
 * the COPYRIGHT.txt file distributed with this work.
 *
 * 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.
 * Modified by
 * @author  Data Virtuality GmbH
 * @date    $Date: 2016-10-21 14:19:11 +0200 (Fr, 21 Okt 2016) $
 *
 */
package org.teiid.resource.adapter.salesforce.transport;

import java.io.ByteArrayInputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.net.*;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.zip.GZIPInputStream;
import java.util.zip.GZIPOutputStream;

import com.sforce.ws.ConnectorConfig;
import com.sforce.ws.transport.JdkHttpTransport;
import com.sforce.ws.transport.LimitingInputStream;
import com.sforce.ws.transport.LimitingOutputStream;
import com.sforce.ws.transport.Transport;
import com.sforce.ws.util.FileUtil;

@SuppressWarnings("nls")
public class SalesforceCustomTransport implements Transport {
    private boolean successful;
    private SalesforceConnectorConfig config;
    private HttpURLConnection connection;

    public void setConfig(ConnectorConfig config) {
        this.config = (SalesforceConnectorConfig)config;
    }

    @Override
    public OutputStream connect(String uri, String soapAction) throws IOException {
        if (soapAction == null) {
            soapAction = "";
        }

        HashMap<String, String> header = new HashMap<>();

        header.put("SOAPAction", "\"" + soapAction + "\"");
        header.put("Content-Type", "text/xml; charset=UTF-8");
        header.put("Accept", "text/xml");

        return connectLocal(uri, header);
    }

    private OutputStream connectLocal(String uri, HashMap<String, String> httpHeaders) throws IOException {
        return connectLocal(uri, httpHeaders, true);
    }

    private OutputStream connectLocal(String uri, HashMap<String, String> httpHeaders, boolean enableCompression)
            throws IOException {
        return wrapOutput(connectRaw(uri, httpHeaders, enableCompression), enableCompression);
    }

    private OutputStream wrapOutput(OutputStream output, boolean enableCompression) throws IOException {
        if (config.getMaxRequestSize() > 0) {
            output = new LimitingOutputStream(config.getMaxRequestSize(), output);
        }

        // when we are writing a zip file we don't bother with compression
        if (enableCompression && config.isCompression()) {
            output = new GZIPOutputStream(output);
        }

        if (config.isTraceMessage()) {
            output = config.teeOutputStream(output);
        }

        return output;
    }

    private OutputStream connectRaw(String uri,
                                    HashMap<String, String> httpHeaders,
                                    boolean enableCompression) throws IOException {

        URL url = new URL(uri);
        connection = JdkHttpTransport.createConnection(config, url, httpHeaders);
        addProxyAuthToConnection(connection, config.getProxyUsername(), config.getProxyPassword());

        connection.setRequestMethod("POST");
        connection.setDoInput(true);
        connection.setDoOutput(true);

        if (enableCompression && config.isCompression()) {
            connection.setRequestProperty("Content-Encoding", "gzip");
            connection.setRequestProperty("Accept-Encoding", "gzip");
        }

        if (config.useChunkedPost()) {
            connection.setRequestProperty("Transfer-Encoding", "chunked");
            connection.setChunkedStreamingMode(4096);
        }

        return connection.getOutputStream();
    }

    @Override
    public InputStream getContent() throws IOException {
        successful = true;
        InputStream in = connection.getInputStream();
        if (connection.getResponseCode() >= 400) {
            successful = false;
        }

        if (!successful) {
            return in;
        }

        String encoding = connection.getContentEncoding();

        if (config.getMaxResponseSize() > 0) {
            in = new LimitingInputStream(config.getMaxResponseSize(), in);
        }

        if ("gzip".equals(encoding)) {
            in = new GZIPInputStream(in);
        }

        if (config.hasMessageHandlers() || config.isTraceMessage()) {
            byte[] bytes = FileUtil.toBytes(in);
            in = new ByteArrayInputStream(bytes);

            if (config.isTraceMessage()) {
                Map<String, List<String>> headersMap = connection.getHeaderFields();
                for (String header : headersMap.keySet()) {
                    config.getTraceStream().print(header);
                    config.getTraceStream().print("=");
                    config.getTraceStream().println(headersMap.get(header));
                }

                config.teeInputStream(bytes);
            }
        }

        return in;
    }

    @Override
    public boolean isSuccessful() {
        return successful;
    }

    @Override
    public OutputStream connect(String endpoint, HashMap<String, String> headers)
            throws IOException {
        return connectLocal(endpoint, headers, true);
    }

    @Override
    public OutputStream connect(String endpoint,
                                HashMap<String, String> httpHeaders, boolean b) throws IOException {
        return connectLocal(endpoint, httpHeaders, b);
    }

    public static void addProxyAuthToConnection(HttpURLConnection connection, String username, String password) {
        if (username != null) {
            connection.setAuthenticator(new Authenticator() {
                @Override
                protected PasswordAuthentication getPasswordAuthentication() {
                    return new PasswordAuthentication(username,
                            password != null
                                    ? password.toCharArray()
                                    : "".toCharArray());
                }
            });
        }
    }
}
