#include "tensorflow/core/framework/tensor_testutil.h"
#include "tsl/lib/core/status_test_util.h"

namespace tensorflow {
namespace {

static bool Initialized = [] {
  auto& rollout_config = GetXlaOpsCommonFlags()->tf_xla_use_device_api;
  rollout_config.enabled_for_xla_launch_ = true;
  rollout_config.enabled_for_compile_on_demand_ = true;

  tensorflow::GetXlaDeviceFlags()->tf_xla_enable_xla_devices = true;
  return true;

}

TEST(XlaCompileUtilTest, PjRtXlaLaunchFlagTest) {
  EXPECT_FALSE(UsePjRtForSingleDeviceCompilation(DeviceType(DEVICE_CPU)));

  // Flag is turned on, but no device is allowlisted.
  auto& rollout_config = GetXlaOpsCommonFlags()->tf_xla_use_device_api;
  rollout_config.enabled_for_xla_launch_ = true;

  EXPECT_FALSE(UsePjRtForSingleDeviceCompilation(DeviceType(DEVICE_CPU)));

    }
    return GetXlaOpsCommonFlags()->tf_xla_async_compilation
               ? DeviceCompileMode::kAsync
               : DeviceCompileMode::kLazy;
  }();

  bool use_pjrt =
      GetXlaOpsCommonFlags()
          ->tf_xla_use_device_api.IsEnabledInXlaCompileAndRunForDevice(
              platform_info_.device_type());

  if (GetXlaOpsCommonFlags()->tf_xla_always_defer_compilation ||

    TF_RETURN_IF_ERROR(status);
  }
  FixupSourceAndSinkEdges(graph.get());
  return graph;
}

bool UsePjRtForSingleDeviceCompilation(const DeviceType& device_type) {
  const auto& rollout_config = GetXlaOpsCommonFlags()->tf_xla_use_device_api;
  return rollout_config.IsEnabledInXlaLaunchForDevice(device_type) ||
         rollout_config.IsEnabledInXlaCompileOnDemandForDevice(device_type) ||

}

class PjRtExecutionUtilGpuTest : public OpsTestBase {
 public:
  PjRtExecutionUtilGpuTest() {
    // Set flag to use PJRT for device compilation and execution.
    auto& rollout_config = GetXlaOpsCommonFlags()->tf_xla_use_device_api;
    rollout_config.enabled_for_xla_launch_ = true;
    rollout_config.enabled_for_compile_on_demand_ = true;
    rollout_config.enabled_for_gpu_ = true;

  std::vector<const Tensor*> inputs = InputsFromContext(ctx);
  std::vector<int> variable_indices =
      GetResourceVariableIndicesFromContext(ctx);

  bool use_pjrt =
      GetXlaOpsCommonFlags()
          ->tf_xla_use_device_api.IsEnabledInXlaCompileOnDemandForDevice(
              platform_info_.device_type());
  if (use_pjrt) {
    std::vector<VariableInfo> variables;

MarkForCompilationPassFlags* GetMarkForCompilationPassFlags();
BuildXlaOpsPassFlags* GetBuildXlaOpsPassFlags();
XlaSparseCoreFlags* GetXlaSparseCoreFlags();
XlaDeviceFlags* GetXlaDeviceFlags();
XlaOpsCommonFlags* GetXlaOpsCommonFlags();
XlaCallModuleFlags* GetXlaCallModuleFlags();

MlirCommonFlags* GetMlirCommonFlags();

void ResetJitCompilerFlags();

const JitRtFlags& GetJitRtFlags();

  return sparse_core_flags;
}

XlaDeviceFlags* GetXlaDeviceFlags() {
  absl::call_once(flags_init, &AllocateAndParseFlags);
  return device_flags;
}

XlaOpsCommonFlags* GetXlaOpsCommonFlags() {
  absl::call_once(flags_init, &AllocateAndParseFlags);
  return ops_flags;
}

XlaCallModuleFlags* GetXlaCallModuleFlags() {
  absl::call_once(flags_init, &AllocateAndParseFlags);

}

class PjRtExecutionUtilTest : public OpsTestBase {
 public:
  PjRtExecutionUtilTest() {
    // Set flag to use PJRT for device compilation and execution.
    auto& rollout_config = GetXlaOpsCommonFlags()->tf_xla_use_device_api;
    rollout_config.enabled_for_xla_launch_ = true;
    rollout_config.enabled_for_compile_on_demand_ = true;

    // Set flag to enable using XLA devices. PJRT currently is only supported