/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you 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.
 */

/*!
 * \file micro_module.cc
 */

#include <tvm/runtime/registry.h>
#include <tvm/runtime/c_runtime_api.h>
#include <tvm/runtime/module.h>
#include <unordered_map>
#include <string>
#include "micro_session.h"
#include "low_level_device.h"
#include "micro_common.h"
#include "../pack_args.h"

namespace air {
namespace runtime {
/*!
 * \brief module for uTVM micro devices
 */
class MicroModuleNode final : public ModuleNode {
 public:
  MicroModuleNode() {}

  ~MicroModuleNode() {}

  const char* type_key() const final {
    return "micro";
  }

  PackedFunc GetFunction(const std::string& name,
                         const ObjectPtr<Object>& sptr_to_self) final;

  /*!
   * \brief initializes module by establishing device connection and loads binary
   * \param binary_path path of the binary to be loaded
   */
  void InitMicroModule(const std::string& binary_path) {
    session_ = MicroSession::Current();
    binary_path_ = binary_path;
    binary_info_ = session_->LoadBinary(binary_path_);
  }

  /*!
   * \brief runs selected function on the micro device
   * \param func_name name of the function to be run
   * \param func_offset offset of the function to be run
   * \param args type-erased arguments passed to the function
   */
  void RunFunction(const std::string& func_name, DevBaseOffset func_offset, const TVMArgs& args) {
    session_->PushToExecQueue(func_offset, args);
  }

 private:
  /*! \brief module binary info */
  BinaryInfo binary_info_;
  /*! \brief path to module binary */
  std::string binary_path_;
  /*! \brief global session pointer */
  ObjectPtr<MicroSession> session_;
};

class MicroWrappedFunc {
 public:
  MicroWrappedFunc(MicroModuleNode* m,
                   ObjectPtr<MicroSession> session,
                   const std::string& func_name,
                   DevBaseOffset func_offset) {
    m_ = m;
    session_ = session;
    func_name_ = func_name;
    func_offset_ = func_offset;
  }

  void operator()(TVMArgs args, TVMRetValue* rv) const {
    m_->RunFunction(func_name_, func_offset_, args);
  }

 private:
  /*! \brief internal module */
  MicroModuleNode* m_;
  /*! \brief reference to the session for this function (to keep the session alive) */
  ObjectPtr<MicroSession> session_;
  /*! \brief name of the function */
  std::string func_name_;
  /*! \brief offset of the function to be called */
  DevBaseOffset func_offset_;
};

PackedFunc MicroModuleNode::GetFunction(
    const std::string& name,
    const ObjectPtr<Object>& sptr_to_self) {
  DevBaseOffset func_offset =
      session_->low_level_device()->ToDevOffset(binary_info_.symbol_map[name]);
  MicroWrappedFunc f(this, session_, name, func_offset);
  return PackedFunc(f);
}

// register loadfile function to load module from Python frontend
TVM_REGISTER_GLOBAL("module.loadfile_micro_dev")
.set_body([](TVMArgs args, TVMRetValue* rv) {
    auto n = make_object<MicroModuleNode>();
    n->InitMicroModule(args[0]);
    *rv = runtime::Module(n);
  });
}  // namespace runtime
}  // namespace air
