#ifndef __JTREE_BRANCH_ADDRESSES_H__
#define __JTREE_BRANCH_ADDRESSES_H__

#include <iostream>

#include "JLang/JException.hh"
#include "JROOT/JRootAddress.hh"
#include "JROOT/JRootClass.hh"
#include "Jeep/JPrint.hh"
#include "Jeep/JStreamToolkit.hh"
#include "TDataMember.h"
#include "TTree.h"

namespace JROOT {

/* to_vector converts a list of Root TObjects into a list
 * of those objects names
 */
inline std::vector<std::string> to_vector(TCollection* col)
{
  std::vector<std::string> names;

  TIter    next(col);
  TObject* object{nullptr};
  while ((object = next()) != nullptr) {
    names.emplace_back(object->GetName());
  }
  return names;
}

using JLANG::JException;

/**
 * Class responsible to setup TTree branch addresses
 * for reading objects of type T.
 *
 */
template <typename T>
class JTreeBranchAddresses : public JRootAddress<T> {

 public:
  /*
   * Set the branch(es) address(es) to read objects of type T.
   *
   * If the tree contains a single branch we assume that branch contains
   * objects of type T.
   *
   * If the tree contains several branches we check whether those branches
   * correspond to the full list of T data members. If that's the case
   * the tree is a flat (one branch per data member) representation of T
   * and setup the branches accordingly. If not, we throw an exception.
   */
  void setBranchAddress(TTree& tree)
  {
    TObjArray* branches = tree.GetListOfBranches();
    // single branch
    if (branches->GetEntries() == 1) {
      tree.SetBranchAddress(branches->First()->GetName(), &this->address);
      return;
    }

    // if more than one branch, check those are the data members of T
    TClass* t_class = TClass::GetClass<T>();
    if (t_class == nullptr) {
      THROW(JException, "Could not get class " << typeid(T).name());
    }

    auto branch_names = to_vector(tree.GetListOfBranches());
    auto all_member_names = to_vector(t_class->GetListOfRealData());

    std::vector<std::string> member_names;

    // remove TObject specific members
    std::copy_if(all_member_names.begin(), all_member_names.end(), std::back_inserter(member_names),
                 [](const std::string& member_name)
                 { return !JRootClass::is_tobject_member(member_name.c_str()); });

    std::sort(branch_names.begin(), branch_names.end());
    std::sort(member_names.begin(), member_names.end());

    // all members must have a corresponding branch
    // (but more branches are allowed, we just don't handle them here)
    if (std::includes(branch_names.begin(), branch_names.end(), member_names.begin(),
                      member_names.end())) {

      this->address = new T;

      auto* base = reinterpret_cast<uint8_t*>(this->address);

      for (const auto& member_name : member_names) {
        auto* member = static_cast<TDataMember*>(t_class->GetDataMember(member_name.c_str()));
        tree.SetBranchAddress(member_name.c_str(),
                              reinterpret_cast<void*>(base + member->GetOffset()));
      }
      return;
    }

    // if we cannot configure the tree, try to give some relevant information
    // to the user who will have to deal with the exception...
    THROW(JException, "Members of type " << typeid(T).name() << "\n"
	                                 << JEEP::JEEPZ() << member_names << "\nBranches in tree\n"
	                                 << JEEP::JEEPZ() << branch_names);
  }
};
}  // namespace JROOT

#endif