package com.sforce.bulk;

import com.sforce.async.BatchInfo;
import com.sforce.async.BatchInfoList;
import com.sforce.async.BatchStateEnum;
import com.sforce.async.BulkConnection;
import com.sforce.async.CSVReader;
import com.sforce.async.JobInfo;
import java.io.IOException;
import java.io.InputStream;
import java.util.ArrayList;

public class UpdateResultStream {
  private JobInfo job;
  
  private BulkConnection bulkConnection;
  
  private StreamHandler handler;
  
  private BatchInfo[] batchList;
  
  private int batchIndex = -1;
  
  private CSVReader resultReader;
  
  public UpdateResultStream(StreamHandler handler, BulkConnection bulkConnection, JobInfo job) throws StreamException {
    this.job = job;
    this.bulkConnection = bulkConnection;
    this.handler = handler;
    loadBatchInfoList(handler, bulkConnection, job);
  }
  
  private void loadBatchInfoList(StreamHandler handler, BulkConnection bulkConnection, JobInfo job) throws StreamException {
    while (handler.shouldContinue()) {
      try {
        BatchInfoList bList = bulkConnection.getBatchInfoList(job.getId());
        this.batchList = bList.getBatchInfo();
        break;
      } catch (Throwable e) {
        handler.error("Failed to get batch list", e);
      } 
    } 
  }
  
  public UpdateResult next() throws StreamException {
    try {
      ArrayList<String> record;
      if (this.resultReader == null || (record = this.resultReader.nextRecord()) == null) {
        this.batchIndex++;
        if (this.batchIndex >= this.batchList.length)
          return null; 
        loadNextBatch();
        record = this.resultReader.nextRecord();
      } 
      if (record == null)
        return null; 
      return new UpdateResult(valueAt(record, 0), booleanAt(record, 1), booleanAt(record, 2), valueAt(record, 3));
    } catch (IOException e) {
      throw new StreamException("Failed to read next record", e);
    } 
  }
  
  private String valueAt(ArrayList<String> record, int index) {
    if (index < record.size())
      return record.get(index); 
    return null;
  }
  
  private boolean booleanAt(ArrayList<String> record, int index) {
    String val = valueAt(record, index);
    return Boolean.parseBoolean(val);
  }
  
  private void waitForNextBatch() throws StreamException {
    int count = 0;
    while (this.handler.shouldContinue()) {
      try {
        BatchInfo bInfo = this.bulkConnection.getBatchInfo(this.job.getId(), this.batchList[this.batchIndex].getId());
        this.handler.info("Batch " + bInfo.getId() + " -- state -- " + bInfo.getState() + " -- try: " + count);
        if (bInfo.getState() == BatchStateEnum.Completed || bInfo.getState() == BatchStateEnum.Failed)
          break; 
        long waitTime = (long)(Math.pow(count, 2.0D) * 1000.0D);
        waitTime = (waitTime > this.handler.getMaxWaitTime()) ? this.handler.getMaxWaitTime() : waitTime;
        Thread.sleep(waitTime);
        count++;
      } catch (Throwable e) {
        this.handler.error("Failed to read result for batch " + this.batchList[this.batchIndex].getId(), e);
      } 
    } 
  }
  
  private void loadNextBatch() throws StreamException {
    waitForNextBatch();
    while (this.handler.shouldContinue()) {
      try {
        InputStream resultStream = 
          this.bulkConnection.getBatchResultStream(this.job.getId(), this.batchList[this.batchIndex].getId());
        this.resultReader = new CSVReader(resultStream);
        this.resultReader.nextRecord();
        break;
      } catch (Throwable e) {
        this.handler.error("Failed to read result for batch " + this.batchList[this.batchIndex].getId(), e);
      } 
    } 
  }
}