package com.acumenvelocity.ath.steps;
   
   import java.util.ArrayList;
   import java.util.List;
   
   import com.acumenvelocity.ath.common.AthUtil;
   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.common.OriginalTuAnnotation;
  import com.acumenvelocity.ath.model.MtResources;
  import com.acumenvelocity.ath.mt.confidence.ConfidenceScoredTranslation;
  import com.acumenvelocity.ath.mt.confidence.HybridDataStructures.HybridEvaluationResult;
  import com.acumenvelocity.ath.mt.confidence.HybridDataStructures.HybridScoredSegment;
  import com.acumenvelocity.ath.mt.confidence.HybridDataStructures.HybridTranslationScore;
  import com.acumenvelocity.ath.mt.confidence.HybridQualityEstimator;
  
  import net.sf.okapi.common.Event;
  import net.sf.okapi.common.IResource;
  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;
  
  /**
   * MT Confidence Scoring Step for Okapi Framework v1.47.0
   * 
   * Translates source segments using multiple MT models (Google Cloud Translate v3 NMT,
   * Translation LLM, and custom AutoML models), evaluates translation quality using
   * Vertex AI (MetricX) and heuristic methods, then attaches confidence-scored
   * alternate translations to each segment.
   */
  public class MtConfidenceScoringStep extends BaseTuBatchProcessingStep {
  
    private final List<MtResources> mtCustomResources;
    private final List<MtResources> modelConfigs = new ArrayList<>();
    private final List<String> sourceTexts = new ArrayList<>();
    private final List<TextFragment> sourceTfs = new ArrayList<>();
    private final List<SegmentInfo> segmentInfos = new ArrayList<>();
    private final boolean mtSendPlainText;
    private final QueryUtil qutil = new QueryUtil();
  
    private HybridEvaluationResult results;
  
    /**
     * Creates a new MT Confidence Scoring Step.
     * 
     * @param mtCustomResources List of user's custom models and glossaries
     */
    public MtConfidenceScoringStep(List<MtResources> mtCustomResources, boolean mtSendPlainText) {
      super();
      this.mtCustomResources = mtCustomResources != null ? mtCustomResources : new ArrayList<>();
      this.mtSendPlainText = mtSendPlainText;
    }
  
    public String getName() {
      return "MT Confidence Scoring";
    }
  
    @Override
    public String getDescription() {
      return "Machine translates source segments using Google Cloud Translate API v3 "
          + "(NMT, Translation LLM) and custom AutoML models/glossaries, then uses "
          + "Google Vertex AI evaluation service and heuristics to calculate confidence "
          + "scores for those translations.";
    }
  
    /**
     * Initialize model configurations with built-in Google models + custom models
     */
    private void initializeModelConfigs() {
      modelConfigs.clear();
  
      // Add Google Cloud Translate v3 built-in models
  
      // 1. NMT (Neural Machine Translation) - standard model
      MtResources nmt = new MtResources();
  
      nmt.setMtModelId("general/nmt");
      nmt.setMtModelProjectId(ControllerUtil.getProjectId());
      nmt.setMtModelProjectLocation(Const.US_CENTRAL1_PROJECT_LOCATION);
  
      modelConfigs.add(nmt);
  
      // 2. Translation LLM - advanced model
      MtResources translationLlm = new MtResources();
  
      translationLlm.setMtModelId("general/translation-llm");
      translationLlm.setMtModelProjectId(ControllerUtil.getProjectId());
      translationLlm.setMtModelProjectLocation(Const.US_CENTRAL1_PROJECT_LOCATION);
  
      modelConfigs.add(translationLlm);
  
     // 3. Add user's custom models and glossaries
     if (!Util.isEmpty(mtCustomResources)) {
       modelConfigs.addAll(mtCustomResources);
     }
 
     Log.info(MtConfidenceScoringStep.class,
         "Initialized {} model configurations (2 built-in + {} custom)",
         modelConfigs.size(), mtCustomResources != null ? mtCustomResources.size() : 0);
   }
 
   /**
    * 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;
     }
     
     for (Segment segment : source.getSegments()) {
       TextFragment srcTf = segment.getContent();
 
       // Skip empty segments
       if (srcTf == null || srcTf.isEmpty()) {
         Log.trace(getClass(), "Skipping empty segment in TU '{}'", tu.getId());
         continue;
       }
       
       sourceTfs.add(srcTf.clone());
     }
     
     if (mtSendPlainText) {
       // Store the original TU with codes, remove codes to improve MT quality
       // CodesReinsertionStep will use OriginalTuAnnotation to get the original source codes
       tu.setAnnotation(new OriginalTuAnnotation(tu.clone(), getSourceLocale()));
       OkapiUtil.removeCodes(tu, true);
     }
 
     // Process each segment in the text unit
     for (Segment segment : source.getSegments()) {
       TextFragment srcTf = segment.getContent();
 
       // Skip empty segments
       if (srcTf == null || srcTf.isEmpty()) {
         Log.trace(getClass(), "Skipping empty segment in TU '{}'", tu.getId());
         continue;
       }
 
       String sourceText = null;
 
       if (mtSendPlainText) {
         sourceText = srcTf.getText();
 
       } else {
         sourceText = qutil.toCodedHTML(srcTf);
       }
 
       sourceTexts.add(sourceText);
 
       // Store segment info for later mapping
       segmentInfos.add(new SegmentInfo(tu, segment.getId()));
 
       Log.trace(getClass(), "Collected segment [{}]: '{}'",
           sourceTexts.size() - 1, sourceText);
     }
   }
 
   /**
    * Post-process text units: apply translations with confidence scores
    */
   private void postProcessTextUnits(HybridEvaluationResult results) {
     if (results == null || results.getSegments().isEmpty()) {
       Log.warn(getClass(), "No evaluation results available");
       return;
     }
 
     List<HybridScoredSegment> scoredSegments = results.getSegments();
 
     if (scoredSegments.size() != segmentInfos.size()) {
       Log.error(getClass(),
           "Mismatch: {} scored segments but {} segment infos",
           scoredSegments.size(), segmentInfos.size());
 
       return;
     }
 
     // Process each scored segment
     for (int i = 0; i < scoredSegments.size(); i++) {
       HybridScoredSegment scoredSegment = scoredSegments.get(i);
       SegmentInfo segInfo = segmentInfos.get(i);
 
       ITextUnit tu = segInfo.textUnit;
       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 targetSegment = target.getSegments().get(segmentId);
 
       if (targetSegment == null) {
         targetSegment = new Segment(segmentId);
         target.append(targetSegment);
         Log.trace(getClass(), "Created target segment '{}' in TU '{}'", segmentId, tu.getId());
       }
 
       // Get all scored translations for this segment
       List<HybridTranslationScore> scores = scoredSegment.getScores();
 
       if (scores.isEmpty()) {
         Log.warn(getClass(), "No translations available for segment {} in TU '{}'",
             i, tu.getId());
 
         continue;
       }
 
       // Create or get AltTranslationsAnnotation for this segment
       AltTranslationsAnnotation ata = targetSegment.getAnnotation(
           AltTranslationsAnnotation.class);
 
       if (ata == null) {
         ata = new AltTranslationsAnnotation();
         targetSegment.setAnnotation(ata);
         Log.trace(getClass(), "Created AltTranslationsAnnotation for segment '{}'", segmentId);
       }
 
       // Add all translations as alternate translations with confidence scores
       // TextFragment sourceTf = new TextFragment(scoredSegment.getSourceText());
       TextFragment sourceTf = scoredSegment.getSourceTf();
 
       for (HybridTranslationScore score : scores) {
         String targetText = score.getTranslation();
         TextFragment targetTf = null;
 
         if (mtSendPlainText) {
           targetTf = new TextFragment(targetText);
 
         } else {
           targetTf = qutil.fromCodedHTMLToFragment(targetText, null);
         }
 
         OkapiUtil.removeExtraCodes(sourceTf.getCodes(), targetTf);
 
         // Align codes and copy metadata from source to target
         TextFragmentUtil.alignAndCopyCodeMetadata(sourceTf, targetTf, true, true);
 
         // Rearrange opening and closing codes
         OkapiUtil.rearrangeCodes(sourceTf.getCodes(), targetTf);
 
         // Create ConfidenceScoredTranslation with all available scores
         ConfidenceScoredTranslation cst = new ConfidenceScoredTranslation(
             getSourceLocale(),
             getTargetLocale(),
             sourceTf,
             sourceTf, // alternate source same as original for MT
             targetTf,
             MatchType.MT, // All are machine translations
             AthUtil.extractLastSection(score.getModelId()),
             score.getConfidence());
 
         // Add metadata
         if (score.isAnomalyFlagged()) {
           cst.setEngine("ANOMALY: " + score.getAnomalyReason());
 
         } else {
           cst.setEngine(score.getMethod());
         }
 
         ata.add(cst);
 
         Log.trace(getClass(),
             "Added alt-trans [{}] confidence={} MetricX={} Heuristic={}: '{}'",
             score.getModelId(),
             score.getConfidence(),
             score.getMetricXScore() != null ? String.format("%.2f", score.getMetricXScore())
                 : "N/A",
 
             score.getHeuristicScore() != null ? String.format("%.3f", score.getHeuristicScore())
                 : "N/A",
 
             score.getTranslation());
       }
 
       // Sort by confidence (already sorted, but ensure it)
       ata.sort();
 
       // Best translation (highest confidence)
       ConfidenceScoredTranslation bestCst = (ConfidenceScoredTranslation) ata.getFirst();
       TextFragment bestTranslation = bestCst.getTarget().getFirstContent();
       targetSegment.setContent(bestTranslation);
 
       Log.debug(getClass(),
           "Set best translation for TU '{}' segment '{}': '{}' (confidence: {})",
           tu.getId(), segmentId, bestTranslation, bestCst.getConfidenceScore());
     }
   }
 
   /**
    * Convert 0-1 double confidence to 0-100 integer percentage
    */
   public static int convertToPercentage(double confidence) {
     return (int) Math.round(confidence * 100.0);
   }
 
   @Override
   protected void clear() {
     sourceTexts.clear();
     segmentInfos.clear();
     results = null;
   }
 
   @Override
   protected void processTuEvents(List<Event> tuEvents) {
     initializeModelConfigs();
 
     // Step 1: Pre-process all text units
     Log.info(getClass(), "Pre-processing {} text units...", tuEvents.size());
 
     for (Event tue : tuEvents) {
       ITextUnit tu = tue.getTextUnit();
       preProcessTextUnit(tu);
     }
 
     Log.info(getClass(), "Collected {} source segments from {} text units",
         sourceTexts.size(), tuEvents.size());
 
     // Step 2: Evaluate translations with all models
     if (!sourceTexts.isEmpty()) {
       try (HybridQualityEstimator hqe = new HybridQualityEstimator(
           ControllerUtil.getProjectId(),
           Const.ATH_GCP_PROJECT_LOCATION,
           mtSendPlainText)) {
 
         Log.info(getClass(),
             "Evaluating translations: {} segments, {} models, {}→{}",
             sourceTexts.size(),
             modelConfigs.size(),
             getSourceLocale(),
             getTargetLocale());
 
         results = hqe.evaluateTranslations(
             sourceTexts,
             sourceTfs,
             getSourceLocale().toString(),
             getTargetLocale().toString(),
             modelConfigs);
 
         Log.info(getClass(),
             "Evaluation complete using strategy: {}",
             results.getStrategy());
 
       } catch (Exception e) {
         Log.error(getClass(), "Error during translation evaluation: {}", e.getMessage(), e);
         // Continue with empty results - segments will remain untranslated
       }
 
     } else {
       Log.warn(getClass(), "No source segments to translate");
     }
 
     // Step 3: Post-process text units with results
     if (results != null) {
       Log.info(getClass(), "Post-processing text units with evaluation results...");
       postProcessTextUnits(results);
     }
 
     Log.info(getClass(),
         "MT Confidence Scoring complete: processed {} TUs, {} segments",
         getNumProcessedTus(), sourceTexts.size());
   }
 }