#include <leptonica/allheaders.h>
#include <pybind11/pybind11.h>
#include <pybind11/stl.h>
#include <tesseract/baseapi.h>
#include <tesseract/genericvector.h>
#include <sstream>

namespace py = pybind11;
using tesseract::OcrEngineMode;
using tesseract::PageIteratorLevel;
using tesseract::PageSegMode;
using tesseract::ResultIterator;
using tesseract::TessBaseAPI;

PYBIND11_MODULE(_pysseract, m) {
    m.doc() = R"pbdoc(
        pysseract
        -----------------------
        .. currentmodule:: pysseract
        .. autosummary::
        
           :toctree: _generate
           Pysseract
           Box
           ResultIterator
           PageIteratorLevel
           PageSegMode
           OcrEngineMode
           apiVersion
           availableLanguages
           defaultDataPath

    )pbdoc";

    m.def("apiVersion", &tesseract::TessBaseAPI::Version, "Tesseract API version as seen in the library");
    m.def("availableLanguages",
          []() {
              tesseract::TessBaseAPI api;
              api.Init(nullptr, nullptr);
              GenericVector<STRING> glangs;
              api.GetAvailableLanguagesAsVector(&glangs);
              std::vector<std::string> langs(glangs.size());
              for (int i = 0; i < glangs.size(); i++) {
                  langs[i] = glangs[i].c_str();
              }
              return langs;
          },
          "return a list of available languages from TESSDATA_PREFIX");
    m.def("defaultDataPath",
          []() {
              tesseract::TessBaseAPI api;
              api.Init(nullptr, nullptr);
              return api.GetDatapath();
          },
          "return TESSDATA_PREFIX");

    py::class_<TessBaseAPI>(m, "Pysseract", R"pbdoc(
        This is the main class for interacting with the Tesseract API
        
        .. code-block:: python

            import pysseract
            t = pysseract.Pysseract()
            t.SetImageFromPath('/Users/hxia/Downloads/001-helloworld.png')
            print(t.utf8Text)
            t.End()
        
        For advanced use cases, please consult ResultIterator.
    )pbdoc")
        .def(py::init([](const char *datapath, const char *language) {
                 TessBaseAPI *api = new (TessBaseAPI);
                 api->Init(datapath, language);
                 return std::unique_ptr<TessBaseAPI>(api);
             }),
             py::arg("datapath"), py::arg("language"))
        .def(py::init([](const char *datapath, const char *language, OcrEngineMode mode) {
                 TessBaseAPI *api = new (TessBaseAPI);
                 api->Init(datapath, language, mode);
                 return std::unique_ptr<TessBaseAPI>(api);
             }),
             py::arg("datapath"), py::arg("language"), py::arg("engineMode"))
        .def(py::init([](const char *datapath, const char *language, OcrEngineMode mode,
                         std::vector<std::string> configs, std::unordered_map<std::string, std::string> settings,
                         bool set_only_non_debug_params) {
                 TessBaseAPI *api = new (TessBaseAPI);

                 char *configs_[configs.size()];
                 for (size_t i = 0; i < configs.size(); i++) {
                     configs[i] = configs[i].c_str();
                 }

                 GenericVector<STRING> vars_vec;
                 GenericVector<STRING> vars_values;
                 for (auto &&entry : settings) {
                     vars_vec.push_back(STRING(entry.first.c_str()));
                     vars_values.push_back(STRING(entry.second.c_str()));
                 }

                 api->Init(datapath, language, mode, configs_, configs.size(), &vars_vec, &vars_values,
                           set_only_non_debug_params);
                 return std::unique_ptr<TessBaseAPI>(api);
             }),
             py::arg("datapath"), py::arg("language"), py::arg("engineMode"), py::arg("configsList"),
             py::arg("settingDict"), py::arg("setOnlyNonDebugParams"), R"pbdoc(
                 detailed initialization for the base API, two conversions have been done as follows
                 char **configs, int configs_size <==> configsList
                 const GenericVector<STRING> *vars_vec, const GenericVector<STRING> *vars_values <==> settingDict
                 )pbdoc")
        .def(py::init([]() {
            TessBaseAPI *api = new (TessBaseAPI);
            api->Init(nullptr, nullptr);
            return std::unique_ptr<TessBaseAPI>(api);
        }))
        .def("Clear", &TessBaseAPI::Clear, R"pbdoc(
        Free up recognition results and any stored image data, without actually
        freeing any recognition data that would be time-consuming to reload.
        Afterwards, you must call SetImage or TesseractRect before doing
        any Recognize or Get* operation.
    )pbdoc")
        .def("End", &TessBaseAPI::End,
             "Close down tesseract and free up all memory, after which the instance should not be reused.")
        .def("__enter__", [](TessBaseAPI &api) { return &api; }, "for use in `with` statement")
        .def("__exit__",
             [](TessBaseAPI &api, py::object exc_type, py::object exc_value, py::object traceback) { api.End(); },
             "for with statement, after exitting with, the instance is not reusable.")
        .def_property_readonly("dataPath", &TessBaseAPI::GetDatapath,
                               R"pbdoc(Read-only: Returns the path where Tesseract model objects are stored)pbdoc")
        .def_property(
            "pageSegMode", &TessBaseAPI::GetPageSegMode, &TessBaseAPI::SetPageSegMode,
            R"pbdoc(This attribute can be used to get or set the page segmentation mode used by the tesseract model)pbdoc")
        .def_property_readonly("utf8Text", &TessBaseAPI::GetUTF8Text,
                               R"pbdoc(Read-only: Return all identified text concatenated into a UTF-8 string)pbdoc")
        .def_property_readonly(
            "UNLVText", &TessBaseAPI::GetUNLVText,
            R"pbdoc(Read-only: Return all identified text according to UNLV format Latin-1 with specific reject and suspect codes)pbdoc")
        .def("GetHOCRText", (char *(TessBaseAPI::*)(int)) & TessBaseAPI::GetHOCRText, py::arg("pagenum"),
             "Make an HTML-formatted string with hOCR. 'pagenum' is 0-based, appears as 1-based in results.")
        .def("GetTSVText", &TessBaseAPI::GetTSVText, py::arg("pagenum"),
             "Make a TSV-formatted string from the internal data structures. 'pagenum' is 0-based, appears as 1-based "
             "in results.")
