// Copyright (C) 2010, Guy Barrand. All rights reserved.
// See the file tools.license for terms.

#ifndef tools_wroot_branch
#define tools_wroot_branch

#include "leaf"
#include "basket"
#include "../cmemT"
#include "../forit"

#include "imutex"

namespace tools {
namespace wroot {

class branch : public virtual ibo {
public: //ibo
  virtual const std::string& store_cls() const {
    static const std::string s_v("TBranch");
    return s_v;
  }
  virtual bool stream(buffer& a_buffer) const {
    unsigned int c;
    if(!a_buffer.write_version(8,c)) return false;
    if(!Named_stream(a_buffer,m_name,m_title)) return false;

    if(!AttFill_stream(a_buffer)) return false;

    int fEntryOffsetLen = 1000;
    int fOffset = 0;
    int fSplitLevel = 0;

    if(!a_buffer.write(fCompress)) return false;
    if(!a_buffer.write(m_basket_size)) return false;
    if(!a_buffer.write(fEntryOffsetLen)) return false;
    if(!a_buffer.write(m_write_basket)) return false;
    int fEntryNumber = (int)m_entry_number;
    if(!a_buffer.write(fEntryNumber)) return false;
    if(!a_buffer.write(fOffset)) return false;
    if(!a_buffer.write(m_max_baskets)) return false;
    if(!a_buffer.write(fSplitLevel)) return false;
    double fEntries = (double)m_entries;
    if(!a_buffer.write(fEntries)) return false;
    double fTotBytes = (double)m_tot_bytes;
    double fZipBytes = (double)m_zip_bytes;
    if(!a_buffer.write(fTotBytes)) return false;
    if(!a_buffer.write(fZipBytes)) return false;

    if(!m_branches.stream(a_buffer)) return false;
    if(!m_leaves.stream(a_buffer)) return false;
    if(!m_baskets.stream(a_buffer)) return false;

    // See TStreamerInfo::ReadBuffer::WriteBasicPointer
    if(!a_buffer.write((char)1)) return false;
    if(!a_buffer.write_fast_array(fBasketBytes,m_max_baskets)) return false;
    if(!a_buffer.write((char)1)) return false;
    if(!a_buffer.write_fast_array(fBasketEntry,m_max_baskets)) return false;

    char isBigFile = 1;
    //GB : begin
   {for(uint32 i=0;i<m_max_baskets;i++) {
      if(fBasketSeek[i]>START_BIG_FILE()) {
        isBigFile = 2;
        break;
      }
    }}
    //GB : end

    if(!a_buffer.write(isBigFile)) return false;
    if(isBigFile==2) {
      if(!a_buffer.write_fast_array(fBasketSeek,m_max_baskets)) return false;
    } else {
      for(uint32 i=0;i<m_max_baskets;i++) {
        if(fBasketSeek[i]>START_BIG_FILE()) { //G.Barrand : add this test.
          m_out << "tools::wroot::branch::stream :"
                << " attempt to write big Seek "
                << fBasketSeek[i] << " on 32 bits."
                << std::endl;
          return false;
        }

        if(!a_buffer.write((seek32)fBasketSeek[i])) return false;
      }
    }

    // fFileName
    if(!a_buffer.write(std::string(""))) return false;

    if(!a_buffer.set_byte_count(c)) return false;
    return true;
  }

public:
  branch(std::ostream& a_out,bool a_byte_swap,uint32 a_compression,
         seek a_seek_directory,const std::string& a_name,const std::string& a_title,bool a_verbose)
  :m_out(a_out)
  ,m_byte_swap(a_byte_swap)
  ,m_verbose(a_verbose)
  ,m_seek_directory(a_seek_directory)
  ,m_name(a_name)
  ,m_title(a_title)
  ,fAutoDelete(false)
  ,fCompress(a_compression)
  ,m_basket_size(32000)
  ,m_write_basket(0)
  ,m_entry_number(0)
  ,m_entries(0)
  ,m_tot_bytes(0)
  ,m_zip_bytes(0)
  ,m_max_baskets(10)
  ,fBasketBytes(0)
  ,fBasketEntry(0)
  ,fBasketSeek(0)
  {
    m_baskets.resize(m_max_baskets,0);
    fBasketBytes = new uint32[m_max_baskets];
    fBasketEntry = new uint32[m_max_baskets];
    fBasketSeek = new seek[m_max_baskets];
   {for(uint32 i=0;i<m_max_baskets;i++) {
      m_baskets[i] = 0;
      fBasketBytes[i] = 0;
      fBasketEntry[i] = 0;
      fBasketSeek[i] = 0;
    }}
    m_baskets[m_write_basket] = new basket(m_out,m_byte_swap,a_seek_directory,m_name,m_title,"TBasket",m_basket_size,m_verbose);
    fBasketEntry[m_write_basket] = (uint32)m_entry_number;
  }
  virtual ~branch(){
    delete [] fBasketBytes;
    delete [] fBasketEntry;
    delete [] fBasketSeek;
    fBasketBytes = 0;
    fBasketEntry = 0;
    fBasketSeek = 0;
  }
protected:
  branch(const branch& a_from)
  :ibo(a_from)
  ,m_out(a_from.m_out)
  ,m_seek_directory(a_from.m_seek_directory)
  {}
  branch& operator=(const branch&){return *this;}
public:
  const std::string& name() const {return m_name;}
  const std::string& title() const {return m_title;}

