package com.acumenvelocity.ath.service;
   
   import java.io.IOException;
   import java.net.URI;
   import java.util.ArrayList;
   import java.util.Date;
   import java.util.List;
   import java.util.Map;
   import java.util.UUID;
  import java.util.concurrent.ExecutorService;
  import java.util.concurrent.Executors;
  import java.util.stream.Collectors;
  
  import org.apache.solr.client.solrj.SolrClient;
  import org.apache.solr.client.solrj.SolrQuery;
  import org.apache.solr.client.solrj.SolrServerException;
  import org.apache.solr.client.solrj.response.QueryResponse;
  import org.apache.solr.common.SolrDocument;
  import org.apache.solr.common.SolrDocumentList;
  import org.apache.solr.common.SolrInputDocument;
  
  import com.acumenvelocity.ath.common.Const;
  import com.acumenvelocity.ath.gcs.AthStorage;
  import com.acumenvelocity.ath.model.DatasetStats;
  import com.acumenvelocity.ath.model.ModelInfo;
  import com.acumenvelocity.ath.model.TrainingFailedResponse.ErrorTypeEnum;
  import com.acumenvelocity.ath.model.TrainingJobStatus;
  import com.acumenvelocity.ath.solr.AthIndex;
  import com.google.api.gax.longrunning.OperationFuture;
  import com.google.cloud.storage.Blob;
  import com.google.cloud.storage.Storage;
  import com.google.cloud.translate.v3.CreateDatasetRequest;
  import com.google.cloud.translate.v3.CreateModelRequest;
  import com.google.cloud.translate.v3.Dataset;
  import com.google.cloud.translate.v3.DatasetInputConfig;
  import com.google.cloud.translate.v3.GcsInputSource;
  import com.google.cloud.translate.v3.ImportDataRequest;
  import com.google.cloud.translate.v3.LocationName;
  import com.google.cloud.translate.v3.Model;
  import com.google.cloud.translate.v3.TranslationServiceClient;
  import com.google.protobuf.Empty;
  
  /**
   * Service layer for AutoML training operations with Solr persistence.
   * Handles the complete training lifecycle including:
   * <ul>
   * <li>GCS bucket scanning for training files</li>
   * <li>AutoML dataset creation and data import</li>
   * <li>Model training orchestration</li>
   * <li>Training job status tracking and persistence in Solr</li>
   * </ul>
   * 
   * <p>
   * This service manages asynchronous training jobs using a thread pool executor.
   * Each training job progresses through the following states:
   * </p>
   * 
   * <pre>
   * SCANNING_BUCKETS → PREPARING_DATASET → IMPORTING_DATA → TRAINING → COMPLETED/FAILED
   * </pre>
   * 
   * <p>
   * <b>AutoML Limits:</b>
   * </p>
   * <ul>
   * <li>Minimum segments: 1,000</li>
   * <li>Maximum segments: 15,000,000</li>
   * <li>Supported file formats: TMX, TSV, CSV, TXT</li>
   * <li>Automatic 80/10/10 train/validation/test split</li>
   * </ul>
   * 
   * <p>
   * <b>Persistence:</b>
   * Training job state is persisted in Solr (ath_automl_jobs core) for:
   * <ul>
   * <li>Durability across service restarts</li>
   * <li>Multi-node coordination</li>
   * <li>Status queries and reporting</li>
   * </ul>
   * </p>
   * 
   * @author Acumen Velocity
   * @version 1.0
   * @since 1.0
   */
  public class AutoMlTrainingService {
  
    private static final int MIN_SEGMENTS = 1_000;
    private static final int MAX_SEGMENTS = 15_000_000;
  
    private static final List<String> SUPPORTED_EXTENSIONS = List.of(
        ".tmx", ".tsv", ".csv", ".txt");
  
    private final ExecutorService executorService;
    private final Storage storageClient;
    private final SolrClient solrClient;
  
    /**
     * Converts a Solr document to a ModelInfo object.
    * 
    * @param doc the Solr document
    * @return ModelInfo object
    */
   private ModelInfo documentToModelInfo(SolrDocument doc) {
     ModelInfo info = new ModelInfo();
 
     String trainingJobId = (String) doc.getFieldValue("trainingJobId");
     info.trainingJobId(trainingJobId != null ? UUID.fromString(trainingJobId) : null);
 
     info.modelId((String) doc.getFieldValue("modelId"));
     info.modelName((String) doc.getFieldValue("modelName"));
     info.datasetId((String) doc.getFieldValue("datasetId"));
     info.srcLang((String) doc.getFieldValue("srcLang"));
     info.trgLang((String) doc.getFieldValue("trgLang"));
     info.createdAt((Date) doc.getFieldValue("completedAt"));
 
     // Build dataset stats if available
     Integer filesProcessed = (Integer) doc.getFieldValue("filesProcessed");
     Long totalSegments = (Long) doc.getFieldValue("totalSegments");
 
     if (filesProcessed != null && totalSegments != null) {
       DatasetStats stats = new DatasetStats()
           .filesProcessed(filesProcessed)
           .totalSegments(totalSegments)
           .trainSegments((Long) doc.getFieldValue("trainSegments"))
           .validationSegments((Long) doc.getFieldValue("validationSegments"))
           .testSegments((Long) doc.getFieldValue("testSegments"));
 
       info.datasetStats(stats);
     }
 
     return info;
   }
 
   /**
    * Internal class to hold training job context for async execution.
    */
   private static class TrainingJobContext {
     final UUID trainingJobId;
     final String modelName;
     final String srcLang;
     final String trgLang;
     final List<URI> gcsBuckets;
     final String projectId;
     final String location;
     // final UUID userId;
 
     TrainingJobContext(UUID trainingJobId, String modelName, String srcLang, String trgLang,
         List<URI> gcsBuckets, String projectId, String location, UUID userId) {
       this.trainingJobId = trainingJobId;
       this.modelName = modelName;
       this.srcLang = srcLang;
       this.trgLang = trgLang;
       this.gcsBuckets = gcsBuckets;
       this.projectId = projectId;
       this.location = location;
       // this.userId = userId;
     }
   }
 
   /*
    * Default constructor.
    * Initializes storage client, Solr client, and thread pool for async operations.
    */
   public AutoMlTrainingService() {
     this.executorService = Executors.newFixedThreadPool(5);
     this.storageClient = AthStorage.getGcs();
     this.solrClient = AthIndex.getSolr().getClient();
   }
 
   /**
    * Constructor with dependency injection for testing.
    * 
    * @param solrClient    the Solr client to use
    * @param storageClient the GCS storage client to use
    */
   public AutoMlTrainingService(SolrClient solrClient, Storage storageClient) {
     this.executorService = Executors.newFixedThreadPool(5);
     this.storageClient = storageClient;
     this.solrClient = solrClient;
   }
 
   /**
    * Checks if a training job with the given ID exists in Solr.
    * 
    * @param trainingJobId the training job UUID
    * @return true if the job exists, false otherwise
    */
   public boolean trainingJobExists(UUID trainingJobId) {
     try {
       SolrQuery query = new SolrQuery("trainingJobId:" + trainingJobId.toString());
       QueryResponse response = solrClient.query(Const.SOLR_CORE_ATH_AUTOML_JOBS, query);
       return response.getResults().getNumFound() > 0;
 
     } catch (Exception e) {
       throw new RuntimeException("Error checking if training job exists", e);
     }
   }
 
   /**
    * Initiates an asynchronous AutoML training job.
    * 
    * <p>
    * This method:
    * <ol>
    * <li>Creates a training job record in Solr</li>
    * <li>Submits the job to the executor service</li>
    * <li>Returns immediately (non-blocking)</li>
    * </ol>
    * </p>
    * 
    * <p>
    * The training process runs asynchronously and updates the job status in Solr as it progresses.
    * </p>
    * 
    * @param trainingJobId UUID for the training job
    * @param modelName     display name for the model
    * @param srcLang       source language ISO code
    * @param trgLang       target language ISO code
    * @param gcsBuckets    list of GCS bucket URIs to scan
    * @param projectId     GCP project ID
    * @param location      GCP region (e.g., "us-central1")
    * @param userId        UUID of the user initiating training
    * @throws IllegalArgumentException if segment count violates AutoML limits
    */
   public void initiateTraining(
       UUID trainingJobId,
       String modelName,
       String srcLang,
       String trgLang,
       List<URI> gcsBuckets,
       String projectId,
       String location,
       UUID userId) {
 
     try {
       // Create initial training job document in Solr
       SolrInputDocument doc = new SolrInputDocument();
       doc.addField("trainingJobId", trainingJobId.toString());
       doc.addField("modelName", modelName);
       doc.addField("srcLang", srcLang);
       doc.addField("trgLang", trgLang);
       doc.addField("projectId", projectId);
       doc.addField("location", location);
       doc.addField("createdBy", userId.toString());
       doc.addField("createdAt", new Date());
       doc.addField("startedAt", new Date());
       doc.addField("status", TrainingJobStatus.SCANNING_BUCKETS.name());
       doc.addField("progress", 0);
       doc.addField("currentPhase", "Scanning GCS buckets for training files");
 
       // Add GCS buckets as multi-valued field
       for (URI bucket : gcsBuckets) {
         doc.addField("gcsBuckets", bucket.toString());
       }
 
       solrClient.add(Const.SOLR_CORE_ATH_AUTOML_JOBS, doc);
       solrClient.commit(Const.SOLR_CORE_ATH_AUTOML_JOBS);
 
       // Submit async training task
       TrainingJobContext context = new TrainingJobContext(
           trainingJobId, modelName, srcLang, trgLang, gcsBuckets, projectId, location, userId);
 
       executorService.submit(() -> executeTraining(context));
 
     } catch (Exception e) {
       throw new RuntimeException("Error initiating training job", e);
     }
   }
 
   /**
    * Gets the current status of a training job from Solr.
    * 
    * @param trainingJobId the training job UUID
    * @return current status enum, or null if job not found
    */
   public TrainingJobStatus getTrainingStatus(UUID trainingJobId) {
     try {
       SolrDocument doc = getJobDocument(trainingJobId);
       if (doc == null) {
         return null;
       }
 
       String statusStr = (String) doc.getFieldValue("status");
       return statusStr != null ? TrainingJobStatus.valueOf(statusStr) : null;
 
     } catch (Exception e) {
       throw new RuntimeException("Error getting training status", e);
     }
   }
 
   /**
    * Gets the progress percentage of a training job.
    * 
    * @param trainingJobId the training job UUID
    * @return progress percentage (0-100)
    */
   public int getProgress(UUID trainingJobId) {
     try {
       SolrDocument doc = getJobDocument(trainingJobId);
       Integer progress = (Integer) doc.getFieldValue("progress");
       return progress != null ? progress : 0;
 
     } catch (Exception e) {
       return 0;
     }
   }
 
   /**
    * Gets the start timestamp of a training job.
    * 
    * @param trainingJobId the training job UUID
    * @return start date
    */
   public Date getStartedAt(UUID trainingJobId) {
     try {
       SolrDocument doc = getJobDocument(trainingJobId);
       return (Date) doc.getFieldValue("startedAt");
 
     } catch (Exception e) {
       return null;
     }
   }
 
   /**
    * Gets the current phase description of a training job.
    * 
    * @param trainingJobId the training job UUID
    * @return human-readable phase description
    */
   public String getCurrentPhase(UUID trainingJobId) {
     try {
       SolrDocument doc = getJobDocument(trainingJobId);
       return (String) doc.getFieldValue("currentPhase");
 
     } catch (Exception e) {
       return null;
     }
   }
 
   /**
    * Gets the AutoML model ID for a completed training job.
    * 
    * @param trainingJobId the training job UUID
    * @return AutoML model resource name
    */
   public String getModelId(UUID trainingJobId) {
     try {
       SolrDocument doc = getJobDocument(trainingJobId);
       return (String) doc.getFieldValue("modelId");
 
     } catch (Exception e) {
       return null;
     }
   }
 
   /**
    * Gets the model name for a completed training job.
    * 
    * @param trainingJobId the training job UUID
    * @return model display name
    */
   public String getModelName(UUID trainingJobId) {
     try {
       SolrDocument doc = getJobDocument(trainingJobId);
       return (String) doc.getFieldValue("modelName");
 
     } catch (Exception e) {
       return null;
     }
   }
 
   /**
    * Gets the AutoML dataset ID for a training job.
    * 
    * @param trainingJobId the training job UUID
    * @return AutoML dataset resource name
    */
   public String getDatasetId(UUID trainingJobId) {
     try {
       SolrDocument doc = getJobDocument(trainingJobId);
       return (String) doc.getFieldValue("datasetId");
 
     } catch (Exception e) {
       return null;
     }
   }
 
   /**
    * Gets the completion timestamp for a training job.
    * 
    * @param trainingJobId the training job UUID
    * @return completion date
    */
   public Date getCompletedAt(UUID trainingJobId) {
     try {
       SolrDocument doc = getJobDocument(trainingJobId);
       return (Date) doc.getFieldValue("completedAt");
 
     } catch (Exception e) {
       return null;
     }
   }
 
   /**
    * Gets the dataset statistics for a training job.
    * 
    * @param trainingJobId the training job UUID
    * @return dataset statistics object
    */
   public DatasetStats getDatasetStats(UUID trainingJobId) {
     try {
       SolrDocument doc = getJobDocument(trainingJobId);
       if (doc == null) {
         return null;
       }
 
       Integer filesProcessed = (Integer) doc.getFieldValue("filesProcessed");
       Long totalSegments = (Long) doc.getFieldValue("totalSegments");
       Long trainSegments = (Long) doc.getFieldValue("trainSegments");
       Long validationSegments = (Long) doc.getFieldValue("validationSegments");
       Long testSegments = (Long) doc.getFieldValue("testSegments");
 
       if (filesProcessed == null || totalSegments == null) {
         return null;
       }
 
       return new DatasetStats()
           .filesProcessed(filesProcessed)
           .totalSegments(totalSegments)
           .trainSegments(trainSegments)
           .validationSegments(validationSegments)
           .testSegments(testSegments);
 
     } catch (Exception e) {
       return null;
     }
   }
 
   /**
    * Gets the error type for a failed training job.
    * 
    * @param trainingJobId the training job UUID
    * @return error type enum
    */
   public ErrorTypeEnum getErrorType(UUID trainingJobId) {
     try {
       SolrDocument doc = getJobDocument(trainingJobId);
       String errorType = (String) doc.getFieldValue("errorType");
       return errorType != null ? ErrorTypeEnum.valueOf(errorType) : null;
 
     } catch (Exception e) {
       return null;
     }
   }
 
   /**
    * Gets the error message for a failed training job.
    * 
    * @param trainingJobId the training job UUID
    * @return error message string
    */
   public String getErrorMessage(UUID trainingJobId) {
     try {
       SolrDocument doc = getJobDocument(trainingJobId);
       return (String) doc.getFieldValue("errorMessage");
 
     } catch (Exception e) {
       return null;
     }
   }
 
   /**
    * Gets the failure timestamp for a failed training job.
    * 
    * @param trainingJobId the training job UUID
    * @return failure date
    */
   public Date getFailedAt(UUID trainingJobId) {
     try {
       SolrDocument doc = getJobDocument(trainingJobId);
       return (Date) doc.getFieldValue("failedAt");
 
     } catch (Exception e) {
       return null;
     }
   }
 
   /**
    * Returns the user ID who owns (initiated) this training job.
    * 
    * @param trainingJobId the training job UUID
    * @return user UUID who created the job
    */
   public UUID getJobOwnerId(UUID trainingJobId) {
     try {
       SolrDocument doc = getJobDocument(trainingJobId);
       String userId = (String) doc.getFieldValue("createdBy");
       return userId != null ? UUID.fromString(userId) : null;
 
     } catch (Exception e) {
       return null;
     }
   }
 
   /**
    * Cancels an in-progress training job.
    * 
    * @param trainingJobId the training job UUID
    */
   public void cancelTraining(UUID trainingJobId) {
     try {
       SolrInputDocument doc = new SolrInputDocument();
       doc.addField("trainingJobId", trainingJobId.toString());
 
       // Partial update - only set status to CANCELLED
       doc.addField("status", Map.of("set", TrainingJobStatus.CANCELLED.name()));
 
       solrClient.add(Const.SOLR_CORE_ATH_AUTOML_JOBS, doc);
       solrClient.commit(Const.SOLR_CORE_ATH_AUTOML_JOBS);
 
       // Note: Actual AutoML operation cancellation would be implemented by
       // checking the cancelled status in the training loop
 
     } catch (Exception e) {
       throw new RuntimeException("Error cancelling training job", e);
     }
   }
 
   /**
    * Gets a list of all successfully trained models from Solr.
    * 
    * @return list of model information objects
    */
   public List<ModelInfo> getAllModels() {
     try {
       SolrQuery query = new SolrQuery("status:" + TrainingJobStatus.COMPLETED.name());
       query.setSort("completedAt", SolrQuery.ORDER.desc); // Newest first
       query.setRows(1000); // Max results
 
       QueryResponse response = solrClient.query(Const.SOLR_CORE_ATH_AUTOML_JOBS, query);
       SolrDocumentList docs = response.getResults();
 
       return docs.stream()
           .map(this::documentToModelInfo)
           .collect(Collectors.toList());
 
     } catch (Exception e) {
       throw new RuntimeException("Error fetching models", e);
     }
   }
 
   /**
    * Executes the complete training workflow asynchronously.
    * Updates Solr at each phase.
    * 
    * @param context the training job context
    */
   private void executeTraining(TrainingJobContext context) {
     try {
       // Phase 1: Scan GCS buckets for training files
       updateJobStatus(context.trainingJobId, TrainingJobStatus.SCANNING_BUCKETS, 10,
           "Scanning GCS buckets for training files");
 
       List<String> trainingFiles = scanGcsBuckets(context.gcsBuckets);
 
       if (trainingFiles.isEmpty()) {
         failJob(context.trainingJobId, ErrorTypeEnum.INVALID_FILES,
             "No valid training files found in specified buckets");
         return;
       }
 
       // Estimate segment count
       int estimatedSegments = trainingFiles.size() * 1000;
 
       if (estimatedSegments < MIN_SEGMENTS) {
         failJob(context.trainingJobId, ErrorTypeEnum.INSUFFICIENT_DATA,
             "Dataset contains approximately " + estimatedSegments +
                 " segments. Minimum required: " + MIN_SEGMENTS);
         return;
       }
 
       if (estimatedSegments > MAX_SEGMENTS) {
         failJob(context.trainingJobId, ErrorTypeEnum.EXCESSIVE_DATA,
             "Dataset contains approximately " + estimatedSegments +
                 " segments. Maximum allowed: " + MAX_SEGMENTS);
         return;
       }
 
       // Check if cancelled
       if (isCancelled(context.trainingJobId)) {
         return;
       }
 
       // Phase 2: Create AutoML dataset
       updateJobStatus(context.trainingJobId, TrainingJobStatus.PREPARING_DATASET, 25,
           "Creating AutoML dataset");
 
       String datasetId = createAutoMLDataset(context);
       updateDatasetId(context.trainingJobId, datasetId);
 
       if (isCancelled(context.trainingJobId)) {
         return;
       }
 
       // Phase 3: Import training data
       updateJobStatus(context.trainingJobId, TrainingJobStatus.IMPORTING_DATA, 40,
           "Importing " + trainingFiles.size() + " files to dataset");
 
       importDataToDataset(context, datasetId, trainingFiles);
 
       if (isCancelled(context.trainingJobId)) {
         return;
       }
 
       // Phase 4: Train model
       updateJobStatus(context.trainingJobId, TrainingJobStatus.TRAINING, 60,
           "Training AutoML model");
 
       String modelId = trainAutoMLModel(context, datasetId);
 
       // Phase 5: Complete
       completeJob(context.trainingJobId, modelId, trainingFiles.size(), estimatedSegments);
 
     } catch (InterruptedException e) {
       // Training was cancelled
       Thread.currentThread().interrupt();
 
     } catch (Exception e) {
       failJob(context.trainingJobId, ErrorTypeEnum.TRAINING_ERROR,
           "AutoML training failed: " + e.getMessage());
     }
   }
 
   /**
    * Scans GCS buckets recursively to find supported training files.
    * 
    * @param gcsBuckets list of GCS bucket URIs
    * @return list of GCS URIs for training files
    */
   private List<String> scanGcsBuckets(List<URI> gcsBuckets) {
     List<String> trainingFiles = new ArrayList<>();
 
     for (URI bucketUri : gcsBuckets) {
       String[] parts = bucketUri.toString().replace("gs://", "").split("/", 2);
       String bucketName = parts[0];
       String prefix = parts.length > 1 ? parts[1] : "";
 
       for (Blob blob : storageClient.list(bucketName,
           Storage.BlobListOption.prefix(prefix)).iterateAll()) {
 
         String fileName = blob.getName().toLowerCase();
 
         if (SUPPORTED_EXTENSIONS.stream().anyMatch(fileName::endsWith)) {
           trainingFiles.add("gs://" + bucketName + "/" + blob.getName());
         }
       }
     }
 
     return trainingFiles;
   }
 
   /**
    * Creates an AutoML Translation dataset.
    * 
    * @param context the training job context
    * @return dataset resource name
    * @throws Exception if dataset creation fails
    */
   private String createAutoMLDataset(TrainingJobContext context) throws Exception {
     try (TranslationServiceClient client = TranslationServiceClient.create()) {
       LocationName parent = LocationName.of(context.projectId, context.location);
 
       Dataset dataset = Dataset.newBuilder()
           .setDisplayName(context.modelName + "-dataset")
           .setSourceLanguageCode(context.srcLang)
           .setTargetLanguageCode(context.trgLang)
           .build();
 
       CreateDatasetRequest request = CreateDatasetRequest.newBuilder()
           .setParent(parent.toString())
           .setDataset(dataset)
           .build();
 
       OperationFuture<Dataset, ?> operationFuture = client.createDatasetAsync(request);
       Dataset createdDataset = operationFuture.get();
 
       return createdDataset.getName();
     }
   }
 
   /**
    * Imports training files into an AutoML dataset.
    * 
    * @param context       the training job context
    * @param datasetId     the dataset resource name
    * @param trainingFiles list of GCS URIs for training files
    * @throws Exception if import fails
    */
   private void importDataToDataset(TrainingJobContext context, String datasetId,
       List<String> trainingFiles) throws Exception {
 
     try (TranslationServiceClient client = TranslationServiceClient.create()) {
       DatasetInputConfig.Builder inputConfigBuilder = DatasetInputConfig.newBuilder();
 
       for (String fileUri : trainingFiles) {
         GcsInputSource gcsSource = GcsInputSource.newBuilder()
             .setInputUri(fileUri)
             .build();
 
         inputConfigBuilder.addInputFiles(
             DatasetInputConfig.InputFile.newBuilder()
                 .setGcsSource(gcsSource)
                 .build());
       }
 
       ImportDataRequest request = ImportDataRequest.newBuilder()
           .setDataset(datasetId)
           .setInputConfig(inputConfigBuilder.build())
           .build();
 
       OperationFuture<Empty, ?> operationFuture = client.importDataAsync(request);
       operationFuture.get();
     }
   }
 
   /**
    * Trains an AutoML Translation model.
    * 
    * @param context   the training job context
    * @param datasetId the dataset resource name
    * @return model resource name
    * @throws Exception if training fails
    */
   private String trainAutoMLModel(TrainingJobContext context, String datasetId) throws Exception {
     try (TranslationServiceClient client = TranslationServiceClient.create()) {
       LocationName parent = LocationName.of(context.projectId, context.location);
 
       Model model = Model.newBuilder()
           .setDisplayName(context.modelName)
           .setDataset(datasetId)
           .build();
 
       CreateModelRequest request = CreateModelRequest.newBuilder()
           .setParent(parent.toString())
           .setModel(model)
           .build();
 
       OperationFuture<Model, ?> operationFuture = client.createModelAsync(request);
 
       // Update progress periodically while training
       while (!operationFuture.isDone()) {
         Thread.sleep(60_000); // Check every minute
 
         if (isCancelled(context.trainingJobId)) {
           operationFuture.cancel(true);
           throw new InterruptedException("Training cancelled by user");
         }
 
         // Increment progress (capped at 95%)
         int currentProgress = getProgress(context.trainingJobId);
         updateProgress(context.trainingJobId, Math.min(95, currentProgress + 2));
       }
 
       Model createdModel = operationFuture.get();
       return createdModel.getName();
     }
   }
 
   /**
    * Checks if a training job has been cancelled.
    * 
    * @param trainingJobId the training job UUID
    * @return true if cancelled, false otherwise
    */
   private boolean isCancelled(UUID trainingJobId) {
     TrainingJobStatus status = getTrainingStatus(trainingJobId);
     return TrainingJobStatus.CANCELLED.equals(status);
   }
 
   /**
    * Updates the status of a training job in Solr.
    * 
    * @param trainingJobId the training job UUID
    * @param status        new status
    * @param progress      progress percentage
    * @param phase         current phase description
    */
   private void updateJobStatus(UUID trainingJobId, TrainingJobStatus status, int progress,
       String phase) {
     try {
       SolrInputDocument doc = new SolrInputDocument();
       doc.addField("trainingJobId", trainingJobId.toString());
       doc.addField("status", Map.of("set", status.name()));
       doc.addField("progress", Map.of("set", progress));
       doc.addField("currentPhase", Map.of("set", phase));
 
       solrClient.add(Const.SOLR_CORE_ATH_AUTOML_JOBS, doc);
       solrClient.commit(Const.SOLR_CORE_ATH_AUTOML_JOBS);
 
     } catch (Exception e) {
       throw new RuntimeException("Error updating job status", e);
     }
   }
 
   /**
    * Updates just the progress field in Solr.
    * 
    * @param trainingJobId the training job UUID
    * @param progress      new progress percentage
    */
   private void updateProgress(UUID trainingJobId, int progress) {
     try {
       SolrInputDocument doc = new SolrInputDocument();
       doc.addField("trainingJobId", trainingJobId.toString());
       doc.addField("progress", Map.of("set", progress));
 
       solrClient.add(Const.SOLR_CORE_ATH_AUTOML_JOBS, doc);
       solrClient.commit(Const.SOLR_CORE_ATH_AUTOML_JOBS);
 
     } catch (Exception e) {
       // Log but don't fail - progress updates are non-critical
       System.err.println("Error updating progress: " + e.getMessage());
     }
   }
 
   /**
    * Updates the dataset ID in Solr.
    * 
    * @param trainingJobId the training job UUID
    * @param datasetId     the dataset resource name
    */
   private void updateDatasetId(UUID trainingJobId, String datasetId) {
     try {
       SolrInputDocument doc = new SolrInputDocument();
       doc.addField("trainingJobId", trainingJobId.toString());
       doc.addField("datasetId", Map.of("set", datasetId));
 
       solrClient.add(Const.SOLR_CORE_ATH_AUTOML_JOBS, doc);
       solrClient.commit(Const.SOLR_CORE_ATH_AUTOML_JOBS);
 
     } catch (Exception e) {
       throw new RuntimeException("Error updating dataset ID", e);
     }
   }
 
   /**
    * Marks a training job as completed successfully in Solr.
    * 
    * @param trainingJobId  the training job UUID
    * @param modelId        the model resource name
    * @param filesProcessed number of files processed
    * @param totalSegments  total number of segments
    */
   private void completeJob(UUID trainingJobId, String modelId, int filesProcessed,
       int totalSegments) {
     try {
       SolrInputDocument doc = new SolrInputDocument();
       doc.addField("trainingJobId", trainingJobId.toString());
       doc.addField("status", Map.of("set", TrainingJobStatus.COMPLETED.name()));
       doc.addField("progress", Map.of("set", 100));
       doc.addField("modelId", Map.of("set", modelId));
       doc.addField("completedAt", Map.of("set", new Date()));
       doc.addField("filesProcessed", Map.of("set", filesProcessed));
       doc.addField("totalSegments", Map.of("set", (long) totalSegments));
       doc.addField("trainSegments", Map.of("set", (long) (totalSegments * 0.8)));
       doc.addField("validationSegments", Map.of("set", (long) (totalSegments * 0.1)));
       doc.addField("testSegments", Map.of("set", (long) (totalSegments * 0.1)));
 
       solrClient.add(Const.SOLR_CORE_ATH_AUTOML_JOBS, doc);
       solrClient.commit(Const.SOLR_CORE_ATH_AUTOML_JOBS);
 
     } catch (Exception e) {
       throw new RuntimeException("Error completing job", e);
     }
   }
 
   /**
    * Marks a training job as failed in Solr.
    * 
    * @param trainingJobId the training job UUID
    * @param errorType     machine-readable error type
    * @param errorMessage  human-readable error message
    */
   private void failJob(UUID trainingJobId, ErrorTypeEnum errorType, String errorMessage) {
     try {
       SolrInputDocument doc = new SolrInputDocument();
       doc.addField("trainingJobId", trainingJobId.toString());
       doc.addField("status", Map.of("set", TrainingJobStatus.FAILED.name()));
       doc.addField("errorType", Map.of("set", errorType.name()));
       doc.addField("errorMessage", Map.of("set", errorMessage));
       doc.addField("failedAt", Map.of("set", new Date()));
 
       solrClient.add(Const.SOLR_CORE_ATH_AUTOML_JOBS, doc);
       solrClient.commit(Const.SOLR_CORE_ATH_AUTOML_JOBS);
 
     } catch (Exception e) {
       throw new RuntimeException("Error failing job", e);
     }
   }
 
   /**
    * Retrieves a training job document from Solr.
    * 
    * @param trainingJobId the training job UUID
    * @return Solr document, or null if not found
    */
   private SolrDocument getJobDocument(UUID trainingJobId) throws SolrServerException, IOException {
     SolrQuery query = new SolrQuery("trainingJobId:" + trainingJobId.toString());
     QueryResponse response = solrClient.query(Const.SOLR_CORE_ATH_AUTOML_JOBS, query);
     SolrDocumentList docs = response.getResults();
 
     return docs.isEmpty() ? null : docs.get(0);
   }
 
 }