// Dstl (c) Crown Copyright 2017
package uk.gov.dstl.baleen.consumers;

import java.io.BufferedWriter;
import java.io.File;
import java.io.IOException;
import java.io.PrintWriter;
import java.nio.charset.StandardCharsets;

import org.apache.commons.io.output.FileWriterWithEncoding;
import org.apache.uima.UimaContext;
import org.apache.uima.analysis_engine.AnalysisEngineProcessException;
import org.apache.uima.fit.descriptor.ConfigurationParameter;
import org.apache.uima.fit.util.JCasUtil;
import org.apache.uima.jcas.JCas;
import org.apache.uima.jcas.tcas.DocumentAnnotation;
import org.apache.uima.resource.ResourceInitializationException;

import uk.gov.dstl.baleen.types.semantic.Entity;
import uk.gov.dstl.baleen.uima.BaleenConsumer;

/**
 * Create a TSV file listing the files processed and how many entities were found in each file
 *
 * @baleen.javadoc
 */
public class EntityCount extends BaleenConsumer {

  /**
   * The file to write results to
   *
   * @baleen.config entityCount.tsv
   */
  public static final String PARAM_OUTPUT_FILE = "outputFile";

  @ConfigurationParameter(name = PARAM_OUTPUT_FILE, defaultValue = "entityCount.tsv")
  private String outputFile = "entityCount.tsv";

  private File output;

  @Override
  public void doInitialize(UimaContext aContext) throws ResourceInitializationException {
    output = new File(outputFile);

    try {
      if (!output.exists() && !output.createNewFile()) {
        throw new IOException("Can not create output file");
      }

      if (!output.canWrite()) {
        throw new IOException("Can not write to output file");
      }
    } catch (IOException ioe) {
      throw new ResourceInitializationException(ioe);
    }
  }

  @Override
  public void doProcess(JCas jCas) throws AnalysisEngineProcessException {
    DocumentAnnotation da = getDocumentAnnotation(jCas);

    try (PrintWriter pw =
        new PrintWriter(
            new BufferedWriter(new FileWriterWithEncoding(output, StandardCharsets.UTF_8, true)))) {
      int count = JCasUtil.select(jCas, Entity.class).size();
      pw.println(da.getSourceUri() + "\t" + count);
    } catch (IOException e) {
      getMonitor().warn("Unable to write to output", e);
    }
  }

  @Override
  public void doDestroy() {
    output = null;
  }
}
