package com.sforce.async;

import com.fasterxml.jackson.core.JsonFactory;
import com.fasterxml.jackson.core.JsonGenerator;
import com.fasterxml.jackson.core.JsonParser;
import com.fasterxml.jackson.core.ObjectCodec;
import com.fasterxml.jackson.databind.DeserializationFeature;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.sforce.ws.ConnectionException;
import com.sforce.ws.ConnectorConfig;
import com.sforce.ws.MessageHandler;
import com.sforce.ws.MessageHandlerWithHeaders;
import com.sforce.ws.bind.CalendarCodec;
import com.sforce.ws.bind.TypeMapper;
import com.sforce.ws.parser.PullParserException;
import com.sforce.ws.parser.XmlInputStream;
import com.sforce.ws.parser.XmlOutputStream;
import com.sforce.ws.transport.Transport;
import com.sforce.ws.util.FileUtil;
import java.io.ByteArrayInputStream;
import java.io.File;
import java.io.FileInputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.net.HttpURLConnection;
import java.net.MalformedURLException;
import java.net.URL;
import java.util.HashMap;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.TimeZone;
import java.util.zip.GZIPInputStream;
import java.util.zip.ZipEntry;
import java.util.zip.ZipOutputStream;
import javax.net.ssl.HttpsURLConnection;
import javax.net.ssl.SSLContext;
import javax.xml.namespace.QName;

public class BulkConnection {
  public static final String NAMESPACE = "http://www.force.com/2009/06/asyncapi/dataload";
  
  public static final String SESSION_ID = "X-SFDC-Session";
  
  public static final String XML_CONTENT_TYPE = "application/xml";
  
  public static final String CSV_CONTENT_TYPE = "text/csv";
  
  public static final String JSON_CONTENT_TYPE = "application/json";
  
  public static final String ZIP_XML_CONTENT_TYPE = "zip/xml";
  
  public static final String ZIP_CSV_CONTENT_TYPE = "zip/csv";
  
  public static final String ZIP_JSON_CONTENT_TYPE = "zip/json";
  
  public static final QName JOB_QNAME = new QName("http://www.force.com/2009/06/asyncapi/dataload", "jobInfo");
  
  public static final QName BATCH_QNAME = new QName("http://www.force.com/2009/06/asyncapi/dataload", "batchInfo");
  
  public static final QName BATCH_LIST_QNAME = new QName("http://www.force.com/2009/06/asyncapi/dataload", "batchInfoList");
  
  public static final QName ERROR_QNAME = new QName("http://www.force.com/2009/06/asyncapi/dataload", "error");
  
  private ConnectorConfig config;
  
  private HashMap<String, String> headers = new HashMap<>();
  
  public static final TypeMapper typeMapper = new TypeMapper(null, null, false);
  
  private static final JsonFactory factory = new JsonFactory((ObjectCodec)new ObjectMapper());
  
  public BulkConnection(ConnectorConfig config) throws AsyncApiException {
    if (config == null)
      throw new AsyncApiException("config can not be null", AsyncExceptionCode.ClientInputError); 
    if (config.getRestEndpoint() == null)
      throw new AsyncApiException("rest endpoint cannot be null", 
          AsyncExceptionCode.ClientInputError); 
    this.config = config;
    if (config.getSessionId() == null)
      throw new AsyncApiException("session ID not found", 
          AsyncExceptionCode.ClientInputError); 
  }
  
  public JobInfo createJob(String object, String operation) throws AsyncApiException {
    JobInfo job = new JobInfo();
    job.setObject(object);
    job.setOperation(OperationEnum.valueOf(operation));
    return createJob(job);
  }
  
  public JobInfo createJob(JobInfo job) throws AsyncApiException {
    String endpoint = getRestEndpoint();
    endpoint = String.valueOf(endpoint) + "job/";
    return createOrUpdateJob(job, endpoint);
  }
  
  public JobInfo createJob(JobInfo job, ContentType contentType) throws AsyncApiException {
    String endpoint = getRestEndpoint();
    endpoint = String.valueOf(endpoint) + "job/";
    return createOrUpdateJob(job, endpoint, contentType);
  }
  
