package com.acumenvelocity.ath.solr.doc;
   
   import java.util.Comparator;
   import java.util.HashMap;
   import java.util.List;
   import java.util.Map;
   
   import org.apache.solr.client.solrj.response.QueryResponse;
   import org.apache.solr.common.SolrDocument;
  import org.apache.solr.common.SolrDocumentList;
  
  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.exception.AthRuntimeException;
  import com.acumenvelocity.ath.model.x.LayeredTextX;
  import com.acumenvelocity.ath.solr.AthIndex;
  import com.acumenvelocity.ath.solr.Solr;
  import com.acumenvelocity.ath.steps.BaseMergerStep;
  
  import net.sf.okapi.common.LocaleId;
  import net.sf.okapi.common.resource.ITextUnit;
  import net.sf.okapi.common.resource.Segment;
  import net.sf.okapi.common.resource.TextContainer;
  import net.sf.okapi.common.resource.TextFragment;
  
  public class SolrDocMergerStep extends BaseMergerStep {
  
    private String docId;
  
    public SolrDocMergerStep(String docId) {
      super();
      this.docId = docId;
    }
  
    @Override
    public String getName() {
      return "Solr Document Merger Step";
    }
  
    @Override
    public String getDescription() {
      return "Merges Text Unit resources with translated segments from Solr";
    }
  
    @Override
    protected void processTu(ITextUnit tu) {
      try {
        LocaleId targetLocale = super.getTargetLocale();
  
        // Check if target exists for this locale, if not create empty target
        if (!tu.hasTarget(targetLocale)) {
          // Create empty target without copying source segmentation
          tu.createTarget(targetLocale, false, ITextUnit.CREATE_EMPTY);
        }
  
        String query = Log.format("docId:\"{}\" AND tuId:\"{}\"",
            docId,
            tu.getId());
  
        // Set up query parameters
        Map<String, Object> queryParams = new HashMap<>();
  
        queryParams.put(Solr.FIELDS, new String[] {
            Const.ATH_PROP_TARGET_JSON,
            Const.ATH_PROP_POSITION // Included for sorting
        });
  
        // Ensure we get all segments for this text unit
        queryParams.put(Solr.ROWS, Integer.MAX_VALUE);
  
        // Query the ath_doc_segments Solr core
        QueryResponse response = AthIndex.getMany(
            Const.SOLR_CORE_ATH_DOC_SEGMENTS,
            query,
            queryParams,
            QueryResponse.class);
  
        SolrDocumentList docList = response.getResults();
  
        if (docList.getNumFound() > 0) {
          // Sort documents by position field
          List<SolrDocument> sortedDocs = new java.util.ArrayList<>(docList);
  
          sortedDocs.sort(Comparator.comparingLong(doc -> {
            Object posObj = doc.get(Const.ATH_PROP_POSITION);
  
            if (posObj instanceof Long) {
              return (Long) posObj;
  
            } else if (posObj instanceof Integer) {
              return ((Integer) posObj).longValue();
            }
  
            return 0L;
          }));
  
          // Always create a new empty target
         TextContainer target = tu.setTarget(getTargetLocale(), new TextContainer());
 
         // Process all segments
         for (int segmentIndex = 0; segmentIndex < sortedDocs.size(); segmentIndex++) {
           SolrDocument doc = sortedDocs.get(segmentIndex);
 
           // Get targetJson field
           String targetJson = doc._getStr(Const.ATH_PROP_TARGET_JSON, null);
 
           if (targetJson != null) {
             // Deserialize JSON to LayeredText
             LayeredTextX targetLt = JacksonUtil.fromJson(targetJson, LayeredTextX.class);
 
             if (targetLt != null) {
               // Convert LayeredText to TextFragment
               TextFragment targetTf = ConversionUtil.toTextFragment(targetLt);
 
               // Create a new segment with the loop index as ID
               Segment targetSegment = new Segment(String.valueOf(segmentIndex), targetTf);
 
               // Append the segment to the target container
               target.getSegments().append(targetSegment);
 
               Log.trace(this.getClass(),
                   "Loaded target segment for docId={}, segmentIndex={}: {}",
                   docId, segmentIndex, targetTf.getText());
 
             } else {
               Log.warn(this.getClass(),
                   "Failed to deserialize targetJson for docId={}, segmentIndex={}",
                   docId, segmentIndex);
             }
 
           } else {
             Log.warn(this.getClass(),
                 "No targetJson found for docId={}, segmentIndex={}",
                 docId, segmentIndex);
           }
         }
 
       } else {
         Log.warn(this.getClass(),
             "No document found in ath_doc for docId={}, tuId={}",
             docId, tu.getId());
       }
 
     } catch (Exception e) {
       AthRuntimeException.logAndThrow(this.getClass(), e);
     }
   }
 
   @Override
   protected void setTargetSegment(ITextUnit tu, Segment sseg, Segment tseg, long position) {
     // Not used, the TU is processed the whole
   }
 
 }