/* Copyright 2024 The OpenXLA Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/

#ifndef XLA_HLO_IR_BACKEND_CONFIG_H_
#define XLA_HLO_IR_BACKEND_CONFIG_H_

#include <functional>
#include <memory>
#include <string>
#include <utility>

#include "absl/base/thread_annotations.h"
#include "absl/log/log.h"
#include "absl/status/status.h"
#include "absl/status/statusor.h"
#include "absl/synchronization/mutex.h"
#include "google/protobuf/message.h"
#include "tsl/platform/human_readable_json.h"
#include "tsl/platform/protobuf.h"

namespace xla {

template <typename T>
using EnableIfProto = typename std::enable_if_t<
    std::is_base_of<tsl::protobuf::Message, T>::value>;

// Returns a string representation of a proto in the format used by
// HloInstruction::raw_backend_config_string.
//
// This is morally equivalent to:
//
//   HloInstruction instr;
//   TF_RETURN_IF_ERROR(instr.set_backend_config(proto));
//   return instr.raw_backend_config_string();
//
absl::StatusOr<std::string> BackendConfigToRawString(
    const tsl::protobuf::Message& proto);

// Clones the provided proto. If the input is nullptr, the result is also
// nullptr.
std::unique_ptr<tsl::protobuf::Message> CloneBackendConfigProto(
    const tsl::protobuf::Message* proto);

// A wrapper around the BackendConfig proto. It can be initialized either with
// a proto object or a string representing the JSON encoding of a proto. Once
// the wrapper is initialized (either during construction or via an assignment)
// it can only be mutated by calling the ApplyFnOnProto method.
//
// When the wrapper is initialized only the provided format is stored. If the
// other format is requested from the wrapper later, it is lazily computed and
// cached internally, before it is returned. Subsequent accesses will directly
// return the cached value.
//
// All accesses are protected via a mutex because instances of this class are
// accessed concurrently during auto tuning.
class BackendConfigWrapper {
 public:
  BackendConfigWrapper() = default;
  explicit BackendConfigWrapper(std::string raw_string)
      : raw_string_(std::move(raw_string)) {}
  explicit BackendConfigWrapper(const tsl::protobuf::Message& proto)
      : proto_(CloneBackendConfigProto(&proto)) {}
  BackendConfigWrapper(const BackendConfigWrapper& other) {
    absl::MutexLock other_lock{other.mutex_};
    proto_ = CloneBackendConfigProto(other.proto_.get());
    raw_string_ = other.raw_string_;
  }

  BackendConfigWrapper& operator=(BackendConfigWrapper&& other);
  bool operator==(const BackendConfigWrapper& other) const;
  bool operator!=(const BackendConfigWrapper& other) const {
    return !(*this == other);
  }

  // Returns a reference to the raw string that corresponds to this backend
  // config.
  //
  // WARNING: This function returns a reference which is valid at the time the
  //          call terminates. If the BackendConfig is reassigned the reference
  //          becomes invalid, which could lead to subtle and hard to detect
  //          bugs, especially in multi-threaded code. The caller is responsible
  //          for ensuring the lifetime of the referenced string.
  //
  //          Prefer to use the safer (but potentially slower) GetProto().
  const std::string& GetRawString() const {
    absl::WriterMutexLock lock{mutex_};
    return GetRawStringWithoutMutex();
  }
  absl::Status GetProto(tsl::protobuf::Message* output_proto) const;

  // Type trait to check if a type has the mutable_custom_call_metadata method.
  template <typename T, typename = void>
  struct has_mutable_custom_call_metadata : std::false_type {};

  template <typename T>
  struct has_mutable_custom_call_metadata<
      T,
      std::void_t<decltype(std::declval<T>().mutable_custom_call_metadata())>>
      : std::true_type {};

  // Applies a function `fn` to the underlying proto. The function receives a
  // mutable pointer to a proto of type `ConfigProto`.
  //
  // If there is no proto initialized, will try to initialize from the
  // raw_string_. If raw_string_ is also empty, will return an error.
  template <typename ConfigProto, EnableIfProto<ConfigProto>* = nullptr>
  absl::Status ApplyFnOnProto(
      const std::function<absl::Status(ConfigProto*)>& fn) {
    absl::WriterMutexLock lock{mutex_};
    if (proto_ == nullptr) {
      if (raw_string_.empty()) {
        return absl::InvalidArgumentError(
            "Has no proto to apply the modifier function on.");
      }
      auto proto = std::make_unique<ConfigProto>();
      absl::Status status =
          tsl::HumanReadableJsonToProto(raw_string_, proto.get());
      if (!status.ok()) {
        // If the proto is unparsable, move the raw string to the custom call
        // metadata field. This preserves the original string for lowering
        // emitters that might depend on it.
        if constexpr (has_mutable_custom_call_metadata<ConfigProto>::value) {
          auto* custom_call_metadata = proto->mutable_custom_call_metadata();
          custom_call_metadata->set_metadata(raw_string_);
          VLOG(1) << "Moving the unparsable backend config " << raw_string_
                  << " to custom call metadata.";
          raw_string_.clear();
        } else {
          return status;
        }
      }
      proto_ = std::move(proto);
    }
    absl::Status status = fn(static_cast<ConfigProto*>(proto_.get()));
    // Invalidate string cache, as the proto might have been mutated.
    raw_string_.clear();
    return status;
  }

  bool empty() const {
    absl::MutexLock lock{mutex_};
    return proto_ == nullptr && raw_string_.empty();
  }

 private:
  const std::string& GetRawStringWithoutMutex() const
      ABSL_EXCLUSIVE_LOCKS_REQUIRED(mutex_);

  // proto_ and raw_string_ must be consistent. If one is set, the other
  // will be lazily initialized when requested. Because this class is accessed
  // concurrently, a mutex is used to protect all access.
  //
  // Unfortunately, all members have to be mutable, since either of them can be
  // the cached one.
  mutable absl::Mutex mutex_;
  mutable std::unique_ptr<tsl::protobuf::Message> proto_
      ABSL_GUARDED_BY(mutex_);
  mutable std::string raw_string_ ABSL_GUARDED_BY(mutex_);
};

}  // namespace xla

#endif  // XLA_HLO_IR_BACKEND_CONFIG_H_