package com.acumenvelocity.ath.solr.doc;
   
   import java.util.Date;
   import java.util.List;
   import java.util.UUID;
   
   import org.apache.solr.client.solrj.SolrClient;
   import org.apache.solr.common.SolrDocument;
   import org.apache.solr.common.SolrInputDocument;
  
  import com.acumenvelocity.ath.common.Const;
  import com.acumenvelocity.ath.common.ConversionUtil;
  import com.acumenvelocity.ath.common.JacksonUtil;
  import com.acumenvelocity.ath.common.Log;
  import com.acumenvelocity.ath.common.SolrUtil;
  import com.acumenvelocity.ath.model.MtTargetInfo;
  import com.acumenvelocity.ath.model.Origin;
  import com.acumenvelocity.ath.model.x.LayeredTextX;
  import com.acumenvelocity.ath.solr.AthIndex;
  import com.acumenvelocity.ath.solr.SolrIndexWriterStep;
  
  import net.sf.okapi.common.Event;
  import net.sf.okapi.common.LocaleId;
  import net.sf.okapi.common.annotation.AltTranslation;
  import net.sf.okapi.common.annotation.AltTranslationsAnnotation;
  import net.sf.okapi.common.query.MatchType;
  import net.sf.okapi.common.resource.ITextUnit;
  import net.sf.okapi.common.resource.Segment;
  
  public class SolrDocWriterStep extends SolrIndexWriterStep {
  
    private UUID docId;
    private String docFileName;
    private UUID userId;
    private boolean newDoc;
    private long position;
  
    public SolrDocWriterStep(UUID docId, String docFileName, UUID userId, boolean newDoc) {
      super(AthIndex.getIndex().getClient(), Const.ATH_SOLR_BATCH_SIZE,
          Const.SOLR_CORE_ATH_DOC_SEGMENTS, true, false);
  
      this.docId = docId;
      this.docFileName = docFileName;
      this.userId = userId;
      this.newDoc = newDoc;
    }
  
    @Override
    public String getName() {
      return "Solr Document Writer Step";
    }
  
    @Override
    public String getDescription() {
      return "Writes text units to Solr in batches";
    }
  
    @Override
    protected Event handleStartDocument(Event event) {
      // Reset the position within the document
      position = 1;
      return super.handleStartDocument(event);
    }
  
    @Override
    protected Event handleEndBatch(Event event) {
      Event ev = super.handleEndBatch(event);
  
      // Update document's updated_at to current time so the modified segments summary returns an
      // empty list (the newly imported document has unmodified segments only)
      SolrDocument existingDoc = SolrUtil.getDocumentByDocId(docId);
      SolrInputDocument doc = SolrUtil.toInputDocument(existingDoc);
      doc.setField(Const.ATH_PROP_UPDATED_AT, new Date());
  
      try {
        AthIndex.getSolr().getClient().add(Const.SOLR_CORE_ATH_DOCS, doc);
        AthIndex.getSolr().getClient().commit(Const.SOLR_CORE_ATH_DOCS);
  
      } catch (Exception e) {
        Log.warn(SolrIndexWriterStep.class, "Unable to update document {} -- {}", docId,
            e.getMessage());
      }
  
      return ev;
    }
  
    @Override
    protected SolrInputDocument getSolrDocument(SolrClient solrClient, String coreName,
        String sourceWithCodes) {
  
      SolrInputDocument doc = null;
  
      if (newDoc) {
        doc = new SolrInputDocument();
  
      } else {
        try {
          doc = SolrUtil.getDocumentBySolrId(solrClient, coreName,
              SolrUtil.buildDocSegSolrId(docId, position));
 
       } catch (Exception e) {
         Log.warn(getClass(), "Error getting a Solr document for '{}' (position: {}) -- {}",
             sourceWithCodes, position, e.getMessage());
       }
 
       if (doc == null) {
         doc = new SolrInputDocument();
       }
     }
 
     return doc;
   }
 
   @Override
   protected void populateSolrDocument(SolrInputDocument doc, ITextUnit tu, Segment sseg,
       Segment tseg, LocaleId srcLoc, LocaleId trgLoc) throws Exception {
 
     // XXX Keep in sync with BaseMergerStep
     if (sseg == null) {
       return;
     }
 
     if (tseg == null && isRequireTarget()) {
       position++;
       return;
     }
 
     // Unique id (docId + position) to overwrite an existing doc segment if the position is the same
     SolrUtil.safeSetField(tu, doc, Const.ATH_PROP_SOLR_ID,
         SolrUtil.buildDocSegSolrId(docId, position));
 
     SolrUtil.safeSetField(tu, doc, Const.ATH_PROP_DOC_SEG_ID, UUID.randomUUID().toString());
 
     SolrUtil.safeSetField(tu, doc, Const.ATH_PROP_DOC_ID, docId.toString());
     SolrUtil.safeSetField(tu, doc, Const.ATH_PROP_TU_ID, tu.getId());
     SolrUtil.safeSetField(tu, doc, Const.ATH_PROP_DOC_FILE_NAME, docFileName);
 
     doc.setField(Const.ATH_PROP_POSITION, position);
 
     SolrUtil.safeSetField(tu, doc, Const.ATH_PROP_SRC_LANG, srcLoc.toString());
     SolrUtil.safeSetField(tu, doc, Const.ATH_PROP_TRG_LANG, trgLoc.toString());
 
     LayeredTextX slt = ConversionUtil.toLayeredText(sseg.getContent(), srcLoc);
     SolrUtil.safeSetField(tu, doc, Const.ATH_PROP_SOURCE, slt.getText());
     SolrUtil.safeSetField(tu, doc, Const.ATH_PROP_SOURCE_WITH_CODES, slt.getTextWithCodes());
     SolrUtil.safeSetField(tu, doc, Const.ATH_PROP_SOURCE_JSON, JacksonUtil.toJson(slt, false));
 
     if (tseg != null) {
       LayeredTextX tlt = ConversionUtil.toLayeredText(tseg.getContent(), trgLoc);
       SolrUtil.safeSetField(tu, doc, Const.ATH_PROP_TARGET, tlt.getText(), false);
 
       SolrUtil.safeSetField(tu, doc, Const.ATH_PROP_TARGET_WITH_CODES, tlt.getTextWithCodes(),
           false);
 
       SolrUtil.safeSetField(tu, doc, Const.ATH_PROP_TARGET_JSON, JacksonUtil.toJson(tlt, false));
 
       AltTranslationsAnnotation ata = tseg.getAnnotation(AltTranslationsAnnotation.class);
 
       // Write origin
       if (ata != null) {
         ata.sort();
         AltTranslation at = ata.getFirst();
 
         if (at.getType() == Const.AL_MATCH_TYPE) {
           SolrUtil.safeSetField(tu, doc, Const.ATH_PROP_ORIGIN, Origin.AL.toString());
 
         } else if (at.getType() == MatchType.MT) {
           SolrUtil.safeSetField(tu, doc, Const.ATH_PROP_ORIGIN, Origin.MT.toString());
 
         } else if (at.getCombinedScore() > 0) {
           SolrUtil.safeSetField(tu, doc, Const.ATH_PROP_ORIGIN, Origin.TM.toString());
           doc.setField(Const.ATH_PROP_TM_MATCH_SCORE, at.getCombinedScore());
         }
 
         List<MtTargetInfo> altTrans = ConversionUtil.toMtTargets(ata);
 
         SolrUtil.safeSetField(tu, doc, Const.ATH_PROP_ALT_TRANS_JSON,
             JacksonUtil.toJson(altTrans, false));
 
         doc.setField(Const.ATH_PROP_ALT_TRANS_INDEX, -1);
       }
     }
 
     if (newDoc) {
       SolrUtil.safeSetField(tu, doc, Const.ATH_PROP_CREATED_BY, userId.toString());
       doc.setField(Const.ATH_PROP_CREATED_AT, new Date());
 
     } else {
       SolrUtil.safeSetField(tu, doc, Const.ATH_PROP_UPDATED_BY, userId.toString());
       doc.setField(Const.ATH_PROP_UPDATED_AT, new Date());
     }
 
     position++;
   }
 
 }