  private JobInfo createOrUpdateJob(JobInfo job, String endpoint) throws AsyncApiException {
    return createOrUpdateJob(job, endpoint, (job.getContentType() == null) ? ContentType.XML : 
        job.getContentType());
  }
  
  private JobInfo createOrUpdateJob(JobInfo job, String endpoint, ContentType contentType) throws AsyncApiException {
    try {
      Transport transport = this.config.createTransport();
      if (contentType == ContentType.ZIP_JSON || contentType == ContentType.JSON) {
        OutputStream out = transport.connect(endpoint, getHeaders("application/json"));
        serializeToJson(out, job);
        out.close();
      } else {
        OutputStream out = transport.connect(endpoint, getHeaders("application/xml"));
        XmlOutputStream xout = new AsyncXmlOutputStream(out, true);
        job.write(JOB_QNAME, xout, typeMapper);
        xout.close();
      } 
      InputStream in = transport.getContent();
      if (transport.isSuccessful()) {
        if (contentType == ContentType.ZIP_JSON || contentType == ContentType.JSON)
          return deserializeJsonToObject(in, JobInfo.class); 
        XmlInputStream xin = new XmlInputStream();
        xin.setInput(in, "UTF-8");
        JobInfo result = new JobInfo();
        result.load(xin, typeMapper);
        return result;
      } 
      parseAndThrowException(in, contentType);
    } catch (PullParserException e) {
      throw new AsyncApiException("Failed to create job ", AsyncExceptionCode.ClientInputError, e);
    } catch (IOException e) {
      throw new AsyncApiException("Failed to create job ", AsyncExceptionCode.ClientInputError, e);
    } catch (ConnectionException e) {
      throw new AsyncApiException("Failed to create job ", AsyncExceptionCode.ClientInputError, e);
    } 
    return null;
  }
  
  static void parseAndThrowException(InputStream in, ContentType type) throws AsyncApiException {
    try {
      AsyncApiException exception;
      if (type == ContentType.XML || type == ContentType.ZIP_XML || type == ContentType.CSV || type == ContentType.ZIP_CSV) {
        exception = new AsyncApiException();
        XmlInputStream xin = new XmlInputStream();
        xin.setInput(in, "UTF-8");
        exception.load(xin, typeMapper);
      } else if (type == ContentType.JSON || type == ContentType.ZIP_JSON) {
        JsonParser parser = factory.createJsonParser(in);
        exception = (AsyncApiException)parser.readValueAs(AsyncApiException.class);
      } else {
        throw new AsyncApiException("Server error returned in unknown format", AsyncExceptionCode.ClientInputError);
      } 
      throw exception;
    } catch (PullParserException e) {
      AsyncApiException exception;
      throw new AsyncApiException("Failed to parse exception ", AsyncExceptionCode.ClientInputError, exception);
    } catch (IOException e) {
      throw new AsyncApiException("Failed to parse exception", AsyncExceptionCode.ClientInputError, e);
    } catch (ConnectionException e) {
      throw new AsyncApiException("Failed to parse exception ", AsyncExceptionCode.ClientInputError, e);
    } 
  }
  
  public void addHeader(String headerName, String headerValue) {
    this.headers.put(headerName, headerValue);
  }
  
  private String getRestEndpoint() {
    String endpoint = this.config.getRestEndpoint();
    endpoint = endpoint.endsWith("/") ? endpoint : (String.valueOf(endpoint) + "/");
    return endpoint;
  }
  
  public BatchInfo createBatchFromStream(JobInfo jobInfo, InputStream input) throws AsyncApiException {
    return createBatchFromStreamImpl(jobInfo, input, false);
  }
  
  public BatchInfo createBatchFromZipStream(JobInfo jobInfo, InputStream zipInput) throws AsyncApiException {
    return createBatchFromStreamImpl(jobInfo, zipInput, true);
  }
  
