package com.acumenvelocity.ath.steps;
   
   import java.util.ArrayList;
   import java.util.List;
   
   import com.acumenvelocity.ath.common.Const;
   import com.acumenvelocity.ath.common.ConversionUtil;
   import com.acumenvelocity.ath.common.Log;
   import com.acumenvelocity.ath.common.OkapiUtil;
  import com.acumenvelocity.ath.common.OkapiWordBreaker;
  import com.acumenvelocity.ath.common.OriginalTuAnnotation;
  import com.acumenvelocity.ath.gemini.GenAi;
  import com.acumenvelocity.ath.model.InlineCodeRef;
  import com.acumenvelocity.ath.model.LayeredSegment;
  import com.acumenvelocity.ath.model.x.LayeredTextX;
  
  import net.sf.okapi.common.Event;
  import net.sf.okapi.common.IResource;
  import net.sf.okapi.common.Util;
  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;
  import net.sf.okapi.common.resource.TextFragmentUtil;
  
  public class CodesReinsertionStep extends BaseTuBatchProcessingStep {
  
    private final List<LayeredSegment> lsegs = new ArrayList<>();
    private final List<LayeredTextX> slts = new ArrayList<>();
    private boolean useCodesReinsertionModel;
    private String codesReinsertionModelName;
  
    public CodesReinsertionStep(boolean useCodesReinsertionModel, String codesReinsertionModelName) {
      super();
      this.useCodesReinsertionModel = useCodesReinsertionModel;
  
      this.codesReinsertionModelName = Util.isEmpty(codesReinsertionModelName)
          ? Const.GEMINI_CODE_REINSERTION_MODEL
          : codesReinsertionModelName;
    }
  
    @Override
    public String getName() {
      return "Code Reinsertion Step";
    }
  
    @Override
    public String getDescription() {
      return "Restores the source part of TextUnit resources from an annotation, employs a LLM to "
          + "insert inline codes in the targets.";
    }
  
    @Override
    protected void clear() {
      lsegs.clear();
      slts.clear();
    }
  
    private void preProcessTextUnit(ITextUnit tu) {
      OriginalTuAnnotation ota = tu.getAnnotation(OriginalTuAnnotation.class);
  
      // OriginalTuAnnotation is set by MtLeveragingStep that removes source codes for MT, for
      // alignment we use the original source codes
      if (ota != null) {
        ITextUnit otu = ota.getTu();
        tu.setSource(otu.getSource());
      }
  
      TextContainer src = tu.getSource();
      TextContainer trg = tu.getTarget(getTargetLocale());
  
      if (src == null) {
        Log.error(this.getClass(), "Source of '{}' is null", tu.getId());
        return;
      }
  
      // Target can be null for not-yet-translated segments, warn and continue
      if (trg == null) {
        Log.warn(this.getClass(), "Target ''{}'' of '{}' is null, creating an empty target",
            getTargetLocale(), tu.getId());
  
        trg = tu.createTarget(getTargetLocale(), false, IResource.COPY_SEGMENTATION);
      }
  
      if (!useCodesReinsertionModel) {
        return;
      }
  
      for (Segment sseg : src.getSegments()) {
        // Don't include in analysis an empty list of source codes
        if (Util.isEmpty(sseg.getContent().getCodes())) {
          continue;
        }
  
        Segment tseg = trg.getSegments().get(sseg.getId());
  
        LayeredTextX slt = ConversionUtil.toLayeredText(sseg.text, getSourceLocale());
        LayeredTextX tlt = ConversionUtil.toLayeredText(tseg.text, getTargetLocale());
  
       LayeredSegment lseg = ConversionUtil.toLayeredSegment(slt, tlt);
 
       lseg.setTrgWordBreakPositions(
           OkapiWordBreaker.getWordBreakPositions(tlt.getText(), getTargetLocale()));
 
       lsegs.add(lseg);
       slts.add(slt);
     }
   }
 
   /**
    * Update the target codes from the LLM API response.
    * 
    * @param trgCodesList
    */
   private void postProcessTextUnits(List<Event> tuEvents, List<List<InlineCodeRef>> trgCodesList) {
     int index = 0;
 
     for (Event tue : tuEvents) {
       ITextUnit tu = tue.getTextUnit();
 
       TextContainer src = tu.getSource();
       TextContainer trg = tu.getTarget(getTargetLocale());
 
       if (src == null) {
         Log.error(this.getClass(), "Source of '{}' is null", tu.getId());
         continue;
       }
 
       if (useCodesReinsertionModel) {
         for (Segment sseg : src.getSegments()) {
           Log.trace(this.getClass(), "sseg: '{}'", sseg);
 
           // Don't include in analysis an empty list of source codes
           if (Util.isEmpty(sseg.getContent().getCodes())) {
             continue;
           }
 
           Segment tseg = trg.getSegments().get(sseg.getId());
           Log.trace(this.getClass(), "tseg: '{}'", tseg);
 
           // Create an empty target segment with the same Id as of the source segment
           if (tseg == null) {
             tseg = new Segment(sseg.getId());
             trg.append(tseg);
             Log.trace(this.getClass(), "Created a missing tseg: '{}'", tseg);
           }
 
           if (index < trgCodesList.size()) {
             List<InlineCodeRef> trgCodes = trgCodesList.get(index);
             LayeredSegment lseg = lsegs.get(index);
             lseg.setTrgCodes(trgCodes);
             LayeredTextX slt = slts.get(index);
             LayeredTextX tlt = ConversionUtil.tltFromLayeredSegment(lseg, slt);
 
             try {
               TextFragment ttf = ConversionUtil.toTextFragment(tlt);
               tseg.setContent(ttf);
 
             } catch (Exception e) {
               Log.warn(getClass(), "Error converting to TextFragment: {}", e.getMessage());
               continue;
             }
 
             index++;
           }
         }
       }
 
       // Fix codes after AI if needed.
       // If AI is not used, this code reinserts target codes.
       for (Segment sseg : src.getSegments()) {
         Segment tseg = trg.getSegments().get(sseg.getId());
         
         if (tseg == null) {
           continue;
         }
 
         // For empty targets we copy source codes
         TextFragment segSource = sseg.getContent();
         TextFragment segTarget = tseg.getContent();
 
         // This normalization is needed even after LLM alignment that can hypothetically create
         // extra codes
         OkapiUtil.removeExtraCodes(segSource.getCodes(), segTarget);
 
         // Align codes and copy metadata from source to target
         TextFragmentUtil.alignAndCopyCodeMetadata(segSource, segTarget, true, true);
         
         // Rearrange opening and closing codes
         OkapiUtil.rearrangeCodes(segSource.getCodes(), segTarget);
       }
     }
   }
 
   @Override
   protected void processTuEvents(List<Event> tuEvents) {
     // Process TUs
     for (Event tue : tuEvents) {
       ITextUnit tu = tue.getTextUnit();
       preProcessTextUnit(tu);
     }
 
     List<List<InlineCodeRef>> trgCodesList = useCodesReinsertionModel
         ? GenAi.reinsertCodes(codesReinsertionModelName, lsegs)
         : null;
 
     postProcessTextUnits(tuEvents, trgCodesList);
   }
 }