  uint64 entries() const {return m_entries;}

  uint64 tot_bytes() const {return m_tot_bytes;}
  void set_tot_bytes(uint64 a_value) {m_tot_bytes = a_value;}

  uint64 zip_bytes() const {return m_zip_bytes;}
  void set_zip_bytes(uint64 a_value) {m_zip_bytes = a_value;}

  void set_basket_size(uint32 a_size) {m_basket_size = a_size;}
  uint32 basket_size() const {return m_basket_size;}

  ////////////////////////////////////////////////
  ////////////////////////////////////////////////
  ////////////////////////////////////////////////
  template <class T>
  leaf_ref<T>* create_leaf_ref(const std::string& a_name,const T& a_ref){
    leaf_ref<T>* lf = new leaf_ref<T>(m_out,a_name,a_ref);
    m_leaves.push_back(lf);
    return lf;
  }

  leaf_string_ref* create_leaf_string_ref(const std::string& a_name,const std::string& a_ref){
    leaf_string_ref* lf = new leaf_string_ref(m_out,a_name,a_ref);
    m_leaves.push_back(lf);
    return lf;
  }
  ////////////////////////////////////////////////
  ////////////////////////////////////////////////
  ////////////////////////////////////////////////

  template <class T>
  leaf<T>* create_leaf(const std::string& a_name){
    leaf<T>* lf = new leaf<T>(m_out,a_name);
    m_leaves.push_back(lf);
    return lf;
  }

  leaf_string* create_leaf_string(const std::string& a_name){
    leaf_string* lf = new leaf_string(m_out,a_name);
    m_leaves.push_back(lf);
    return lf;
  }

  leaf_element* create_leaf_element(const std::string& a_name,int a_id,int a_type){
    leaf_element* lf = new leaf_element(m_out,a_name,a_id,a_type);
    m_leaves.push_back(lf);
    return lf;
  }

  // for row_wise ntuple :
  template <class T>
  leaf_std_vector_ref<T>* create_leaf_std_vector_ref(const std::string& a_name,
                                                     base_leaf& a_leaf_count,const std::vector<T>& a_ref) {
    leaf_std_vector_ref<T>* lf = new leaf_std_vector_ref<T>(m_out,a_name,a_leaf_count,a_ref);
    m_leaves.push_back(lf);
    return lf;
  }

  const std::vector<base_leaf*>& leaves() const {return m_leaves;}

