package com.sforce.bulk;

import com.sforce.async.BatchInfo;
import com.sforce.async.BulkConnection;
import com.sforce.async.ConcurrencyMode;
import com.sforce.async.ContentType;
import com.sforce.async.JobInfo;
import com.sforce.async.OperationEnum;
import java.io.ByteArrayInputStream;
import java.io.StringWriter;

public class UpdateStream {
  private StreamHandler handler;
  
  private JobInfo job;
  
  private BulkConnection bulkConnection;
  
  private String[] fieldNames;
  
  private StringWriter writer;
  
  private CsvWriter csvWriter;
  
  private int recordCount;
  
  public static UpdateStream create(StreamHandler handler) throws StreamException {
    return new UpdateStream(handler);
  }
  
  private UpdateStream(StreamHandler handler) throws StreamException {
    this.handler = handler;
    if (handler.getConfig().getAuthEndpoint() == null)
      throw new StreamException("AuthEndpoint not set in config"); 
    if (!handler.getConfig().getAuthEndpoint().contains("/services/Soap/u/"))
      throw new StreamException("Not a valid partner AuthEndpoint " + handler.getConfig().getAuthEndpoint() + 
          " This URL should contain /services/Soap/u/"); 
    LoginHelper loginHelper = new LoginHelper(handler);
    while (handler.shouldContinue()) {
      try {
        loginHelper.doLogin();
        break;
      } catch (Throwable e) {
        handler.error("Failed to login ", e);
      } 
    } 
  }
  
  public void start(String object, OperationEnum operation, ConcurrencyMode concurrencyMode, String[] fieldNames) throws StreamException {
    if (fieldNames == null || fieldNames.length == 0)
      throw new StreamException("field names can not be null/empty"); 
    this.fieldNames = fieldNames;
    while (this.handler.shouldContinue()) {
      try {
        this.bulkConnection = new BulkConnection(this.handler.getConfig());
        this.job = new JobInfo();
        this.job.setObject(object);
        this.job.setOperation(operation);
        this.job.setConcurrencyMode(concurrencyMode);
        this.job.setContentType(ContentType.CSV);
        this.handler.info("Creating bulk api job");
        this.job = this.bulkConnection.createJob(this.job);
        this.handler.info("Bulk api job created with ID : " + this.job.getId());
        break;
      } catch (Throwable e) {
        this.handler.error("Failed to create job ", e);
      } 
    } 
  }
  
  public UpdateResultStream close() throws StreamException {
    if (this.writer != null)
      createBatch(); 
    while (this.handler.shouldContinue()) {
      try {
        this.handler.info("Closing job");
        this.job = this.bulkConnection.closeJob(this.job.getId());
        this.handler.info("Job closed");
        break;
      } catch (Throwable e) {
        this.handler.error("Failed to close job ", e);
      } 
    } 
    return new UpdateResultStream(this.handler, this.bulkConnection, this.job);
  }
  
  public void write(String... values) throws StreamException {
    if (this.job == null)
      throw new StreamException("start() not called"); 
    if (this.writer == null) {
      this.writer = new StringWriter();
      this.csvWriter = new CsvWriter(this.fieldNames, this.writer);
    } 
    this.csvWriter.writeRecord(values);
    this.recordCount++;
    if (this.recordCount > this.handler.getMaxRecordsInBatch()) {
      this.csvWriter.endDocument();
      createBatch();
    } 
  }
  
  private void createBatch() throws StreamException {
    while (this.handler.shouldContinue()) {
      try {
        this.handler.info("Creating Batch");
        BatchInfo batch = this.bulkConnection.createBatchFromStream(this.job, 
            new ByteArrayInputStream(this.writer.getBuffer().toString().getBytes()));
        this.handler.info("Batch created with ID: " + batch.getId());
        this.writer = null;
        this.csvWriter = null;
        this.recordCount = 0;
        break;
      } catch (Throwable e) {
        this.handler.error("Failed to close job ", e);
      } 
    } 
  }
}