#if TESSERACT_VERSION >= (4 << 16 | 1 << 8)
        .def("GetAltoText", (char *(TessBaseAPI::*)(int)) & TessBaseAPI::GetAltoText, py::arg("pagenum"),
             "Make an ALTO XML string from internal data. 'pagenum' is 0-based, appears as 1-based in results.")
        .def("GetLSTMBoxText", &TessBaseAPI::GetLSTMBoxText, py::arg("pagenum"),
             "Make a box file for LSTM training from the internal data structures. 'pagenum' is 0-based, appears as "
             "1-based in results.")
        .def("GetWordStrBoxText", &TessBaseAPI::GetWordStrBoxText, py::arg("pagenum"),
             "Make a string formatted in the same style as Tesseract training data. 'pagenum' is 0-based, appears as "
             "1-based in results.")
#endif
        .def("GetOsdText", &TessBaseAPI::GetOsdText, py::arg("pagenum"),
             "Recognised text is returned as UTF-8. 'pagenum' is 0-based, appears as 1-based in results.")
        .def("GetIterator",
             [](TessBaseAPI &api) {
                 api.Recognize(nullptr);
                 return api.GetIterator();
             },
             "Returns the iterator over boxes found in a given source image")
        .def("GetVariableAsString",
             [](TessBaseAPI &api, const char *name) {
                 STRING str;
                 bool res = api.GetVariableAsString(name, &str);
                 if (res) {
                     return str.c_str();
                 }
                 return "";
             },
             "Get value of named variable as a string, if it exists.")
        .def("SetSourceResolution", &TessBaseAPI::SetSourceResolution, py::arg("ppi"),
             "Set the pixel-per-inch value for the source image")
        .def("SetImageFromPath",
             [](TessBaseAPI &api, const char *imgpath) {
                 Pix *image = pixRead(imgpath);
                 api.SetImage(image);
             },
             py::arg("imgpath"), "Read an image from a given fully-qualified file path")
        .def("SetImageFromBytes",
             [](TessBaseAPI &api, const std::string &bytes) {
                 Pix *image = pixReadMem((unsigned char *)bytes.data(), bytes.size());
                 api.SetImage(image);
             },
             py::arg("bytes"), "Read an image from a string of bytes")
        .def("SetVariable", &TessBaseAPI::SetVariable, py::arg("name"), py::arg("value"),
             "Note: Must be called after Init(). Only works for non-init variables.")
        .def("SetRectangle", &TessBaseAPI::SetRectangle, py::arg("left"), py::arg("top"), py::arg("width"),
             py::arg("height"), "Restrict recognition to a sub-rectangle of the image. Call after SetImage.");

    py::enum_<PageIteratorLevel>(m, "PageIteratorLevel", "Enumeration of page iteration level settings")
        .value("BLOCK", PageIteratorLevel::RIL_BLOCK, "ResultIterator at the block level")
        .value("PARA", PageIteratorLevel::RIL_PARA, "ResultIterator at the paragraph level")
        .value("TEXTLINE", PageIteratorLevel::RIL_TEXTLINE, "ResultIterator at the text line level")
        .value("WORD", PageIteratorLevel::RIL_WORD, "ResultIterator at the word level")
        .value("SYMBOL", PageIteratorLevel::RIL_SYMBOL, "ResultIterator at the symbol level");

    py::enum_<PageSegMode>(m, "PageSegMode", "Enumeration of page segmentation settings")
        .value("OSD_ONLY", PageSegMode::PSM_OSD_ONLY, "Segment the page in \"OSD only\" mode")
        .value("AUTO_OSD", PageSegMode::PSM_AUTO_OSD, "Segment the page in \"Auto OSD\" mode")
        .value("AUTO_ONLY", PageSegMode::PSM_AUTO_ONLY, "Segment the page in \"Automatic only\" mode")
        .value("AUTO", PageSegMode::PSM_AUTO, "Segment the page in \"Automatic\" mode")
        .value("SINGLE_COLUMN", PageSegMode::PSM_SINGLE_COLUMN, "Segment the page in \"Single column\" mode")
        .value("SINGLE_BLOCK_VERT_TEXT", PageSegMode::PSM_SINGLE_BLOCK_VERT_TEXT,
               "Segment the page in \"Single block of vertical text\" mode")
        .value("SINGLE_BLOCK", PageSegMode::PSM_SINGLE_BLOCK, "Segment the page in \"Single block\" mode")
        .value("SINGLE_LINE", PageSegMode::PSM_SINGLE_LINE, "Segment the page in \"Single line\" mode")
        .value("SINGLE_WORD", PageSegMode::PSM_SINGLE_WORD, "Segment the page in \"Single word\" mode")
        .value("CIRCLE_WORD", PageSegMode::PSM_CIRCLE_WORD, "Segment the page in \"Circle word\" mode")
        .value("SINGLE_CHAR", PageSegMode::PSM_SINGLE_CHAR, "Segment the page in \"Single character\" mode")
        .value("SPARSE_TEXT", PageSegMode::PSM_SPARSE_TEXT, "Segment the page in \"Sparse text\" mode")
        .value("SPARSE_TEXT_OSD", PageSegMode::PSM_SPARSE_TEXT_OSD, "Segment the page in \"Sparse text OSD\" mode")
        .value("RAW_LINE", PageSegMode::PSM_RAW_LINE, "Segment the page in \"Raw line\" mode")
        .value("COUNT", PageSegMode::PSM_COUNT, "Segment the page in \"Count\" mode");

    py::enum_<OcrEngineMode>(m, "OcrEngineMode", "Enumeration of Engine Mode")
        .value("TESSERACT_ONLY", OcrEngineMode::OEM_TESSERACT_ONLY, "Run Tesseract only - fastest; deprecated")
        .value("LSTM_ONLY", OcrEngineMode::OEM_LSTM_ONLY, "Run just the LSTM line recognizer")
        .value("TESSERACT_LSTM_COMBINED", OcrEngineMode::OEM_TESSERACT_LSTM_COMBINED,
               "Run the LSTM recognizer, but allow fallback to Tesseract when things get difficult. deprecated")
        .value("DEFAULT", OcrEngineMode::OEM_DEFAULT, R"pbdoc(
            Specify this mode when calling init_*(),
            to indicate that any of the above modes
            should be automatically inferred from the
            variables in the language-specific config,
            command-line configs, or if not specified
            in any of the above should be set to the
            default OEM_TESSERACT_ONLY.
            )pbdoc")
        .value("COUNT", OcrEngineMode::OEM_COUNT, "Number of OEMs");

    py::class_<Box>(m, "Box", R"pbdoc(The bounding box structure)pbdoc")
        .def_readonly("left", &Box::x, "Leftmost co-ordinate of the box")
        .def_readonly("top", &Box::y, "Topmost co-ordinate of the box")
        .def_readonly("width", &Box::w, "Box width")
        .def_readonly("height", &Box::h, "Box height")
        .def("__repr__",
             [](const Box &box) {
                 std::ostringstream os;
                 os << "class Box<left:" << box.x << ",top:" << box.y << ",width:" << box.w << ",height:" << box.h
                    << ">";
                 return os.str();
             })
        .def_property_readonly("valid",
                               [](const Box &box) {
                                   if (box.x < 0 || box.y <= 0 || box.w <= 0 || box.h <= 0) return false;
                                   return true;
                               },
                               "Returns a boolean indicating whether the dimensions of the box are valid");

    py::class_<ResultIterator>(m, "ResultIterator", R"pbdoc(
        Internal Iterator that yields result at chosen level. If you're familiar with C/C++ iterators, the methods of this class should look familiar.

        .. code-block:: python

            t: pysseract.Pysseract
            resIter = t.GetIterator()
            LEVEL = pysseract.PageIteratorLevel.TEXTLINE
            while True:
                box = resIter.BoundingBox(LEVEL)
                text = resIter.GetUTF8Text(LEVEL)
                if not resIter.Next(LEVEL):
                    break

        For more examples, please consult https://github.com/tesseract-ocr/tesseract/wiki/APIExample
        )pbdoc",
                               py::dynamic_attr())
        .def("Begin", &ResultIterator::Begin,
             "Moves the iterator to point to the start of the page to begin an iteration")
        .def("Next", &ResultIterator::Next, py::arg("pageIterLv"),
             "Moves to the start of the next object at the given level in the page hierarchy in the appropriate "
             "reading order and returns false if the end of the page was reached. NOTE that RIL_SYMBOL will skip "
             "non-text blocks, but all other PageIteratorLevel level values will visit each non-text block once. Think "
             "of non text blocks as containing a single para, with a single line, with a single imaginary word. Calls "
             "to Next with different levels may be freely intermixed. This function iterates words in right-to-left "
             "scripts correctly, if the appropriate language has been loaded into Tesseract.")
        .def("Empty", &ResultIterator::Empty, py::arg("pageIterLv"),
             "Returns a boolean flag indicating whether the iterator is empty at the given PageIteratorLevel")
