/*
/* Copyright 2018-2026 contributors to the OpenLineage project
/* SPDX-License-Identifier: Apache-2.0
*/

package io.openlineage.spark.agent.lifecycle;

import io.openlineage.client.OpenLineage;
import io.openlineage.client.OpenLineage.InputDataset;
import io.openlineage.client.utils.DatasetIdentifier;
import io.openlineage.spark.agent.util.PathUtils;
import io.openlineage.spark.api.AbstractInputDatasetBuilder;
import io.openlineage.spark.api.OpenLineageContext;
import java.io.IOException;
import java.net.URI;
import java.util.ArrayList;
import java.util.Collection;
import java.util.List;
import java.util.stream.Collectors;
import lombok.NonNull;
import lombok.extern.slf4j.Slf4j;
import org.apache.hadoop.fs.Path;
import org.apache.hadoop.mapreduce.Job;
import org.apache.spark.rdd.HadoopRDD;
import org.apache.spark.rdd.NewHadoopRDD;
import org.apache.spark.rdd.RDD;

/** Build a collection of {@link InputDataset}s from a {@link HadoopRDD} */
@Slf4j
public class HadoopRDDInputDatasetBuilder extends AbstractInputDatasetBuilder<RDD<?>> {

  public HadoopRDDInputDatasetBuilder(@NonNull OpenLineageContext context) {
    super(context);
  }

  @Override
  public boolean isDefinedAt(RDD<?> x) {
    return x instanceof HadoopRDD || x instanceof NewHadoopRDD;
  }

  @Override
  public Collection<InputDataset> apply(RDD<?> x) {
    return findInputs(x).stream().map(this::buildInputDataset).collect(Collectors.toList());
  }

  protected InputDataset buildInputDataset(URI uri) {
    DatasetIdentifier di = PathUtils.fromURI(uri);
    return new OpenLineage.InputDatasetBuilder()
        .name(di.getName())
        .namespace(di.getNamespace())
        .build();
  }

  protected List<URI> findInputs(RDD<?> rdd) {
    List<URI> result = new ArrayList<>();
    Path[] inputPaths = getInputPaths(rdd);
    if (inputPaths != null) {
      for (Path path : inputPaths) {
        result.add(getDatasetUri(path.toUri()));
      }
    }
    return result;
  }

  protected Path[] getInputPaths(RDD<?> rdd) {
    Path[] inputPaths = null;
    if (rdd instanceof HadoopRDD) {
      inputPaths =
          org.apache.hadoop.mapred.FileInputFormat.getInputPaths(
              ((HadoopRDD<?, ?>) rdd).getJobConf());
    } else if (rdd instanceof NewHadoopRDD) {
      try {
        inputPaths =
            org.apache.hadoop.mapreduce.lib.input.FileInputFormat.getInputPaths(
                new Job(((NewHadoopRDD<?, ?>) rdd).getConf()));
      } catch (IOException e) {
        log.error("Openlineage spark agent could not get input paths", e);
      }
    }
    return inputPaths;
  }

  // exposed for testing
  protected URI getDatasetUri(URI pathUri) {
    return pathUri;
  }
}
