package com.acumenvelocity.ath.steps;
   
   import com.acumenvelocity.ath.common.OkapiUtil;
   import com.acumenvelocity.ath.common.OriginalTuAnnotation;
   
   import net.sf.okapi.common.Event;
   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.steps.leveraging.LeveragingStep;
  
  /**
   * Skips machine translation for text units where all target segments already have non-empty
   * content.
   */
  public class MtLeveragingStep extends LeveragingStep {
  
    private boolean mtSendPlainText;
  
    public MtLeveragingStep(boolean mtSendPlainText) {
      super();
      this.mtSendPlainText = mtSendPlainText;
    }
  
    @Override
    protected Event handleTextUnit(Event event) {
      // Get the text unit from the event
      ITextUnit textUnit = event.getTextUnit();
  
      if (textUnit == null) {
        return event; // Return unchanged if no text unit
      }
  
      if (mtSendPlainText) {
        // Store the original TU with codes, remove codes to improve MT quality
        // CodesReinsertionStep will use OriginalTuAnnotation to get the original source codes
        textUnit.setAnnotation(new OriginalTuAnnotation(textUnit.clone(), getSourceLocale()));
        OkapiUtil.removeCodes(textUnit, true);
      }
  
      // Get the target container for the specified locale
      TextContainer target = textUnit.getTarget(getTargetLocale());
  
      if (target == null || target.getSegments() == null) {
        return super.handleTextUnit(event); // Proceed with MT if no target or segments
      }
  
      // Get source container to align segments
      TextContainer source = textUnit.getSource();
  
      if (source == null || source.getSegments() == null) {
        return super.handleTextUnit(event); // Proceed with MT if no source or segments
      }
  
      // Check if all target segments have non-empty content
      boolean allSegmentsTranslated = true;
  
      for (Segment sourceSegment : source.getSegments()) {
        if (sourceSegment == null) {
          continue; // Skip null source segments
        }
  
        Segment targetSegment = target.getSegments().get(sourceSegment.getId());
  
        if (targetSegment == null || targetSegment.getContent() == null ||
            targetSegment.getContent().isEmpty()) {
  
          allSegmentsTranslated = false;
          break; // Exit loop if any segment is untranslated
        }
      }
  
      // If all segments are translated, return the event unchanged to skip MT
      if (allSegmentsTranslated) {
        return event;
      }
  
      // Otherwise, proceed with machine translation in the superclass
      return super.handleTextUnit(event);
    }
  }