#include <sstream>

#include <pybind11/pybind11.h>
#include <mecab.h>

#include "utils.h"

namespace py = pybind11;

class Iterator {
private:
  const MeCab::Node *cursor;

public:
  explicit Iterator(const MeCab::Node* cursor) : cursor(cursor) {}

  const MeCab::Node& operator*() const { return *cursor; }
  Iterator& operator++() { cursor = cursor->next; return *this; }
  bool operator==(const Iterator& rhs) const { return cursor == rhs.cursor;}
};

void initialize_lattice(py::module &m) {
  // Parameters for MeCab::Lattice::request_type
  m.attr("MECAB_ONE_BEST") = 1;
  m.attr("MECAB_NBEST") = 2;
  m.attr("MECAB_PARTIAL") = 4;
  m.attr("MECAB_MARGINAL_PROB") = 8;
  m.attr("MECAB_ALTERNATIVE") = 16;
  m.attr("MECAB_ALL_MORPHS") = 32;
  m.attr("MECAB_ALLOCATE_SENTENCE") = 64;

  // Parameters for MeCab::Lattice::boundary_constraint_type
  m.attr("MECAB_ANY_BOUNDARY") = 0;
  m.attr("MECAB_TOKEN_BOUNDARY") = 1;
  m.attr("MECAB_INSIDE_TOKEN") = 2;

  // Reference: https://taku910.github.io/mecab/doxygen/classMeCab_1_1Lattice.html
  py::class_<MeCab::Lattice>(m, "Lattice")
    .def(py::init([]() {
      return std::unique_ptr<MeCab::Lattice>(MeCab::Lattice::create());
    }))
    .def("clear", &MeCab::Lattice::clear)
    .def("is_available", &MeCab::Lattice::is_available)
    .def("bos_node", &MeCab::Lattice::bos_node, py::return_value_policy::reference)
    .def("eos_node", &MeCab::Lattice::eos_node, py::return_value_policy::reference)
    .def("begin_nodes", py::overload_cast<size_t>(&MeCab::Lattice::begin_nodes, py::const_), py::return_value_policy::reference)
    .def("end_nodes", py::overload_cast<size_t>(&MeCab::Lattice::end_nodes, py::const_), py::return_value_policy::reference)
    .def("sentence", &MeCab::Lattice::sentence)
    .def("set_sentence", py::overload_cast<const char *>(&MeCab::Lattice::set_sentence))
    .def("set_sentence", py::overload_cast<const char *, size_t>(&MeCab::Lattice::set_sentence))
    .def("size", &MeCab::Lattice::size)
    .def("set_Z", &MeCab::Lattice::set_Z)
    .def("Z", &MeCab::Lattice::Z)
    .def("set_theta", &MeCab::Lattice::set_theta)
    .def("theta", &MeCab::Lattice::theta)
    .def("next", &MeCab::Lattice::next)
    .def("request_type", &MeCab::Lattice::request_type)
    .def("has_request_type", &MeCab::Lattice::has_request_type)
    .def("set_request_type", &MeCab::Lattice::set_request_type)
    .def("add_request_type", &MeCab::Lattice::add_request_type)
    .def("remove_request_type", &MeCab::Lattice::remove_request_type)
    .def("new_node", &MeCab::Lattice::newNode, py::return_value_policy::reference)
    .def("to_string", py::overload_cast<>(&MeCab::Lattice::toString), py::return_value_policy::copy)
    .def("to_string", py::overload_cast<const MeCab::Node *>(&MeCab::Lattice::toString), py::return_value_policy::copy)
    .def("to_string", py::overload_cast<char *, size_t>(&MeCab::Lattice::toString))
    .def("to_string", py::overload_cast<const MeCab::Node *, char *, size_t>(&MeCab::Lattice::toString))
    .def("enum_nbest_as_string", py::overload_cast<size_t>(&MeCab::Lattice::enumNBestAsString), py::return_value_policy::copy)
    .def("enum_nbest_as_string", py::overload_cast<size_t, char *, size_t>(&MeCab::Lattice::enumNBestAsString))
    .def("has_constraint", &MeCab::Lattice::has_constraint)
    .def("boundary_constraint", &MeCab::Lattice::boundary_constraint)
    .def("feature_constraint", &MeCab::Lattice::feature_constraint)
    .def("set_boundary_constraint", &MeCab::Lattice::set_boundary_constraint)
    .def("set_feature_constraint", &MeCab::Lattice::set_feature_constraint)
    .def("set_result", &MeCab::Lattice::set_result)
    .def("what", &MeCab::Lattice::what, py::return_value_policy::copy)
    .def("set_what", &MeCab::Lattice::set_what)
    .def("__len__",  &MeCab::Lattice::size)
    .def("__iter__", [](const MeCab::Lattice &self) {
      Iterator begin = Iterator(self.bos_node()->next);
      Iterator end = Iterator(self.eos_node());
      return py::make_iterator(begin, end);
    }, py::keep_alive<0, 1>()) // 0 -> iterator, 1 -> self
    .def("__repr__", [](const MeCab::Lattice &self) {
      std::stringstream stream;
      stream << "Lattice(";
      stream << "text=\"" << escape(self.sentence()) << "\"";
      stream << ")";
      return stream.str();
    })
  ;
}