  private BatchInfo createBatchFromStreamImpl(JobInfo jobInfo, InputStream input, boolean isZip) throws AsyncApiException {
    try {
      String endpoint = getRestEndpoint();
      Transport transport = this.config.createTransport();
      endpoint = String.valueOf(endpoint) + "job/" + jobInfo.getId() + "/batch";
      String contentType = getContentTypeString(jobInfo.getContentType(), isZip);
      HashMap<String, String> httpHeaders = getHeaders(contentType);
      boolean allowZipToBeGzipped = false;
      OutputStream out = transport.connect(endpoint, httpHeaders, !(!allowZipToBeGzipped && isZip));
      FileUtil.copy(input, out);
      InputStream result = transport.getContent();
      if (!transport.isSuccessful())
        parseAndThrowException(result, jobInfo.getContentType()); 
      if (jobInfo.getContentType() == ContentType.JSON || jobInfo.getContentType() == ContentType.ZIP_JSON)
        return deserializeJsonToObject(result, BatchInfo.class); 
      return BatchRequest.loadBatchInfo(result);
    } catch (IOException e) {
      throw new AsyncApiException("Failed to create batch", AsyncExceptionCode.ClientInputError, e);
    } catch (PullParserException e) {
      throw new AsyncApiException("Failed to create batch", AsyncExceptionCode.ClientInputError, e);
    } catch (ConnectionException e) {
      throw new AsyncApiException("Failed to create batch", AsyncExceptionCode.ClientInputError, e);
    } 
  }
  
  public BatchInfo createBatchFromDir(JobInfo job, InputStream batchContent, File attachmentDir) throws AsyncApiException {
    List<File> files = FileUtil.listFilesRecursive(attachmentDir, false);
    Map<String, File> fileMap = new HashMap<>(files.size());
    String rootPath = String.valueOf(attachmentDir.getAbsolutePath()) + File.separator;
    for (File f : files) {
      String name = f.getAbsolutePath().replace(rootPath, "");
      fileMap.put(name, f);
    } 
    return createBatchWithFileAttachments(job, batchContent, fileMap);
  }
  
  public BatchInfo createBatchWithFileAttachments(JobInfo jobInfo, InputStream batchContent, File rootDirectory, String... files) throws AsyncApiException {
    Map<String, File> fileMap = new HashMap<>(files.length);
    byte b;
    int i;
    String[] arrayOfString;
    for (i = (arrayOfString = files).length, b = 0; b < i; ) {
      String fileName = arrayOfString[b];
      File f = new File(rootDirectory, fileName);
      fileMap.put(fileName, f);
      b++;
    } 
    return createBatchWithFileAttachments(jobInfo, batchContent, fileMap);
  }
  
  public BatchInfo createBatchWithFileAttachments(JobInfo jobInfo, InputStream batchContent, Map<String, File> attachedFiles) throws AsyncApiException {
    Map<String, InputStream> inputStreamMap = new HashMap<>(attachedFiles.size());
    for (Map.Entry<String, File> entry : attachedFiles.entrySet()) {
      File file = entry.getValue();
      try {
        inputStreamMap.put(entry.getKey(), new FileInputStream(file));
      } catch (IOException e) {
        throw new AsyncApiException("Failed to create batch. Could not read file : " + file, 
            AsyncExceptionCode.ClientInputError, e);
      } 
    } 
    return createBatchWithInputStreamAttachments(jobInfo, batchContent, inputStreamMap);
  }
  