  void reset() {
    m_baskets.clear_objs();
    delete [] fBasketBytes;
    delete [] fBasketEntry;
    delete [] fBasketSeek;
    fBasketBytes = 0;
    fBasketEntry = 0;
    fBasketSeek = 0;

    m_max_baskets = 10;

    m_write_basket = 0;
    m_entry_number = 0;
    m_entries = 0;
    m_tot_bytes = 0;
    m_zip_bytes = 0;

    m_baskets.resize(m_max_baskets,0);
    fBasketBytes = new uint32[m_max_baskets];
    fBasketEntry = new uint32[m_max_baskets];
    fBasketSeek = new seek[m_max_baskets];
   {for(uint32 i=0;i<m_max_baskets;i++) {
      m_baskets[i] = 0;
      fBasketBytes[i] = 0;
      fBasketEntry[i] = 0;
      fBasketSeek[i] = 0;
    }}
    fBasketEntry[m_write_basket] = (uint32)m_entry_number;
    m_baskets[m_write_basket] = new basket(m_out,m_byte_swap,m_seek_directory,m_name,m_title,"TBasket",m_basket_size,m_verbose);

   {tools_vforit(branch*,m_branches,it) (*it)->reset();}
  }

  bool fill(ifile& a_file,uint32& a_nbytes,uint32& a_add_bytes,uint32& a_nout) {
    a_nbytes = 0;
    a_add_bytes = 0;
    a_nout = 0;

    if(m_write_basket>=m_max_baskets) {
      m_out << "tools::wroot::branch::fill :"
            << " potential overflow : m_write_basket (" << m_write_basket << ")"
	    << " >= m_max_baskets (" << m_max_baskets << ")."
            << std::endl;
      return false;
    }

    basket* bk = m_baskets[m_write_basket];
    if(!bk) {
      m_out << "tools::wroot::branch::fill :"
            << " m_baskets[m_write_basket] should not be null."
            << std::endl;
      return false;
    }

    buffer& buf = bk->datbuf();

    buf.reset_objs_map();

    uint32 lold = buf.length();

    bk->update(bk->key_length()+lold);
    m_entries++;
    m_entry_number++;

    if(!fill_leaves(buf)) {
      m_out << "tools::wroot::branch::fill :"
            << " fill_leaves() failed."
            << std::endl;
      return false;
    }

    uint32 lnew   = buf.length();
    uint32 nbytes = lnew - lold;
    uint32 nsize  = 0;

    uint32 add_bytes  = 0;
    uint32 nout  = 0;

    // Should we create a new basket?
    // Compare expected next size with m_basket_size.
    if((lnew+2*nsize+nbytes)>=m_basket_size) {
      if(!bk->write_on_file(a_file,m_write_basket,nout)) {//it set bk m_seek_key, set a_file.m_END.
                                                          //Does a a_file.set_END().
        m_out << "tools::wroot::branch::fill :"
              << " basket.write_on_file() failed."
              << std::endl;
        return false;
      }
      fBasketBytes[m_write_basket]  = bk->number_of_bytes();
    //fBasketEntry[m_write_basket]  = // not done here since we do not know here the first entry.
      fBasketSeek[m_write_basket]   = bk->seek_key();

      add_bytes = bk->object_size() + bk->key_length();

      delete bk;
      m_baskets[m_write_basket] = 0;

      m_tot_bytes += add_bytes;
      m_zip_bytes += nout;

      m_write_basket++;
      if(!check_alloc_fBasketXxx()) return false;

      m_baskets[m_write_basket] = new basket(m_out,m_byte_swap,
                                             m_seek_directory,m_name,m_title,"TBasket",m_basket_size,
                                             m_verbose);
      fBasketEntry[m_write_basket] = (uint32)m_entry_number;
    }
    a_nbytes = nbytes;
    a_add_bytes = add_bytes;
    a_nout = nout;
    return true;
  }