#if TESSERACT_VERSION >= (4 << 16 | 1 << 8)
        .def("GetBestLSTMSymbolChoices", &ResultIterator::GetBestLSTMSymbolChoices,
             "Returns the LSTM choices for every LSTM timestep for the current word.")
#endif
        .def("BoundingBox",
             [](const ResultIterator &ri, const PageIteratorLevel &lv) {
                 Box box;
                 // it is actually left, bottom, right, top
                 ri.BoundingBox(lv, &box.x, &box.h, &box.w, &box.y);
                 box.w -= box.x;
                 box.h = box.y - box.h;
                 return box;
             },
             py::arg("pageIterLv"), "Returns the bounding box of the current item")
        .def("IsAtBeginningOf", &ResultIterator::IsAtBeginningOf, py::arg("pageIterLv"),
             "IsAtBeginningOf() returns whether we're at the logical beginning of the given level. (as opposed to "
             "ResultIterator's left-to-right top-to-bottom order).")
        .def("IsAtFinalElement", &ResultIterator::IsAtFinalElement, py::arg("pageIterLv"), py::arg("element"),
             "Implement PageIterator's IsAtFinalElement correctly in a BiDi context. For instance, "
             "IsAtFinalElement(RIL_PARA, RIL_WORD) returns whether we point at the last word in a paragraph.")
        .def("ParagraphIsLtr", &ResultIterator::ParagraphIsLtr,
             "Return whether the current paragraph's dominant reading direction is left-to-right (as opposed to "
             "right-to-left).")
        .def("BlanksBeforeWord", &ResultIterator::BlanksBeforeWord,
             "Returns whether there are any blank spaces before the start of the current text object")
        .def("GetUTF8Text", &ResultIterator::GetUTF8Text, py::arg("pageIterLv"),
             "Returns the text of the current object at the specified page hierarchy level in UTF-8 format")
        .def("Confidence", &ResultIterator::Confidence, py::arg("pageIterLv"),
             "Return the confidence level expressed by the model for the current object at the specified page "
             "hierarchy level");
/**
 * VERSION_INFO is set from setup.py
 **/
#ifdef VERSION_INFO
    m.attr("__version__") = VERSION_INFO;
#else
    m.attr("__version__") = "dev";
#endif
}