  public BatchInfo createBatchWithInputStreamAttachments(JobInfo jobInfo, InputStream batchContent, Map<String, InputStream> attachments) throws AsyncApiException {
    if (batchContent != null && attachments.get("request.txt") != null)
      throw new AsyncApiException("Request content cannot be included as both input stream and attachment", 
          AsyncExceptionCode.ClientInputError); 
    try {
      String endpoint = getRestEndpoint();
      endpoint = String.valueOf(endpoint) + "job/" + jobInfo.getId() + "/batch";
      Transport transport = this.config.createTransport();
      ZipOutputStream zipOut = new ZipOutputStream(transport.connect(endpoint, getHeaders(getContentTypeString(
                jobInfo.getContentType(), true)), false));
      try {
        if (batchContent != null) {
          zipOut.putNextEntry(new ZipEntry("request.txt"));
          FileUtil.copy(batchContent, zipOut, false);
        } 
        for (Map.Entry<String, InputStream> entry : attachments.entrySet()) {
          zipOut.putNextEntry(new ZipEntry(entry.getKey()));
          FileUtil.copy(entry.getValue(), zipOut, false);
        } 
      } finally {
        zipOut.close();
      } 
      InputStream result = transport.getContent();
      if (jobInfo.getContentType() == ContentType.JSON || jobInfo.getContentType() == ContentType.ZIP_JSON)
        return deserializeJsonToObject(result, BatchInfo.class); 
      return BatchRequest.loadBatchInfo(result);
    } catch (IOException e) {
      throw new AsyncApiException("Failed to create batch", AsyncExceptionCode.ClientInputError, e);
    } catch (PullParserException e) {
      throw new AsyncApiException("Failed to create batch", AsyncExceptionCode.ClientInputError, e);
    } catch (ConnectionException e) {
      throw new AsyncApiException("Failed to create batch", AsyncExceptionCode.ClientInputError, e);
    } 
  }
  
  public BatchInfo createBatchFromForeignCsvStream(JobInfo jobInfo, InputStream input, String charSet) throws AsyncApiException {
    try {
      String endpoint = getRestEndpoint();
      Transport transport = this.config.createTransport();
      endpoint = String.valueOf(endpoint) + "job/" + jobInfo.getId() + "/batch";
      String contentType = getContentTypeString(ContentType.CSV, false);
      if (charSet != null)
        contentType = String.valueOf(contentType) + ";charset=" + charSet; 
      HashMap<String, String> httpHeaders = getHeaders(contentType);
      boolean allowZipToBeGzipped = false;
      OutputStream out = transport.connect(endpoint, httpHeaders, false);
      FileUtil.copy(input, out);
      InputStream result = transport.getContent();
      if (!transport.isSuccessful())
        parseAndThrowException(result, jobInfo.getContentType()); 
      return BatchRequest.loadBatchInfo(result);
    } catch (IOException e) {
      throw new AsyncApiException("Failed to create batch", AsyncExceptionCode.ClientInputError, e);
    } catch (PullParserException e) {
      throw new AsyncApiException("Failed to create batch", AsyncExceptionCode.ClientInputError, e);
    } catch (ConnectionException e) {
      throw new AsyncApiException("Failed to create batch", AsyncExceptionCode.ClientInputError, e);
    } 
  }
  
  public void createTransformationSpecFromStream(JobInfo jobInfo, InputStream input) throws AsyncApiException {
    try {
      String endpoint = getRestEndpoint();
      Transport transport = this.config.createTransport();
      endpoint = String.valueOf(endpoint) + "job/" + jobInfo.getId() + "/spec";
      String contentType = getContentTypeString(ContentType.CSV, false);
      HashMap<String, String> httpHeaders = getHeaders(contentType);
      boolean allowZipToBeGzipped = false;
      OutputStream out = transport.connect(endpoint, httpHeaders, false);
      FileUtil.copy(input, out);
      InputStream result = transport.getContent();
      if (!transport.isSuccessful())
        parseAndThrowException(result, jobInfo.getContentType()); 
    } catch (IOException e) {
      throw new AsyncApiException("Failed to create transformation specification", AsyncExceptionCode.ClientInputError, e);
    } catch (ConnectionException e) {
      throw new AsyncApiException("Failed to create transformation specification", AsyncExceptionCode.ClientInputError, e);
    } 
  }
  
  private String getContentTypeString(ContentType contentType, boolean isZip) throws AsyncApiException {
    ContentType ct = (contentType == null) ? ContentType.XML : contentType;
    if (isZip) {
      switch (ct) {
        case ZIP_CSV:
          return "zip/csv";
        case ZIP_XML:
          return "zip/xml";
        case ZIP_JSON:
          return "zip/json";
      } 
      throw new AsyncApiException("Invalid zip content type: " + contentType, 
          AsyncExceptionCode.ClientInputError);
    } 
    switch (ct) {
      case XML:
        return "application/xml";
      case null:
        return "text/csv";
      case JSON:
        return "application/json";
    } 
    throw new AsyncApiException("Not expecting zip content type: " + contentType, 
        AsyncExceptionCode.ClientInputError);
  }
  