  ///////////////////////////////////////////////////////////////////////////
  /// for parallelization : /////////////////////////////////////////////////
  ///////////////////////////////////////////////////////////////////////////
  bool add_basket(ifile& a_file,basket& a_basket,uint32& a_add_bytes,uint32& a_nout) {
    //done on a (main) ntuple/branch.

    if(m_write_basket>=m_max_baskets) {
      m_out << "tools::wroot::branch::add_basket :"
            << " potential overflow : m_write_basket (" << m_write_basket << ")"
	    << " >= m_max_baskets (" << m_max_baskets << ")."
            << std::endl;
      return false;
    }

    uint32 nout;
    if(!a_basket.write_on_file(a_file,m_write_basket,nout)) { //a_file is the main file (a true file).
      //the upper set a_basket.m_seek_key, set a_file.m_END. Does a a_file.set_END().
      m_out << "tools::wroot::branch::add_basket :"
            << " basket.write_on_file() failed."
            << std::endl;
      return false;
    }

    fBasketBytes[m_write_basket]  = a_basket.number_of_bytes();
    fBasketEntry[m_write_basket]  = (uint32)m_entry_number;
    fBasketSeek[m_write_basket]   = a_basket.seek_key();

    m_entries += a_basket.nev();
    m_entry_number += a_basket.nev();

    uint32 add_bytes = a_basket.object_size() + a_basket.key_length();

    delete m_baskets[m_write_basket];
    m_baskets[m_write_basket] = 0;

    m_write_basket++;
    if(!check_alloc_fBasketXxx()) return false;

    m_baskets[m_write_basket] = new basket(m_out,m_byte_swap,
                                           m_seek_directory,m_name,m_title,"TBasket",m_basket_size,
                                           m_verbose);
    fBasketEntry[m_write_basket] = (uint32)m_entry_number;

    a_add_bytes = add_bytes;
    a_nout = nout;
    return true;
  }

  class iadd_basket {
  public:
    virtual ~iadd_basket() {}
  public:
    virtual bool add_basket(basket*) = 0;
  };

  bool pfill(iadd_basket& a_badd,uint32 a_nev) {
    //done on a pntuple/branch.

    //a_nbytes = 0;
    //a_add_bytes = 0;
    //a_nout = 0;

    // this method uses always m_baskets[0].

    basket* bk = m_baskets[m_write_basket];
    if(!bk) {
      m_out << "tools::wroot::branch::parallel_fill :"
            << " get_basket failed."
            << std::endl;
      return false;
    }

    buffer& buf = bk->datbuf();

    uint32 lold = buf.length();

    bk->update(bk->key_length()+lold);

    //m_entries++;      //not used by parallel branches.
    //m_entry_number++; //idem.

    if(!fill_leaves(buf)) {
      m_out << "tools::wroot::branch::parallel_fill :"
            << " fill_leaves() failed."
            << std::endl;
      return false;
    }

    uint32 lnew   = buf.length();
    uint32 nbytes = lnew - lold;

    // Should we create a new basket?
    bool store_basket = false;
    if(a_nev) {
      store_basket = (bk->nev()>=a_nev);              // trigger add_basket per number of entries reached.
    } else {
      store_basket = ((lnew+nbytes)>=m_basket_size);  // trigger add_basket when reaching m_basket_size.
    }

    if(store_basket) {
      if(!a_badd.add_basket(bk)) {  //we give ownership of bk.
        m_out << "tools::wroot::branch::parallel_fill :"
              << " main_branch.add_basket() failed."
              << std::endl;
        return false;
      }
      //delete bk; //no, it is deleted by the a_badd.add_basket().
      bool main_branch_byte_swap = m_byte_swap; //NOTE : it assumes that pntuple/branch and main/branch have same byte_swap.
      bool main_branch_verbose = m_verbose;
      m_baskets[m_write_basket] = new basket(m_out,main_branch_byte_swap,
                                             m_seek_directory,m_name,m_title,"TBasket",m_basket_size,
                                             main_branch_verbose);
    }

    return true;
  }

