/*
 * Copyright (C) 2017 The Android Open Source Project
 *
 * 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 HIDL_CALLBACK_UTIL_H_
#define HIDL_CALLBACK_UTIL_H_

#include <set>

#include <hidl/HidlSupport.h>

namespace {
// Type of callback invoked by the death handler.
using on_death_cb_function = std::function<void(uint64_t)>;

// Private class used to keep track of death of individual
// callbacks stored in HidlCallbackHandler.
template <typename CallbackType>
class HidlDeathHandler : public android::hardware::hidl_death_recipient {
 public:
  HidlDeathHandler(const on_death_cb_function& user_cb_function)
      : cb_function_(user_cb_function) {}
  ~HidlDeathHandler() = default;

  // Death notification for callbacks.
  void serviceDied(
      uint64_t cookie,
      const android::wp<android::hidl::base::V1_0::IBase>& /* who */) override {
    cb_function_(cookie);
  }

 private:
  on_death_cb_function cb_function_;

  DISALLOW_COPY_AND_ASSIGN(HidlDeathHandler);
};
}  // namespace

namespace android {
namespace hardware {
namespace wifi {
namespace V1_0 {
namespace implementation {
namespace hidl_callback_util {
template <typename CallbackType>
// Provides a class to manage callbacks for the various HIDL interfaces and
// handle the death of the process hosting each callback.
class HidlCallbackHandler {
 public:
  HidlCallbackHandler()
      : death_handler_(new HidlDeathHandler<CallbackType>(
            std::bind(&HidlCallbackHandler::onObjectDeath,
                      this,
                      std::placeholders::_1))) {}
  ~HidlCallbackHandler() = default;

  bool addCallback(const sp<CallbackType>& cb) {
    // TODO(b/33818800): Can't compare proxies yet. So, use the cookie
    // (callback proxy's raw pointer) to track the death of individual clients.
    uint64_t cookie = reinterpret_cast<uint64_t>(cb.get());
    if (cb_set_.find(cb) != cb_set_.end()) {
      LOG(WARNING) << "Duplicate death notification registration";
      return true;
    }
    if (!cb->linkToDeath(death_handler_, cookie)) {
      LOG(ERROR) << "Failed to register death notification";
      return false;
    }
    cb_set_.insert(cb);
    return true;
  }

  const std::set<android::sp<CallbackType>>& getCallbacks() { return cb_set_; }

  // Death notification for callbacks.
  void onObjectDeath(uint64_t cookie) {
    CallbackType* cb = reinterpret_cast<CallbackType*>(cookie);
    const auto& iter = cb_set_.find(cb);
    if (iter == cb_set_.end()) {
      LOG(ERROR) << "Unknown callback death notification received";
      return;
    }
    cb_set_.erase(iter);
    LOG(DEBUG) << "Dead callback removed from list";
  }

  void invalidate() {
    for (const sp<CallbackType>& cb : cb_set_) {
      if (!cb->unlinkToDeath(death_handler_)) {
        LOG(ERROR) << "Failed to deregister death notification";
      }
    }
    cb_set_.clear();
  }

 private:
  std::set<sp<CallbackType>> cb_set_;
  sp<HidlDeathHandler<CallbackType>> death_handler_;

  DISALLOW_COPY_AND_ASSIGN(HidlCallbackHandler);
};

}  // namespace hidl_callback_util
}  // namespace implementation
}  // namespace V1_0
}  // namespace wifi
}  // namespace hardware
}  // namespace android
#endif  // HIDL_CALLBACK_UTIL_H_