  private HashMap<String, String> getHeaders(String contentType) {
    HashMap<String, String> newMap = new HashMap<>();
    for (Map.Entry<String, String> entry : this.headers.entrySet())
      newMap.put(entry.getKey(), entry.getValue()); 
    newMap.put("Content-Type", contentType);
    newMap.put("X-SFDC-Session", this.config.getSessionId());
    return newMap;
  }
  
  public BatchRequest createBatch(JobInfo job) throws AsyncApiException {
    try {
      String endpoint = getRestEndpoint();
      Transport transport = this.config.createTransport();
      endpoint = String.valueOf(endpoint) + "job/" + job.getId() + "/batch";
      ContentType ct = job.getContentType();
      if (ct != null && ct != ContentType.XML && ct != ContentType.JSON)
        throw new AsyncApiException(
            "This method can only be used with xml or JSON content type", AsyncExceptionCode.ClientInputError); 
      String jobContentType = "";
      if (ct == null) {
        jobContentType = "application/xml";
      } else {
        OutputStream out;
        switch (ct) {
          case JSON:
            jobContentType = "application/json";
            out = transport.connect(endpoint, getHeaders(jobContentType));
            return new BatchRequest(transport, out);
        } 
        jobContentType = "application/xml";
      } 
      OutputStream outputStream = transport.connect(endpoint, getHeaders(jobContentType));
      return new BatchRequest(transport, outputStream);
    } catch (IOException e) {
      throw new AsyncApiException("Failed to create batch", AsyncExceptionCode.ClientInputError, e);
    } catch (ConnectionException x) {
      throw new AsyncApiException("Failed to create batch", AsyncExceptionCode.ClientInputError, x);
    } 
  }
  
  public CsvBatchRequest createCsvBatch(JobInfo job) throws AsyncApiException {
    try {
      String endpoint = getRestEndpoint();
      Transport transport = this.config.createTransport();
      endpoint = String.valueOf(endpoint) + "job/" + job.getId() + "/batch";
      ContentType ct = job.getContentType();
      if (ct != null && ct != ContentType.CSV)
        throw new AsyncApiException(
            "This method can only be used with csv content type", AsyncExceptionCode.ClientInputError); 
      OutputStream out = transport.connect(endpoint, getHeaders("text/csv"));
      return new CsvBatchRequest(transport, out);
    } catch (IOException e) {
      throw new AsyncApiException("Failed to create batch", AsyncExceptionCode.ClientInputError, e);
    } catch (ConnectionException e) {
      throw new AsyncApiException("Failed to create batch", AsyncExceptionCode.ClientInputError, e);
    } 
  }
  
  public TransformationSpecRequest createTransformationSpec(JobInfo job) throws AsyncApiException {
    try {
      String endpoint = getRestEndpoint();
      Transport transport = this.config.createTransport();
      endpoint = String.valueOf(endpoint) + "job/" + job.getId() + "/spec";
      ContentType ct = job.getContentType();
      if (ct != null && ct != ContentType.CSV)
        throw new AsyncApiException(
            "This method can only be used with csv content type", AsyncExceptionCode.ClientInputError); 
      boolean allowZipToBeGzipped = false;
      OutputStream out = transport.connect(endpoint, getHeaders("text/csv"), false);
      return new TransformationSpecRequest(transport, out);
    } catch (IOException e) {
      throw new AsyncApiException("Failed to create transformation spec", AsyncExceptionCode.ClientInputError, e);
    } catch (ConnectionException e) {
      throw new AsyncApiException("Failed to create transformation spec", AsyncExceptionCode.ClientInputError, e);
    } 
  }
  
  public BatchInfoList getBatchInfoList(String jobId) throws AsyncApiException {
    return getBatchInfoList(jobId, ContentType.XML);
  }
  
