#include "tensorflow/core/framework/cancellation.h"

struct TFE_CancellationManager;
typedef struct TFE_CancellationManager TFE_CancellationManager;

namespace tensorflow {
DEFINE_CONVERSION_FUNCTIONS(tensorflow::CancellationManager,
                            TFE_CancellationManager);
DEFINE_CONVERSION_FUNCTIONS(tensorflow::CancellationManager*,
                            TFE_CancellationManager*);
}  // namespace tensorflow

TF_CAPI_EXPORT extern bool TFE_CancellationManagerIsCancelling(
    TFE_CancellationManager*);
TF_CAPI_EXPORT extern void TFE_CancellationManagerStartCancel(
    TFE_CancellationManager*);
TF_CAPI_EXPORT extern TFE_CancellationToken TFE_CancellationManagerGetToken(
    TFE_CancellationManager*);
TF_CAPI_EXPORT extern bool TFE_CancellationManagerRegisterCallback(
    TFE_CancellationManager*, TFE_CancellationToken token,

}

TFE_CancellationManager* TFE_NewCancellationManager() {
  return tensorflow::wrap(new tensorflow::CancellationManager);
}

void TFE_CancellationManagerStartCancel(
    TFE_CancellationManager* cancellation_manager) {
  tensorflow::unwrap(cancellation_manager)->StartCancel();
}

bool TFE_CancellationManagerIsCancelled(
    TFE_CancellationManager* cancellation_manager) {

  void StartExecute(TFE_Context* context, const char* operation_name,
                    std::vector<TFE_TensorHandle*> inputs,
                    const TFE_OpAttrs* attributes, int expected_max_outputs,
                    CancellationManager& cancellation_manager,
                    absl::optional<int64_t> step_id = absl::nullopt);
  // Block until the previous `StartExecute` operation has executed. Forwards

                    const std::vector<ParallelTensor*>& inputs,
                    const char* operation_name, const TFE_OpAttrs* attributes,
                    int expected_max_outputs,
                    CancellationManager& cancellation_manager,
                    std::optional<int64_t> step_id = std::nullopt) const;

  void StartExecute(TFE_Context* context,
                    const std::vector<std::vector<TFE_TensorHandle*>>& inputs,

  virtual const tensorflow::AbstractOpAttrs* GetOpAttrs() const = 0;
  virtual void AddAttrs(const AbstractOpAttrs* op_attrs) = 0;

  virtual void SetCancellationManager(
      CancellationManager* cancellation_manager) = 0;

  // Returns the stack trace set by `SetStackTrace` if exists.
  virtual absl::optional<ManagedStackTrace> GetStackTrace() = 0;

  virtual void SetStepId(int64_t step_id) = 0;

  TF_DeleteBuffer(result2);
  TFE_MonitoringDeleteSampler2(sampler2);

  TFE_MonitoringDeleteBuckets(buckets);
  TF_DeleteStatus(status);
}

TEST(CAPI, CancellationManager) {
  TFE_CancellationManager* c_mgr = TFE_NewCancellationManager();
  EXPECT_FALSE(TFE_CancellationManagerIsCancelled(c_mgr));

  TFE_CancelCallback callback1;
  callback1.callback = [](void* context) {

  TFE_OpSetAttrShape(handle_op.get(), "shape", /*dims=*/nullptr, /*num_dims=*/0,
                     status.get());
  ASSERT_TRUE(TF_GetCode(status.get()) == TF_OK) << TF_Message(status.get());
  CancellationManager cancellation_manager;
  parallel_device.StartExecute(context.get(), std::vector<ParallelTensor*>(),
                               "VarHandleOp", TFE_OpGetAttrs(handle_op.get()),

    ],
    visibility = [
        "//tensorflow:internal",
    ],
    deps = [
        ":tfe_tensorhandle_internal",
    ],
)

cc_header_only_library(
    name = "tfe_cancellationmanager_internal_hdrs_only",
    extra_deps = [
        "@com_google_absl//absl/strings",
    ],
    visibility = [
        "//tensorflow:internal",
    ],
    deps = [