#pragma once

#include <cstdlib>
#include <initializer_list>
#include <string>
#include <utility>

#include "R_ext/Arith.h"
#include "cpp4r/R.hpp"
#include "cpp4r/attribute_proxy.hpp"
#include "cpp4r/list.hpp"
#include "cpp4r/r_vector.hpp"

namespace cpp4r {

class named_arg;
namespace writable {
class data_frame;
}

class data_frame : public list {
  using list::list;
  friend class writable::data_frame;

  static R_xlen_t calc_nrow(SEXP x) noexcept {
    auto nms = Rf_getAttrib(x, R_RowNamesSymbol);
    bool has_short_rownames =
        (Rf_isInteger(nms) && Rf_xlength(nms) == 2 && INTEGER(nms)[0] == NA_INTEGER);
    if (has_short_rownames) return static_cast<R_xlen_t>(abs(INTEGER(nms)[1]));
    if (!Rf_isNull(nms)) return Rf_xlength(nms);
    if (Rf_xlength(x) == 0) return 0;
    return Rf_xlength(VECTOR_ELT(x, 0));
  }

 public:
  R_xlen_t nrow() const noexcept { return calc_nrow(*this); }
  R_xlen_t ncol() const noexcept { return size(); }
};

namespace writable {

class data_frame : public cpp4r::data_frame {
  writable::list set_data_frame_attributes(writable::list&& x) {
    return set_data_frame_attributes(std::move(x), calc_nrow(x));
  }

  writable::list set_data_frame_attributes(writable::list&& x, R_xlen_t nrow) {
    x.attr(R_RowNamesSymbol) = {NA_INTEGER, -static_cast<int>(nrow)};
    x.attr(R_ClassSymbol) = "data.frame";
    return std::move(x);
  }

 public:
  data_frame(const SEXP data) : cpp4r::data_frame(set_data_frame_attributes(data)) {}
  data_frame(const SEXP data, bool is_altrep)
      : cpp4r::data_frame(set_data_frame_attributes(data), is_altrep) {}
  data_frame(const SEXP data, bool is_altrep, R_xlen_t nrow)
      : cpp4r::data_frame(set_data_frame_attributes(data, nrow), is_altrep) {}
  data_frame(std::initializer_list<list> il)
      : cpp4r::data_frame(set_data_frame_attributes(writable::list(il))) {}
  data_frame(std::initializer_list<named_arg> il)
      : cpp4r::data_frame(set_data_frame_attributes(writable::list(il))) {}

  using cpp4r::data_frame::ncol;
  using cpp4r::data_frame::nrow;

  attribute_proxy<data_frame> attr(const char* name) const { return {*this, name}; }
  attribute_proxy<data_frame> attr(const std::string& name) const {
    return {*this, name.c_str()};
  }
  attribute_proxy<data_frame> attr(SEXP name) const { return {*this, name}; }
  attribute_proxy<data_frame> names() const { return {*this, R_NamesSymbol}; }
};

}  // namespace writable

}  // namespace cpp4r