  public BatchInfoList getBatchInfoList(String jobId, ContentType contentType) throws AsyncApiException {
    try {
      String endpoint = String.valueOf(getRestEndpoint()) + "job/" + jobId + "/batch/";
      URL url = new URL(endpoint);
      InputStream stream = doHttpGet(url);
      if (contentType == ContentType.JSON || contentType == ContentType.ZIP_JSON)
        return deserializeJsonToObject(stream, BatchInfoList.class); 
      XmlInputStream xin = new XmlInputStream();
      xin.setInput(stream, "UTF-8");
      BatchInfoList result = new BatchInfoList();
      result.load(xin, typeMapper);
      return result;
    } catch (IOException e) {
      throw new AsyncApiException("Failed to get batch info list ", AsyncExceptionCode.ClientInputError, e);
    } catch (PullParserException e) {
      throw new AsyncApiException("Failed to get batch info list ", AsyncExceptionCode.ClientInputError, e);
    } catch (ConnectionException e) {
      throw new AsyncApiException("Failed to get batch info list ", AsyncExceptionCode.ClientInputError, e);
    } 
  }
  
  public BatchInfo getBatchInfo(String jobId, String batchId) throws AsyncApiException {
    return getBatchInfo(jobId, batchId, ContentType.XML);
  }
  
  public BatchInfo getBatchInfo(String jobId, String batchId, ContentType contentType) throws AsyncApiException {
    try {
      String endpoint = String.valueOf(getRestEndpoint()) + "job/" + jobId + "/batch/" + batchId;
      URL url = new URL(endpoint);
      InputStream stream = doHttpGet(url);
      if (contentType == ContentType.JSON || contentType == ContentType.ZIP_JSON)
        return deserializeJsonToObject(stream, BatchInfo.class); 
      XmlInputStream xin = new XmlInputStream();
      xin.setInput(stream, "UTF-8");
      BatchInfo result = new BatchInfo();
      result.load(xin, typeMapper);
      return result;
    } catch (IOException e) {
      throw new AsyncApiException("Failed to parse batch info ", AsyncExceptionCode.ClientInputError, e);
    } catch (PullParserException e) {
      throw new AsyncApiException("Failed to parse batch info ", AsyncExceptionCode.ClientInputError, e);
    } catch (ConnectionException e) {
      throw new AsyncApiException("Failed to parse batch info ", AsyncExceptionCode.ClientInputError, e);
    } 
  }
  
  public BatchResult getBatchResult(String jobId, String batchId) throws AsyncApiException {
    return getBatchResult(jobId, batchId, ContentType.XML);
  }
  
  public BatchResult getBatchResult(String jobId, String batchId, ContentType contentType) throws AsyncApiException {
    try {
      InputStream stream = doHttpGet(buildBatchResultURL(jobId, batchId));
      if (contentType == ContentType.JSON || contentType == ContentType.ZIP_JSON) {
        BatchResult batchResult = new BatchResult();
        Result[] results = deserializeJsonToObject(stream, (Class)Result[].class);
        batchResult.setResult(results);
        return batchResult;
      } 
      XmlInputStream xin = new XmlInputStream();
      xin.setInput(stream, "UTF-8");
      BatchResult result = new BatchResult();
      result.load(xin, typeMapper);
      return result;
    } catch (PullParserException e) {
      throw new AsyncApiException("Failed to parse result ", AsyncExceptionCode.ClientInputError, e);
    } catch (IOException e) {
      throw new AsyncApiException("Failed to get result ", AsyncExceptionCode.ClientInputError, e);
    } catch (ConnectionException e) {
      throw new AsyncApiException("Failed to get result ", AsyncExceptionCode.ClientInputError, e);
    } 
  }
  
  public InputStream getBatchResultStream(String jobId, String batchId) throws AsyncApiException {
    try {
      String endpoint = String.valueOf(getRestEndpoint()) + "job/" + jobId + "/batch/" + batchId + "/result";
      URL url = new URL(endpoint);
      return doHttpGet(url);
    } catch (IOException e) {
      throw new AsyncApiException("Failed to get result ", AsyncExceptionCode.ClientInputError, e);
    } 
  }
  
