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

package io.openlineage.spark.agent.lifecycle.plan;

import io.openlineage.client.OpenLineage.Dataset;
import io.openlineage.client.utils.DatasetIdentifier;
import io.openlineage.spark.agent.util.CatalogDatasetFacetUtils;
import io.openlineage.spark.agent.util.PathUtils;
import io.openlineage.spark.api.DatasetFactory;
import io.openlineage.spark.api.OpenLineageContext;
import io.openlineage.spark.api.QueryPlanVisitor;
import java.util.Collections;
import java.util.List;
import lombok.NonNull;
import org.apache.spark.sql.catalyst.catalog.HiveTableRelation;
import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan;

/**
 * Simple visitor to support {@link HiveTableRelation}s in the plan. Both input and output {@link
 * Dataset}s are supported.
 *
 * @param <D>
 */
public class HiveTableRelationVisitor<D extends Dataset>
    extends QueryPlanVisitor<HiveTableRelation, D> {

  private final DatasetFactory<D> factory;

  public HiveTableRelationVisitor(@NonNull OpenLineageContext context, DatasetFactory<D> factory) {
    super(context);
    this.factory = factory;
  }

  @Override
  public List<D> apply(LogicalPlan x) {
    if (!context.getSparkSession().isPresent()) {
      return Collections.emptyList();
    }

    HiveTableRelation hiveTable = (HiveTableRelation) x;
    DatasetIdentifier datasetId =
        PathUtils.fromCatalogTable(hiveTable.tableMeta(), context.getSparkSession().get());

    return Collections.singletonList(
        CatalogDatasetFacetUtils.getCatalogDatasetFacetForHive(context)
            .map(
                catalogDatasetFacet ->
                    factory.getDataset(datasetId, x.schema(), catalogDatasetFacet))
            .orElse(factory.getDataset(datasetId, x.schema())));
  }
}