  bool end_pfill(iadd_basket& a_badd) {
    //done on a pntuple/branch.

    basket* bk = m_baskets[m_write_basket];
    if(!bk) {
      m_out << "tools::wroot::branch::end_pfill :"
            << " m_baskets[m_write_basket] should not be null."
            << std::endl;
      return false;
    }

    buffer& buf = bk->datbuf();
    uint32 lold = buf.length();

    if(lold) {
      if(!a_badd.add_basket(bk)) {  //we give ownership of bk.
        m_out << "tools::wroot::branch::parallel_fill :"
              << " main_branch.add_basket() failed."
              << std::endl;
        return false;
      }
    } else {
      delete bk;
    }

    m_baskets[m_write_basket] = 0; // no need to recreate a basket since it is end of fill.

    return true;
  }

  ///////////////////////////////////////////////////////////////////////////
  ///////////////////////////////////////////////////////////////////////////
  ///////////////////////////////////////////////////////////////////////////
protected:
  bool check_alloc_fBasketXxx() {
    if(m_write_basket>=m_max_baskets) {
      //Increase BasketEntry buffer of a minimum of 10 locations
      // and a maximum of 50 per cent of current size
      uint32 newsize = max_of<uint32>(10,uint32(1.5*m_max_baskets));
      if(newsize>=START_BIG_FILE()) {
        //we are going to have pb with uint32[] indexing.
        m_out << "tools::wroot::branch::add_basket :"
              << " new size for fBasket[Bytes,Entry,Seek] arrays"
              << " is too close of 32 bits limit."
              << std::endl;
        m_out << "tools::wroot::branch::add_basket :"
              << " you have to work with larger basket size."
              << std::endl;
        return false;
      }

      m_baskets.resize(newsize,0);

      if(!cmem_realloc<uint32>(fBasketBytes,newsize,m_max_baskets,true)) {
        m_out << "tools::wroot::branch::add_basket : realloc failed." << std::endl;
        return false;
      }
      if(!cmem_realloc<uint32>(fBasketEntry,newsize,m_max_baskets,true)){
        m_out << "tools::wroot::branch::add_basket : realloc failed." << std::endl;
        return false;
      }
      if(!cmem_realloc<seek>(fBasketSeek,newsize,m_max_baskets,true)){
        m_out << "tools::wroot::branch::add_basket : realloc failed." << std::endl;
        return false;
      }
      m_max_baskets = newsize;
    }

    m_baskets[m_write_basket] = 0;
    fBasketBytes[m_write_basket] = 0;
    fBasketEntry[m_write_basket] = 0;
    fBasketSeek[m_write_basket]  = 0;

    return true;
  }

  virtual bool fill_leaves(buffer& a_buffer) { //virtual for branch_element.
    tools_vforit(base_leaf*,m_leaves,it) {
      if(!(*it)->fill_buffer(a_buffer)) return false;
    }
    return true;
  }
protected:
  std::ostream& m_out;
  bool m_byte_swap;
  bool m_verbose;
  seek m_seek_directory;
  obj_array<basket> m_baskets;
  /// for parallelization :
public:
  std::vector<basket*> m_parallel_baskets;
protected:
  //Object
  //uint32 m_bits;
  //Named
  std::string m_name;
  std::string m_title;

  bool fAutoDelete;
private: //be sure that sub branches are not used for the moment.
  obj_array<branch> m_branches;
protected:
  obj_array<base_leaf> m_leaves;
  uint32 fCompress;      //  Compression level and algorithm
  uint32 m_basket_size;  //  Initial Size of Basket Buffer
  uint32 m_write_basket; //  Last basket number written
  uint64 m_entry_number; // Current entry number (last one filled in this branch)
  uint64 m_entries;      //  Number of entries
  uint64 m_tot_bytes;
  uint64 m_zip_bytes;
  uint32 m_max_baskets;
  uint32* fBasketBytes;  //[m_max_baskets] Length of baskets on file
  uint32* fBasketEntry;  //[m_max_baskets] Table of first entry in eack basket
  seek* fBasketSeek;     //[m_max_baskets] Addresses of baskets on file
};

}}

#endif