  public URL buildBatchResultURL(String jobId, String batchId) throws AsyncApiException {
    try {
      return new URL(String.valueOf(getRestEndpoint()) + "job/" + jobId + "/batch/" + batchId + "/result");
    } catch (MalformedURLException e) {
      throw new AsyncApiException("Failed to construct URL for getting batch results: " + e.getMessage(), 
          AsyncExceptionCode.ClientInputError, e);
    } 
  }
  
  public InputStream getBatchRequestInputStream(String jobId, String batchId) throws AsyncApiException {
    try {
      String endpoint = String.valueOf(getRestEndpoint()) + "job/" + jobId + "/batch/" + batchId + "/request";
      URL url = new URL(endpoint);
      return doHttpGet(url);
    } catch (IOException e) {
      throw new AsyncApiException("Failed to get request ", AsyncExceptionCode.ClientInputError, e);
    } 
  }
  
  public QueryResultList getQueryResultList(String jobId, String batchId) throws AsyncApiException {
    return getQueryResultList(jobId, batchId, ContentType.XML);
  }
  
  public QueryResultList getQueryResultList(String jobId, String batchId, ContentType contentType) throws AsyncApiException {
    InputStream stream = getBatchResultStream(jobId, batchId);
    try {
      if (contentType == ContentType.JSON || contentType == ContentType.ZIP_JSON) {
        String[] results = deserializeJsonToObject(stream, (Class)String[].class);
        QueryResultList list = new QueryResultList();
        list.setResult(results);
        return list;
      } 
      XmlInputStream xin = new XmlInputStream();
      xin.setInput(stream, "UTF-8");
      QueryResultList result = new QueryResultList();
      result.load(xin, typeMapper);
      return result;
    } catch (ConnectionException e) {
      throw new AsyncApiException("Failed to parse query result list ", AsyncExceptionCode.ClientInputError, e);
    } catch (PullParserException e) {
      throw new AsyncApiException("Failed to parse query result list ", AsyncExceptionCode.ClientInputError, e);
    } catch (IOException e) {
      throw new AsyncApiException("Failed to parse query result list ", AsyncExceptionCode.ClientInputError, e);
    } 
  }
  
  public InputStream getQueryResultStream(String jobId, String batchId, String resultId) throws AsyncApiException {
    try {
      return doHttpGet(buildQueryResultURL(jobId, batchId, resultId));
    } catch (IOException e) {
      throw new AsyncApiException("Failed to get result ", AsyncExceptionCode.ClientInputError, e);
    } 
  }
  
  public URL buildQueryResultURL(String jobId, String batchId, String resultId) throws AsyncApiException {
    try {
      return new URL(String.valueOf(getRestEndpoint()) + "job/" + jobId + "/batch/" + batchId + "/result" + "/" + resultId);
    } catch (MalformedURLException e) {
      throw new AsyncApiException("Failed to construct URL for getting query result: " + e.getMessage(), 
          AsyncExceptionCode.ClientInputError, e);
    } 
  }
  
  private InputStream doHttpGet(URL url) throws IOException, AsyncApiException {
    InputStream in;
    HttpURLConnection connection = this.config.createConnection(url, null);
    SSLContext sslContext = this.config.getSslContext();
    if (sslContext != null && connection instanceof HttpsURLConnection)
      ((HttpsURLConnection)connection).setSSLSocketFactory(sslContext.getSocketFactory()); 
    connection.setRequestProperty("X-SFDC-Session", this.config.getSessionId());
    boolean success = true;
    try {
      in = connection.getInputStream();
    } catch (IOException e) {
      success = false;
      in = connection.getErrorStream();
    } 
    String encoding = connection.getHeaderField("Content-Encoding");
    if ("gzip".equals(encoding))
      in = new GZIPInputStream(in); 
    if (this.config.isTraceMessage() || this.config.hasMessageHandlers()) {
      byte[] bytes = FileUtil.toBytes(in);
      in = new ByteArrayInputStream(bytes);
      if (this.config.hasMessageHandlers()) {
        Iterator<MessageHandler> it = this.config.getMessagerHandlers();
        while (it.hasNext()) {
          MessageHandler handler = it.next();
          if (handler instanceof MessageHandlerWithHeaders) {
            ((MessageHandlerWithHeaders)handler).handleRequest(url, new byte[0], null);
            ((MessageHandlerWithHeaders)handler).handleResponse(url, bytes, connection.getHeaderFields());
            continue;
          } 
          handler.handleRequest(url, new byte[0]);
          handler.handleResponse(url, bytes);
        } 
      } 
      if (this.config.isTraceMessage()) {
        this.config.getTraceStream().println(url.toExternalForm());
        Map<String, List<String>> headers = connection.getHeaderFields();
        for (Map.Entry<String, List<String>> entry : headers.entrySet()) {
          StringBuffer sb = new StringBuffer();
          List<String> values = entry.getValue();
          if (values != null)
            for (String v : values)
              sb.append(v);  
          this.config.getTraceStream().println(String.valueOf(entry.getKey()) + ": " + sb.toString());
        } 
        this.config.teeInputStream(bytes);
      } 
    } 
    if (!success) {
      ContentType type = null;
      String contentTypeHeader = connection.getContentType();
      if (contentTypeHeader != null)
        if (contentTypeHeader.contains("application/xml")) {
          type = ContentType.XML;
        } else if (contentTypeHeader.contains("application/json")) {
          type = ContentType.JSON;
        }  
      parseAndThrowException(in, type);
    } 
    return in;
  }
  
