package net.sf.okapi.steps.llmsentencealigner;
   
   import java.util.ArrayList;
   import java.util.Iterator;
   import java.util.List;
   
   import com.acumenvelocity.ath.common.AlignmentData.CombinedAlignment;
   import com.acumenvelocity.ath.common.AlignmentData.CombinedAlignmentInput;
   import com.acumenvelocity.ath.common.AlignmentData.CombinedAlignmentOutput;
  import com.acumenvelocity.ath.common.AlignmentData.ParagraphAlignment;
  import com.acumenvelocity.ath.common.AlignmentData.ParagraphWithSegments;
  import com.acumenvelocity.ath.common.AlignmentData.SegmentInfo;
  import com.acumenvelocity.ath.common.AlignmentData.SentenceAlignment;
  import com.acumenvelocity.ath.common.ConversionUtil;
  import com.acumenvelocity.ath.common.Log;
  import com.acumenvelocity.ath.common.OkapiUtil;
  import com.acumenvelocity.ath.gemini.GenAi;
  import com.acumenvelocity.ath.model.InlineCode;
  import com.acumenvelocity.ath.model.x.LayeredTextX;
  import com.acumenvelocity.ath.steps.BaseAlignerStep;
  
  import net.sf.okapi.common.IParameters;
  import net.sf.okapi.common.IResource;
  import net.sf.okapi.common.UsingParameters;
  import net.sf.okapi.common.Util;
  import net.sf.okapi.common.exceptions.OkapiException;
  import net.sf.okapi.common.filters.IFilter;
  import net.sf.okapi.common.resource.AlignmentStatus;
  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.common.resource.TextUnitUtil;
  
  @UsingParameters(LlmSentenceAlignerParameters.class)
  public class LlmSentenceAlignerStep extends BaseAlignerStep {
  
    private LlmSentenceAlignerParameters params;
  
    private final List<List<SegmentDataWithCodes>> sourceSegmentData = new ArrayList<>();
    private final List<List<SegmentDataWithCodes>> targetSegmentData = new ArrayList<>();
  
    public LlmSentenceAlignerStep(IFilter targetFilter) {
      super(targetFilter);
      params = new LlmSentenceAlignerParameters();
    }
  
    private static class SegmentDataWithCodes {
      String text;
      List<InlineCode> codes = new ArrayList<>();
    }
  
    @Override
    public String getName() {
      return "LLM Sentence Alignment";
    }
  
    @Override
    public String getDescription() {
      return "Aligns paragraphs and sentences using LLM. Handles crossed paragraphs and different document structures.";
    }
  
    @Override
    public LlmSentenceAlignerParameters getParameters() {
      return params;
    }
  
    @Override
    public void setParameters(IParameters params) {
      this.params = (LlmSentenceAlignerParameters) params;
    }
  
    @Override
    protected boolean isSegmentSource() {
      return params.isSegmentSource();
    }
  
    @Override
    protected boolean isSegmentTarget() {
      return params.isSegmentTarget();
    }
  
    @Override
    protected boolean isUseCustomSourceRules() {
      return params.isUseCustomSourceRules();
    }
  
    @Override
    protected boolean isUseCustomTargetRules() {
      return params.isUseCustomTargetRules();
    }
  
    @Override
    protected String getCustomSourceRulesPath() {
      return params.getCustomSourceRulesPath();
    }
  
    @Override
   protected String getCustomTargetRulesPath() {
     return params.getCustomTargetRulesPath();
   }
 
   @Override
   protected boolean isCollapseWhitespace() {
     return params.isCollapseWhitespace();
   }
 
   @Override
   protected void clear() {
     super.clear();
     sourceSegmentData.clear();
     targetSegmentData.clear();
   }
 
   @Override
   protected void performAlignment(List<ITextUnit> sourceTUs, List<ITextUnit> targetTUs) {
     // Single LLM call for both paragraph and sentence alignment
     CombinedAlignmentOutput alignmentOutput = performCombinedAlignment(sourceTUs, targetTUs);
 
     // Apply alignments
     applyCombinedAlignments(alignmentOutput, sourceTUs, targetTUs);
 
     // Set alignment origin metadata
     for (ITextUnit tu : sourceTUs) {
       OkapiUtil.setAlOrigin(tu, getSourceLocale(), getTargetLocale());
     }
   }
 
   private CombinedAlignmentOutput performCombinedAlignment(List<ITextUnit> sourceTUs,
       List<ITextUnit> targetTUs) {
     CombinedAlignmentInput input = new CombinedAlignmentInput();
     input.sourceLanguage = getSourceLocale().toString();
     input.targetLanguage = getTargetLocale().toString();
     input.task = "Align paragraphs first, then align sentences within each paragraph pair";
 
     // Build source paragraphs with segments (codes removed)
     for (ITextUnit srcTu : sourceTUs) {
       ParagraphWithSegments pws = new ParagraphWithSegments();
       pws.position = input.sourceParagraphs.size();
       pws.id = srcTu.getId();
       pws.context = srcTu.getName();
 
       List<SegmentDataWithCodes> segDataList = new ArrayList<>();
 
       for (Segment seg : srcTu.getSource().getSegments()) {
         LayeredTextX slt = ConversionUtil.toLayeredText(seg.text, getSourceLocale());
 
         SegmentInfo si = new SegmentInfo();
         si.position = pws.segments.size();
         si.text = slt.getText();
         pws.segments.add(si);
 
         SegmentDataWithCodes sdc = new SegmentDataWithCodes();
         sdc.text = slt.getText();
         sdc.codes = slt.getCodes();
         segDataList.add(sdc);
       }
 
       input.sourceParagraphs.add(pws);
       sourceSegmentData.add(segDataList);
     }
 
     // Build target paragraphs with segments (codes removed)
     for (ITextUnit trgTu : targetTUs) {
       ParagraphWithSegments pws = new ParagraphWithSegments();
       pws.position = input.targetParagraphs.size();
       pws.id = trgTu.getId();
       pws.context = trgTu.getName();
 
       List<SegmentDataWithCodes> segDataList = new ArrayList<>();
 
       for (Segment seg : trgTu.getSource().getSegments()) {
         LayeredTextX tlt = ConversionUtil.toLayeredText(seg.text, getTargetLocale());
 
         SegmentInfo si = new SegmentInfo();
         si.position = pws.segments.size();
         si.text = tlt.getText();
         pws.segments.add(si);
 
         SegmentDataWithCodes sdc = new SegmentDataWithCodes();
         sdc.text = tlt.getText();
         sdc.codes = tlt.getCodes();
         segDataList.add(sdc);
       }
 
       input.targetParagraphs.add(pws);
       targetSegmentData.add(segDataList);
     }
 
     // LLM-based alignment
     try {
       return GenAi.alignParagraphsAndSentences(params.getAlignmentModelName(), input);
 
     } catch (Exception e) {
       Log.error(getClass(), "LLM alignment failed: {}", e.getMessage(), e);
       throw new OkapiException("LLM alignment failed", e);
     }
   }
 
   private void applyCombinedAlignments(CombinedAlignmentOutput output, List<ITextUnit> sourceTUs,
       List<ITextUnit> targetTUs) {
 
     for (CombinedAlignment combined : output.alignments) {
       ParagraphAlignment paraAlign = combined.paragraphAlignment;
 
       if (params.isLogAlignmentDetails()) {
         Log.debug(getClass(), "Paragraph alignment: type={}, src={}, trg={}",
             paraAlign.type, paraAlign.sourceParagraphPositions, paraAlign.targetParagraphPositions);
       }
 
       if (Util.isEmpty(paraAlign.sourceParagraphPositions)
           || Util.isEmpty(paraAlign.targetParagraphPositions)) {
 
         Log.warn(getClass(), "Problematic para alignment -- source: {}, target: {}",
             paraAlign.sourceParagraphPositions, paraAlign.targetParagraphPositions);
 
         continue;
       }
 
       // Collect source and target TUs for this paragraph pair
       List<ITextUnit> srcTuGroup = new ArrayList<>();
       List<List<SegmentDataWithCodes>> srcSegDataGroup = new ArrayList<>();
 
       for (int srcParaIndex : paraAlign.sourceParagraphPositions) {
         srcTuGroup.add(sourceTUs.get(srcParaIndex));
         srcSegDataGroup.add(sourceSegmentData.get(srcParaIndex));
       }
 
       List<ITextUnit> trgTuGroup = new ArrayList<>();
       List<List<SegmentDataWithCodes>> trgSegDataGroup = new ArrayList<>();
 
       for (int trgParaIndex : paraAlign.targetParagraphPositions) {
         trgTuGroup.add(targetTUs.get(trgParaIndex));
         trgSegDataGroup.add(targetSegmentData.get(trgParaIndex));
       }
 
       applySentenceAlignmentsToParagraphPair(
           srcTuGroup, trgTuGroup, srcSegDataGroup, trgSegDataGroup,
           combined.sentenceAlignments);
     }
   }
 
   private void applySentenceAlignmentsToParagraphPair(
       List<ITextUnit> srcTuGroup, List<ITextUnit> trgTuGroup,
       List<List<SegmentDataWithCodes>> srcSegDataGroup,
       List<List<SegmentDataWithCodes>> trgSegDataGroup,
       List<SentenceAlignment> sentenceAlignments) {
 
     if (Util.isEmpty(srcTuGroup) || Util.isEmpty(trgTuGroup)) {
       Log.warn(getClass(), "Problematic para alignment -- source: {}, target: {}",
           srcTuGroup, trgTuGroup);
 
       return;
     }
 
     if (Util.isEmpty(srcSegDataGroup) || Util.isEmpty(trgSegDataGroup)) {
       Log.warn(getClass(), "Problematic sentence alignment -- source: {}, target: {}",
           srcSegDataGroup, trgSegDataGroup);
 
       return;
     }
 
     // Case 1: 1:1 paragraph match
     if (srcTuGroup.size() == 1 && trgTuGroup.size() == 1) {
       applySentenceAlignmentsToSinglePair(
           srcTuGroup.get(0), trgTuGroup.get(0),
           srcSegDataGroup.get(0), trgSegDataGroup.get(0),
           sentenceAlignments);
 
       return;
     }
 
     // Case 2: Multi-paragraph match - merge source TUs
     ITextUnit primarySrcTu = srcTuGroup.get(0);
     List<SegmentDataWithCodes> mergedSrcSegData = new ArrayList<>(srcSegDataGroup.get(0));
 
     for (int i = 1; i < srcTuGroup.size(); i++) {
       ITextUnit additionalTu = srcTuGroup.get(i);
 
       for (Segment seg : additionalTu.getSource().getSegments()) {
         primarySrcTu.getSource().append(seg.clone());
       }
 
       mergedSrcSegData.addAll(srcSegDataGroup.get(i));
     }
 
     // Merge target TUs if multiple
     ITextUnit primaryTrgTu = trgTuGroup.isEmpty() ? null : trgTuGroup.get(0);
 
     List<SegmentDataWithCodes> mergedTrgSegData = trgTuGroup.isEmpty()
         ? new ArrayList<>()
         : new ArrayList<>(trgSegDataGroup.get(0));
 
     if (primaryTrgTu != null && trgTuGroup.size() > 1) {
       for (int i = 1; i < trgTuGroup.size(); i++) {
         ITextUnit additionalTu = trgTuGroup.get(i);
 
         for (Segment seg : additionalTu.getSource().getSegments()) {
           primaryTrgTu.getSource().append(seg.clone());
         }
 
         mergedTrgSegData.addAll(trgSegDataGroup.get(i));
       }
     }
 
     if (primaryTrgTu != null) {
       applySentenceAlignmentsToSinglePair(
           primarySrcTu, primaryTrgTu, mergedSrcSegData, mergedTrgSegData,
           sentenceAlignments);
     }
   }
 
   private void applySentenceAlignmentsToSinglePair(
       ITextUnit sourceTu, ITextUnit targetTu,
       List<SegmentDataWithCodes> srcSegData, List<SegmentDataWithCodes> trgSegData,
       List<SentenceAlignment> sentenceAlignments) {
 
     TextContainer srcCont = sourceTu.getSource();
     TextContainer trgCont = sourceTu.createTarget(getTargetLocale(), false, IResource.CREATE_EMPTY);
     trgCont.clear();
 
     List<Segment> srcSegments = new ArrayList<>(srcCont.getSegments().asList());
     int nextSegmentId = 1;
     int currentSrcPos = 0;
 
     for (SentenceAlignment align : sentenceAlignments) {
       if (Util.isEmpty(align.sourcePositions) || Util.isEmpty(align.targetPositions)) {
 
         Log.warn(getClass(), "Problematic sentence alignment -- source: {}, target: {}",
             align.sourcePositions, align.targetPositions);
 
         continue;
       }
 
       if (params.isLogAlignmentDetails()) {
         Log.debug(getClass(), "Sentence alignment: type={}, src={}, trg={}",
             align.type, align.sourcePositions, align.targetPositions);
       }
 
       switch (align.type) {
       case "MATCH":
         int srcPos = align.sourcePositions.get(0);
         int trgPos = align.targetPositions.get(0);
 
         String segId = srcSegments.get(srcPos).getId();
         SegmentDataWithCodes trgData = trgSegData.get(trgPos);
 
         // Restore codes to target text
         TextFragment trgFrag = ConversionUtil.toTextFragment(
             new LayeredTextX().text(trgData.text).codes(trgData.codes)
                 .language(getTargetLocale().toString()));
 
         trgCont.append(new Segment(segId, trgFrag));
         currentSrcPos = srcPos + 1;
         break;
 
       case "MULTI_MATCH":
         int firstSrcPos = align.sourcePositions.get(0);
         String groupId = srcSegments.get(firstSrcPos).getId();
 
         // Merge source segments
         for (int i = 0; i < align.sourcePositions.size() - 1; i++) {
           srcCont.getSegments().joinWithNext(firstSrcPos);
         }
 
         // Merge target segments with codes
         List<SegmentDataWithCodes> trgGroup = new ArrayList<>();
 
         for (int tPos : align.targetPositions) {
           trgGroup.add(trgSegData.get(tPos));
         }
 
         TextFragment mergedFrag = mergeSegmentDataWithCodes(trgGroup);
 
         trgCont.append(new Segment(groupId, mergedFrag));
         currentSrcPos = firstSrcPos + 1;
         break;
 
       case "DELETED":
         int delSrcPos = align.sourcePositions.get(0);
         String delSrcId = srcSegments.get(delSrcPos).getId();
 
         trgCont.append(new Segment(delSrcId, new TextFragment("")));
         currentSrcPos = delSrcPos + 1;
         break;
 
       case "INSERTED":
         int insTrgPos = align.targetPositions.get(0);
         SegmentDataWithCodes insTrgData = trgSegData.get(insTrgPos);
 
         String newId = generateUniqueSegmentId(sourceTu, nextSegmentId++);
 
         TextFragment emptySrcFrag = new TextFragment("");
         srcCont.getSegments().insert(currentSrcPos, new Segment(newId, emptySrcFrag));
 
         TextFragment insTrgFrag = ConversionUtil.toTextFragment(
             new LayeredTextX().text(insTrgData.text).codes(insTrgData.codes)
                 .language(getTargetLocale().toString()));
 
         trgCont.append(new Segment(newId, insTrgFrag));
         currentSrcPos++;
         break;
       }
     }
 
     // Verify segment counts match
     int srcCount = srcCont.getSegments().count();
     int trgCount = trgCont.getSegments().count();
 
     if (srcCount != trgCount) {
       throw new OkapiException(
           String.format("Segment count mismatch in TU '%s': source=%d, target=%d",
               sourceTu.getId(), srcCount, trgCount));
     }
 
     // Verify all segment IDs match
     Iterator<Segment> srcIt = srcCont.getSegments().iterator();
     Iterator<Segment> trgIt = trgCont.getSegments().iterator();
 
     while (srcIt.hasNext() && trgIt.hasNext()) {
       Segment srcSeg = srcIt.next();
       Segment trgSeg = trgIt.next();
 
       if (!srcSeg.getId().equals(trgSeg.getId())) {
         throw new OkapiException(
             String.format("Segment ID mismatch in TU '%s': source='%s', target='%s'",
                 sourceTu.getId(), srcSeg.getId(), trgSeg.getId()));
       }
     }
 
     trgCont.setHasBeenSegmentedFlag(true);
     trgCont.getSegments().setAlignmentStatus(AlignmentStatus.ALIGNED);
 
     // Align and copy code metadata from source to target
     if (params.isLogAlignmentDetails()) {
       Log.debug(getClass(), "Aligning codes between source and target for TU: {}",
           sourceTu.getId());
     }
 
     srcIt = srcCont.getSegments().iterator();
     trgIt = trgCont.getSegments().iterator();
 
     while (srcIt.hasNext() && trgIt.hasNext()) {
       Segment srcSeg = srcIt.next();
       Segment trgSeg = trgIt.next();
 
       if (params.isUseCodesReinsertionModel()) {
         TextUnitUtil.removeCodes(trgSeg.getContent());
 
       } else {
         // LLM-based code re-insertion does this normalization, called only for no-LLM
         OkapiUtil.removeExtraCodes(srcSeg.getContent().getCodes(), trgSeg.getContent());
 
         // Align codes and copy metadata from source to target
         TextFragmentUtil.alignAndCopyCodeMetadata(srcSeg.text, trgSeg.text, true, true);
 
         // Rearrange opening and closing codes
         OkapiUtil.rearrangeCodes(srcSeg.getContent().getCodes(), trgSeg.getContent());
       }
     }
   }
 
   private TextFragment mergeSegmentDataWithCodes(List<SegmentDataWithCodes> segDataList) {
     if (segDataList.isEmpty()) {
       return new TextFragment("");
     }
 
     if (segDataList.size() == 1) {
       SegmentDataWithCodes data = segDataList.get(0);
 
       return ConversionUtil.toTextFragment(
           new LayeredTextX().text(data.text).codes(data.codes)
               .language(getTargetLocale().toString()));
     }
 
     // Merge texts and codes
     StringBuilder mergedText = new StringBuilder();
     List<InlineCode> mergedCodes = new ArrayList<>();
 
     int cumulativeOffset = 0;
 
     for (int i = 0; i < segDataList.size(); i++) {
       SegmentDataWithCodes data = segDataList.get(i);
 
       if (i > 0) {
         mergedText.append(" ");
         cumulativeOffset++;
       }
 
       mergedText.append(data.text);
 
       // Adjust code positions
       for (InlineCode code : data.codes) {
         InlineCode adjustedCode = new InlineCode();
         adjustedCode.setId(code.getId());
         adjustedCode.setPosition(code.getPosition() + cumulativeOffset);
         adjustedCode.setTagType(code.getTagType());
         adjustedCode.setType(code.getType());
         adjustedCode.setData(code.getData());
         adjustedCode.setOuterData(code.getOuterData());
         adjustedCode.setFlag(code.getFlag());
         adjustedCode.setDisplayText(code.getDisplayText());
         adjustedCode.setOriginalId(code.getOriginalId());
         mergedCodes.add(adjustedCode);
       }
 
       cumulativeOffset += data.text.length();
     }
 
     return ConversionUtil.toTextFragment(
         new LayeredTextX().text(mergedText.toString()).codes(mergedCodes)
             .language(getTargetLocale().toString()));
   }
 
   private String generateUniqueSegmentId(ITextUnit tu, int counter) {
     String candidateId;
 
     do {
       candidateId = tu.getId() + "_seg_" + counter++;
 
     } while (tu.getSource().getSegments().get(candidateId) != null);
 
     return candidateId;
   }
 }