package com.acumenvelocity.ath.solr;
   
   import java.util.ArrayList;
   import java.util.Collection;
   import java.util.List;
   
   import org.apache.solr.client.solrj.SolrClient;
   import org.apache.solr.common.SolrInputDocument;
   
  import com.acumenvelocity.ath.common.Log;
  import com.acumenvelocity.ath.common.OkapiUtil;
  import com.acumenvelocity.ath.steps.BaseMergerStep;
  
  import net.sf.okapi.common.Event;
  import net.sf.okapi.common.LocaleId;
  import net.sf.okapi.common.pipeline.BasePipelineStep;
  import net.sf.okapi.common.resource.ISegments;
  import net.sf.okapi.common.resource.ITextUnit;
  import net.sf.okapi.common.resource.Segment;
  import net.sf.okapi.common.resource.TextContainer;
  
  /**
   * Abstract base class for indexing text units into a Solr instance using batched operations.
   * Subclasses provide custom field population logic for index documents.
   */
  public abstract class SolrIndexWriterStep extends BasePipelineStep {
  
    private SolrClient solrClient;
    private String coreName;
    private boolean autoCommit;
    private int tuBatchSize;
    private List<ITextUnit> tuBatch;
    private boolean requireTarget = true;
  
    protected abstract SolrInputDocument getSolrDocument(SolrClient solrClient, String coreName,
        String sourceWithCodes);
  
    /**
     * Implement this method in subclasses to define how index document fields are set based on the
     * provided text unit and segments.
     * 
     * @param doc    The index document to fill with data.
     * @param tu     The text unit being processed.
     * @param sseg   The source segment.
     * @param srcLoc The source locale.
     * @param tseg   The corresponding target segment (may be null).
     * @param trgLoc The target locale.
     * @throws Exception If an error occurs during field population.
     */
    protected abstract void populateSolrDocument(SolrInputDocument doc, ITextUnit tu,
        Segment sseg, Segment tseg, LocaleId srcLoc, LocaleId trgLoc) throws Exception;
  
    /**
     * Constructs the step with batch capacity, index identifier, and auto-commit option.
     * 
     * @param tuBatchSize Maximum number of TUs per TU batch.
     * @param coreName    The identifier for the Solr index.
     * @param autoCommit  Whether to commit changes automatically after each batch.
     */
    public SolrIndexWriterStep(SolrClient solrClient, int tuBatchSize, String coreName,
        boolean autoCommit, boolean requireTarget) {
  
      this.solrClient = solrClient;
      this.tuBatch = new ArrayList<>();
      this.tuBatchSize = tuBatchSize;
      this.coreName = coreName;
      this.autoCommit = autoCommit;
      this.requireTarget = requireTarget;
    }
  
    @Override
    public String getName() {
      return "Solr Indexing Processor";
    }
  
    @Override
    public String getDescription() {
      return "Handles batch indexing of text units into a Solr collection.";
    }
  
    @Override
    protected Event handleStartBatch(Event event) {
      tuBatch.clear();
      return event;
    }
  
    @Override
    protected Event handleTextUnit(Event event) {
      ITextUnit tu = event.getTextUnit();
  
      if (tu != null) {
        tuBatch.add(tu);
  
        if (tuBatch.size() >= tuBatchSize) {
          processTuBatch();
        }
      }
      return event;
    }
 
   @Override
   protected Event handleEndBatch(Event event) {
     if (!tuBatch.isEmpty()) {
       processTuBatch();
     }
 
     return event;
   }
 
   /**
    * Executes the indexing operation for the current pending batch.
    */
   private void processTuBatch() {
     if (tuBatch.isEmpty()) {
       Log.warn(SolrIndexWriterStep.class, "Attempted to index an empty batch");
       return;
     }
 
     if (solrClient == null) {
       throw new IllegalStateException("Index client must be configured prior to use.");
     }
 
     try {
       Collection<SolrInputDocument> indexDocs = generateIndexDocuments(tuBatch);
 
       if (indexDocs.size() < tuBatch.size()) {
         Log.warn(SolrIndexWriterStep.class, "Generated {} index documents from {} text units",
             indexDocs.size(), tuBatch.size());
       }
 
       Log.debug(SolrIndexWriterStep.class, "Indexing {} documents into Solr", indexDocs.size());
 
       if (!indexDocs.isEmpty()) {
         solrClient.add(coreName, indexDocs);
 
         if (autoCommit) {
           solrClient.commit(coreName);
         }
       }
 
       Log.debug(SolrIndexWriterStep.class, "Completed indexing of {} documents", indexDocs.size());
 
     } catch (Exception e) {
       throw new RuntimeException("Indexing operation failed", e);
 
     } finally {
       tuBatch.clear();
     }
   }
 
   /**
    * Transforms a list of text units into a collection of index documents.
    * 
    * @param tus The text units to transform.
    * @return Collection of prepared index documents.
    */
   private Collection<SolrInputDocument> generateIndexDocuments(List<ITextUnit> tus) {
     Collection<SolrInputDocument> docs = new ArrayList<>();
 
     for (ITextUnit tu : tus) {
       if (tu == null) {
         continue;
       }
 
       // XXX Keep in sync with BaseMergerStep
       if (!BaseMergerStep.checkTu(tu)) {
         continue;
       }
 
       Collection<LocaleId> trgLocs = tu.getTargetLocales();
 
       if (trgLocs.isEmpty() && requireTarget) {
         Log.debug(SolrIndexWriterStep.class,
             "Omitting text unit without output languages since target content is required");
         continue;
       }
 
       if (trgLocs.isEmpty()) {
         processTuSegments(tu, null, docs);
 
       } else {
         for (LocaleId lang : trgLocs) {
           TextContainer target = tu.getTarget(lang);
 
           if (requireTarget && (target == null || !target.hasText(false))) {
             Log.debug(SolrIndexWriterStep.class,
                 "Omitting text unit for language {} due to missing content", lang);
 
             continue;
           }
 
           processTuSegments(tu, lang, docs);
         }
       }
     }
 
     return docs;
   }
 
   /**
    * Handles segment processing for a given text unit and optional output language.
    * 
    * @param tu     The text unit to process.
    * @param trgLoc The output language (may be null).
    * @param docs   The collection to add generated documents to.
    */
   private void processTuSegments(ITextUnit tu, LocaleId trgLoc,
       Collection<SolrInputDocument> docs) {
 
     TextContainer source = tu.getSource();
     TextContainer target = OkapiUtil.safeGetTarget(tu, trgLoc);
 
     ISegments sourceSegments = source.getSegments();
 
     for (Segment sseg : sourceSegments) {
       Segment tseg = null;
 
       if (target != null) {
         tseg = target.getSegments().get(sseg.getId());
 
         if (requireTarget && tseg == null) {
           Log.debug(SolrIndexWriterStep.class, "Omitting segment {} due to missing output segment",
               sseg.getId());
 
           continue;
         }
       }
 
       if (!requireTarget && tseg == null) {
         Log.debug(SolrIndexWriterStep.class, "Proceeding with segment {} without output segment",
             sseg.getId());
       }
 
       // sseg.toString() => sseg.text.toText(), codes included
       SolrInputDocument doc = getSolrDocument(solrClient, coreName, sseg.toString());
 
       try {
         LocaleId effectiveOutputLang = (trgLoc != null) ? trgLoc : getTargetLocale();
         populateSolrDocument(doc, tu, sseg, tseg, getSourceLocale(), effectiveOutputLang);
 
       } catch (Exception e) {
         Log.warn(SolrIndexWriterStep.class, "Unable to populate index document for segment {}",
             sseg.getId(), e);
 
         continue;
       }
 
       docs.add(doc);
     }
   }
 
   public SolrClient getSolrClient() {
     return solrClient;
   }
 
   public String getCoreName() {
     return coreName;
   }
 
   public boolean isAutoCommit() {
     return autoCommit;
   }
 
   public int getTuBatchSize() {
     return tuBatchSize;
   }
 
   public List<ITextUnit> getTuBatch() {
     return tuBatch;
   }
 
   public boolean isRequireTarget() {
     return requireTarget;
   }
 
 }