  public JobInfo getJobStatus(String jobId) throws AsyncApiException {
    return getJobStatus(jobId, ContentType.XML);
  }
  
  public JobInfo getJobStatus(String jobId, ContentType contentType) throws AsyncApiException {
    try {
      String endpoint = getRestEndpoint();
      endpoint = String.valueOf(endpoint) + "job/" + jobId;
      URL url = new URL(endpoint);
      InputStream in = doHttpGet(url);
      if (contentType == ContentType.JSON || contentType == ContentType.ZIP_JSON)
        return deserializeJsonToObject(in, JobInfo.class); 
      JobInfo result = new JobInfo();
      XmlInputStream xin = new XmlInputStream();
      xin.setInput(in, "UTF-8");
      result.load(xin, typeMapper);
      return result;
    } catch (PullParserException e) {
      throw new AsyncApiException("Failed to get job status ", AsyncExceptionCode.ClientInputError, e);
    } catch (IOException e) {
      throw new AsyncApiException("Failed to get job status ", AsyncExceptionCode.ClientInputError, e);
    } catch (ConnectionException e) {
      throw new AsyncApiException("Failed to get job status ", AsyncExceptionCode.ClientInputError, e);
    } 
  }
  
  static void serializeToJson(OutputStream out, Object value) throws IOException {
    JsonGenerator generator = factory.createJsonGenerator(out);
    ObjectMapper mapper = new ObjectMapper();
    mapper.setDateFormat(CalendarCodec.getDateFormat());
    mapper.writeValue(generator, value);
  }
  
  static <T> T deserializeJsonToObject(InputStream in, Class<T> tmpClass) throws IOException {
    ObjectMapper mapper = new ObjectMapper();
    mapper.configure(DeserializationFeature.ACCEPT_SINGLE_VALUE_AS_ARRAY, true);
    mapper.setTimeZone(TimeZone.getTimeZone("GMT"));
    return (T)mapper.readValue(in, tmpClass);
  }
  
  public JobInfo abortJob(String jobId) throws AsyncApiException {
    JobInfo job = new JobInfo();
    job.setId(jobId);
    job.setState(JobStateEnum.Aborted);
    return updateJob(job);
  }
  
  public JobInfo closeJob(String jobId) throws AsyncApiException {
    JobInfo job = new JobInfo();
    job.setId(jobId);
    job.setState(JobStateEnum.Closed);
    return updateJob(job);
  }
  
  public JobInfo updateJob(JobInfo job) throws AsyncApiException {
    return updateJob(job, ContentType.XML);
  }
  
  public JobInfo updateJob(JobInfo job, ContentType contentType) throws AsyncApiException {
    String endpoint = getRestEndpoint();
    endpoint = String.valueOf(endpoint) + "job/" + job.getId();
    return createOrUpdateJob(job, endpoint, contentType);
  }
  
  public ConnectorConfig getConfig() {
    return this.config;
  }
}