package com.acumenvelocity.ath.steps;
   
   import java.util.ArrayList;
   import java.util.List;
   
   import com.acumenvelocity.ath.common.Const;
   import com.acumenvelocity.ath.common.ControllerUtil;
   import com.acumenvelocity.ath.common.Log;
   import com.acumenvelocity.ath.common.OkapiUtil;
  import com.acumenvelocity.ath.gct.v3.AthTranslation;
  
  import net.sf.okapi.common.Event;
  import net.sf.okapi.common.IParameters;
  import net.sf.okapi.common.IResource;
  import net.sf.okapi.common.MimeTypeMapper;
  import net.sf.okapi.common.Util;
  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;
  import net.sf.okapi.common.resource.TextContainer;
  import net.sf.okapi.common.resource.TextFragment;
  import net.sf.okapi.common.resource.TextFragmentUtil;
  import net.sf.okapi.lib.translation.QueryUtil;
  
  public class BatchMtStep extends BaseTuBatchProcessingStep {
  
    private BatchMtParameters params = new BatchMtParameters();
    private final List<String> sourceSegments = new ArrayList<>();
    private final List<SegmentInfo> segmentInfos = new ArrayList<>();
    private final QueryUtil qutil = new QueryUtil();
  
    @Override
    public String getName() {
      return "Batch Google Cloud Translation v3";
    }
  
    @Override
    public String getDescription() {
      return "Translates a batch of collected TU segments with the "
          + "Google Cloud Translate v3 service";
    }
  
    @Override
    public BatchMtParameters getParameters() {
      return params;
    }
  
    @Override
    public void setParameters(IParameters params) {
      this.params = (BatchMtParameters) params;
    }
  
    @Override
    protected void clear() {
      sourceSegments.clear();
      segmentInfos.clear();
    }
  
    /**
     * Pre-process text units: collect source segments
     */
    private void preProcessTextUnit(ITextUnit tu) {
      TextContainer source = tu.getSource();
  
      if (source == null) {
        Log.error(getClass(), "Source of TU '{}' is null", tu.getId());
        return;
      }
  
      // Process each segment in the text unit
      for (Segment segment : source.getSegments()) {
        TextFragment content = segment.getContent();
  
        // Skip empty segments
        if (content == null || content.isEmpty()) {
          Log.trace(getClass(), "Skipping empty segment in TU '{}'", tu.getId());
          continue;
        }
  
        String sourceText = null;
  
        if (params.isMtSendPlainText()) {
          sourceText = content.getText();
          params.setMimeType(MimeTypeMapper.PLAIN_TEXT_MIME_TYPE);
  
        } else {
          sourceText = qutil.toCodedHTML(content);
          params.setMimeType(MimeTypeMapper.HTML_MIME_TYPE);
        }
  
        sourceSegments.add(sourceText);
  
        // Store segment info for later mapping
        segmentInfos.add(new SegmentInfo(tu, segment.getId()));
  
        Log.trace(getClass(), "Collected segment [{}]: '{}'",
            sourceSegments.size() - 1, sourceText);
      }
   }
 
   /**
    * Post-process text units: set target segments
    */
   private void postProcessTextUnits(List<String> targetSegments) {
     if (Util.isEmpty(targetSegments)) {
       Log.warn(getClass(), "No evaluation results available");
       return;
     }
 
     if (targetSegments.size() != segmentInfos.size()) {
       Log.error(getClass(), "Mismatch: {} target segments, but {} segment infos",
           targetSegments.size(), segmentInfos.size());
 
       return;
     }
 
     // Process each target segment
     for (int i = 0; i < targetSegments.size(); i++) {
       String targetText = targetSegments.get(i);
       SegmentInfo segInfo = segmentInfos.get(i);
 
       ITextUnit tu = segInfo.textUnit;
       TextContainer source = tu.getSource();
       String segmentId = segInfo.segmentId;
 
       // Get or create target container
       TextContainer target = tu.getTarget(getTargetLocale());
 
       if (target == null) {
         target = tu.createTarget(getTargetLocale(), false, IResource.COPY_SEGMENTATION);
         Log.trace(getClass(), "Created target container for TU '{}'", tu.getId());
       }
 
       // Get or create target segment
       Segment sseg = source.getSegments().get(segmentId);
       Segment tseg = target.getSegments().get(segmentId);
 
       if (tseg == null) {
         tseg = new Segment(segmentId);
         target.append(tseg);
 
         Log.trace(getClass(), "Created target segment '{}' in TU '{}'", segmentId, tu.getId());
       }
 
       TextFragment targetTf;
 
       if (params.isMtSendPlainText()) {
         targetTf = new TextFragment(targetText);
 
       } else {
         targetTf = qutil.fromCodedHTMLToFragment(targetText, null);
         TextFragment segSource = sseg.getContent();
 
         OkapiUtil.removeExtraCodes(segSource.getCodes(), targetTf);
 
         // Align codes and copy metadata from source to target
         TextFragmentUtil.alignAndCopyCodeMetadata(segSource, targetTf, true, true);
 
         // Rearrange opening and closing codes
         OkapiUtil.rearrangeCodes(segSource.getCodes(), targetTf);
       }
 
       tseg.setContent(targetTf);
 
       // Mark as MT
       AltTranslationsAnnotation ata = new AltTranslationsAnnotation();
 
       ata.add(getSourceLocale(), getTargetLocale(), sseg.getContent(), sseg.getContent(), targetTf,
           MatchType.MT, 95, Const.MT_PROVIDER_GOOGLE_MT_V3, 95, 95);
 
       tseg.setAnnotation(ata);
     }
   }
 
   @Override
   protected void processTuEvents(List<Event> tuEvents) {
     for (Event tue : tuEvents) {
       ITextUnit tu = tue.getTextUnit();
       preProcessTextUnit(tu);
     }
 
     Log.info(getClass(), "Collected {} source segments from {} text units",
         sourceSegments.size(), tuEvents.size());
 
     List<String> targetSegments = AthTranslation.translateBatch(
         sourceSegments,
         getSourceLocale().toString(),
         getTargetLocale().toString(),
         MimeTypeMapper.PLAIN_TEXT_MIME_TYPE,
         ControllerUtil.getProjectId(),
         params.getProjectLocation(),        
         params.getModelProjectId(),
         params.getModelProjectLocation(),
         params.getModelId(),        
         params.getGlossaryProjectId(),
         params.getGlossaryProjectLocation(),
         params.getGlossaryId());
 
     // translateBatchWithModel(
     // AthTranslation.getClient(),
     // sourceSegments,
     // getSourceLocale().toString(),
     // getTargetLocale().toString());
 
     postProcessTextUnits(targetSegments);
   }
 
 }