//  ************************************************************************************************
//
//  BornAgain: simulate and fit reflection and scattering
//
//! @file      PyCore/Embed/PyInterpreter.cpp
//! @brief     Implements functions to expose Python-interpreter functionality to C++.
//!
//! @homepage  http://www.bornagainproject.org
//! @license   GNU General Public License v3 or higher (see COPYING)
//! @copyright Forschungszentrum Jülich GmbH 2018
//! @authors   Scientific Computing Group at MLZ (see CITATION, AUTHORS)
//
//  ************************************************************************************************

/* Embedded Python Interpreter

Note that Python objects are structures allocated on the heap,
accessed through pointers of type `PyObject*`.

References:
- Python C-API <https://docs.python.org/3/c-api>
- Python ABI stability <https://docs.python.org/3/c-api/stable.html>
- Numpy C-API <https://numpy.org/doc/stable/reference/c-api>;
  <https://numpy.org/doc/stable/reference/c-api/array.html>.
- Python Extension Patterns <https://pythonextensionpatterns.readthedocs.io>
- "Python behind the scenes" series <https://tenthousandmeters.com/tag/python-behind-the-scenes>
- Python's garbage collector <https://rushter.com/blog/python-garbage-collector>
*/

#include "PyCore/Embed/PyInterpreter.h"
#include "Base/Util/Assert.h"
#include "PyCore/Embed/PyCore.h"
#include "PyCore/Embed/PyObjectPtr.h"
#include <algorithm> // transform
#include <chrono>    // time_since_epoch
#include <csignal>   // SIGINT
#include <cstddef>   // NULL
#include <cstring>   // memcpy
#include <iostream>  // cerr
#include <memory>    // unique_ptr
#include <stdexcept> // runtime_error

namespace {

std::string toString(const wchar_t* const s)
{
    if (!s)
        return "";
    std::wstring w(s);
    std::string result(w.length(), 0);
    std::transform(w.begin(), w.end(), result.begin(), [](wchar_t c) { return (char)c; });
    return result;
}

//! Converts PyObject into vector of strings, if possible, or throws exception
std::vector<std::string> toVectorString(PyObject* py_object)
{
    std::vector<std::string> result;

    if (PyTuple_Check(py_object)) {
        for (Py_ssize_t i = 0; i < PyTuple_Size(py_object); i++) {
            PyObject* value = PyTuple_GetItem(py_object, i);
            result.push_back(PyInterpreter::pyStrtoString(value));
        }
    } else if (PyList_Check(py_object)) {
        for (Py_ssize_t i = 0; i < PyList_Size(py_object); i++) {
            PyObject* value = PyList_GetItem(py_object, i);
            result.push_back(PyInterpreter::pyStrtoString(value));
        }
    } else {
        throw std::runtime_error(PyInterpreter::errorDescription(
            "PyInterpreter: Cannnot convert the given Python object "
            "to vector<string>."));
    }

    return result;
}

//! Returns seconds since epoch which is used to make unique names
std::string secondsSinceEpoch()
{
    auto now = std::chrono::system_clock::now();
    auto epoch = now.time_since_epoch();
    auto seconds = std::chrono::duration_cast<std::chrono::seconds>(epoch);
    return std::to_string(seconds.count());
}

} // namespace


// NOTE: "Python stable ABI" denotes the functions which use _only_ the Python's stable ABI;
// see <https://docs.python.org/3/c-api/stable.html>

// Python stable ABI
void PyInterpreter::initialize()
{
    if (!Py_IsInitialized())
        Py_Initialize();
}

// Python stable ABI
bool PyInterpreter::isInitialized()
{
    return static_cast<bool>(Py_IsInitialized());
}

// Python stable ABI
void PyInterpreter::finalize()
{
    // undo all initializations made by Py_Initialize() and subsequent use
    // of Python/C API functions, and destroy all sub-interpreters.
    // This is a no-op when called for a second time.
    Py_Finalize();
}

// Python stable ABI
bool PyInterpreter::checkError()
{
    if (PyErr_Occurred()) {
        // print a standard traceback to sys.stderr and clear the error indicator
        std::cerr << "---PyInterpreter: Error in Python interpreter:\n";
        PyErr_Print();
        std::cerr << "\n---\n";
        return true;
    }
    return false;
}

// Python stable ABI
void PyInterpreter::addPythonPath(const std::string& path)
{
    if (!path.empty()) {
        PyObject* sysPath = PySys_GetObject("path");
        PyList_Append(sysPath, PyUnicode_FromString(path.c_str())); // add to `PYTHONPATH`
    }
}

// Python stable ABI
void PyInterpreter::setPythonPath(const std::string& path)
{
    // returns 0 on success, -1 on error
    const int result = PySys_SetObject((char*)"path", PyUnicode_FromString(path.c_str()));
    if (result != 0) {
        PyInterpreter::checkError();
        throw std::runtime_error("PyInterpreter.setPythonPath: Cannot set the Python path.");
    }
}

// Python stable ABI
PyObjectPtr PyInterpreter::import(const std::string& pymodule_name, const std::string& path)
{
    ASSERT(!pymodule_name.empty());

    addPythonPath(path);

    // import the module
    PyObject* pymodule = PyImport_ImportModule(pymodule_name.c_str());
    if (!pymodule || !PyModule_Check(pymodule)) {
        checkError();
        throw std::runtime_error(errorDescription("PyInterpreter: Cannot load Python module '"
                                                  + pymodule_name + "' (given path = '" + path
                                                  + "')"));
    }

    // returns a _new_ reference; ie. caller is responsible for the ref-count
    return {pymodule};
}

// Python stable ABI
void PyInterpreter::DecRef(PyObject* py_object)
{
    // Decrementing Python reference-count without
    // Python initialized leads to memory access violation (segmentation fault)
    ASSERT(PyInterpreter::isInitialized());

    Py_XDECREF(py_object);
}

std::string PyInterpreter::pyStrtoString(PyObject* py_object)
{
    std::string result;
    PyObject* pyStr = PyUnicode_AsEncodedString(py_object, "utf-8", "replace");
    if (pyStr) {
        result = std::string(PyBytes_AsString(pyStr));
        Py_DecRef(pyStr);
    }

    return result;
}

std::string PyInterpreter::getStrAttribute(PyObject* py_module, const std::string& attr_name)
{
    std::string result_str;
    PyObject* pAttr = PyObject_GetAttrString(py_module, attr_name.c_str());

    if (!pAttr) {
        Py_DecRef(pAttr);
        checkError();
        throw std::runtime_error(
            errorDescription("PyInterpreter: Cannot get Python attribute '" + attr_name + "'"));
    }

    result_str = PyInterpreter::pyStrtoString(pAttr);
    Py_DecRef(pAttr);
    return result_str;
}

void PyInterpreter::callFunction(PyObject* py_module, const std::string& fn_name)
{

    // NOTE: The Python function signature is assumed to be `fn() -> str`
    std::string result_str;
    PyObject *pFunc, *pRetValue;

    pFunc = PyObject_GetAttrString(py_module, fn_name.c_str());

    if (!pFunc || !PyCallable_Check(pFunc)) {
        Py_DecRef(pFunc);
        checkError();
        throw std::runtime_error(
            errorDescription("PyInterpreter: Cannot call Python function '" + fn_name + "'"));
    }

    // Call the function with no arguments
    pRetValue = PyObject_CallObject(pFunc, NULL);
    result_str = PyInterpreter::pyStrtoString(pRetValue);
    Py_DecRef(pRetValue);
    Py_DecRef(pFunc);
}

std::string PyInterpreter::runtimeInfo()
{
    std::string result;

    // Runtime environment
    result = std::string(60, '=') + "\n";
    // Embedded Python details
    result += "Py_GetProgramName(): " + toString(Py_GetProgramName()) + "\n";
    result += "Py_GetProgramFullPath(): " + toString(Py_GetProgramFullPath()) + "\n";
    result += "Py_GetPath(): " + toString(Py_GetPath()) + "\n";
    result += "Py_GetPythonHome(): " + toString(Py_GetPythonHome()) + "\n";

    // Runtime Python's sys.path
    PyObject* sysPath = PySys_GetObject((char*)"path");
    std::vector<std::string> content{toVectorString(sysPath)};
    result += "sys.path: ";
    for (const std::string& s : content)
        result += s + ",";
    result += "\n";

    return result;
}

// Attempt to retrieve Python stack trace
// Ref: <https://stackoverflow.com/q/1796510>
std::string PyInterpreter::stackTrace()
{
    std::string result;

    if (PyErr_Occurred()) {
        PyObject *ptype, *pvalue, *ptraceback, *pystr;

        PyErr_Fetch(&ptype, &pvalue, &ptraceback);
        pystr = PyObject_Str(pvalue);
        if (char* str = PyBytes_AsString(pystr))
            result += std::string(str) + "\n";
        Py_DecRef(pystr);

        PyObject* module_name = PyUnicode_FromString("traceback");
        PyObject* py_traceback_module = PyImport_Import(module_name);
        Py_DecRef(module_name);

        if (py_traceback_module) {
            result += "\n";
            PyObject* pyth_func = PyObject_GetAttrString(py_traceback_module, "format_exception");
            if (pyth_func && PyCallable_Check(pyth_func)) {
                PyObject* pyth_val =
                    PyObject_CallFunctionObjArgs(pyth_func, ptype, pvalue, ptraceback, NULL);
                Py_DecRef(pyth_func);
                if (pyth_val) {
                    pystr = PyObject_Str(pyth_val);
                    if (char* str = PyBytes_AsString(pystr))
                        result += std::string(str);
                    Py_DecRef(pystr);
                    Py_DecRef(pyth_val);
                }
            }
            result += "\n";
        }

        Py_DecRef(py_traceback_module);
    }

    result += "\n";
    result += PyInterpreter::runtimeInfo();
    result += "\n";

    return result;
}

std::string PyInterpreter::errorDescription(const std::string& title)
{
    std::string msg = title + "\n" + PyInterpreter::stackTrace() + "\n";
    return msg;
}


PyObjectPtr PyInterpreter::BornAgain::import(const std::string& path)
{
    PyInterpreter::addPythonPath(path);

#ifndef _WIN32
    // store ctrl-C handler before Numpy messes it up
    PyOS_sighandler_t sighandler = PyOS_getsig(SIGINT);
#endif

    PyObject* ba_pymodule = PyImport_ImportModule("bornagain");

#ifndef _WIN32
    PyOS_setsig(SIGINT, sighandler); // restore previous ctrl-C handler
#endif

    if (!ba_pymodule || !PyModule_Check(ba_pymodule)) {
        checkError();
        throw std::runtime_error(
            errorDescription("PyInterpreter: Cannot load 'bornagain' Python module "
                             "(given path = '"
                             + path + "')"));
    }

    return {ba_pymodule};
}

PyObjectPtr PyInterpreter::BornAgain::importScript(const std::string& script,
                                                   const std::string& path)
{
    // import bornagain Python module
    PyObjectPtr pBAModule{PyInterpreter::BornAgain::import(path)};

    // introduce a new module to execute the Python script
    const std::string tmp_module_name = "_ba_tmp_module_" + secondsSinceEpoch() + "_";
    PyObject* pModule = PyImport_AddModule(tmp_module_name.c_str()); // borrowed reference
    if (pModule == NULL) {
        throw std::runtime_error(errorDescription("Cannot add a temporary Python module"));
        return nullptr;
    }

    PyObject* pGlobalDict = PyModule_GetDict(pModule); // borrowed reference
    PyDict_SetItemString(pGlobalDict, "bornagain", pBAModule.get());

    // execute the function in the module namespace
    PyObject* pExeResult = PyRun_String(script.c_str(), Py_file_input, pGlobalDict, pGlobalDict);

    if (pExeResult == NULL) {
        throw std::runtime_error(errorDescription("Failed to run the script"));
        return nullptr;
    }

    Py_DecRef(pExeResult);

    return PyObjectPtr{pModule, /* borrowed = */ true};
}


PyObjectPtr PyInterpreter::BornAgain::callScriptFunction(const std::string& functionName,
                                                         const std::string& script,
                                                         const std::string& path)
{
    PyObjectPtr tmpModule{PyInterpreter::BornAgain::importScript(script, path)};

    // locate the function (it is an attribute of the module)
    PyObject* pFunc = PyObject_GetAttrString(tmpModule.get(), functionName.c_str());
    if (!pFunc)
        throw std::runtime_error("PyInterpreter::BornAgain: "
                                 "Cannot locate the compiled function '''\n"
                                 + functionName + "\n'''");

    // call the function with signature f() -> PythonObject
    PyObject* pResult = PyObject_CallFunctionObjArgs(pFunc, NULL);
    Py_DecRef(pFunc);

    if (!pResult)
        throw std::runtime_error("Failed executing Python function '" + functionName + "'");

    return {pResult};
}

std::vector<std::string> PyInterpreter::BornAgain::listOfFunctions(const std::string& script,
                                                                   const std::string& path)
{
    PyObjectPtr tmpModule{PyInterpreter::BornAgain::importScript(script, path)};

    if (!tmpModule.valid())
        throw std::runtime_error("PyInterpreter::BornAgain: "
                                 "Cannot obtain the dictionary from the script module");

    PyObject* pDict = PyModule_GetDict(tmpModule.get());
    if (!pDict) {
        checkError();
        throw std::runtime_error("PyInterpreter::BornAgain: "
                                 "Cannot obtain the dictionary from the script module");
    }

    PyObject *key, *value;
    Py_ssize_t pos = 0;
    std::vector<std::string> fn_names;
    while (PyDict_Next(pDict, &pos, &key, &value)) {
        if (PyCallable_Check(value)) {
            std::string func_name{PyInterpreter::pyStrtoString(key)};
            // do not import callables whose names begin with '__'
            if (func_name.find("__") == std::string::npos)
                fn_names.push_back(func_name);
        }
    }

    Py_DecRef(pDict);
    return fn_names;
}


PyObjectPtr PyInterpreter::Fabio::import()
{
    return {PyInterpreter::import("fabio")};
}

PyObjectPtr PyInterpreter::Fabio::open(const std::string& filename, PyObjectPtr& fabio_module)
{
    // load an image via calling `fabio.load` function which takes a
    // filename (Python string) and returns a Numpy array.

    if (!fabio_module.valid() || !PyModule_Check(fabio_module.get())) {
        throw std::runtime_error(errorDescription("PyInterpreter.fabio: Invalid Python module "
                                                  "(expected 'fabio' module)"));
    }

    PyObject* pFunc = PyObject_GetAttrString(fabio_module.get(), (char*)"open");
    if (!pFunc || !PyCallable_Check(pFunc)) {
        PyInterpreter::checkError();
        throw std::runtime_error(
            errorDescription("PyInterpreter.fabio: The function 'fabio.open' is not callable"));
    }

    // convert the filename to a Python unicode string
    PyObject* pFilename = PyUnicode_FromString(filename.c_str());
    if (!pFilename) {
        PyInterpreter::checkError();
        throw std::runtime_error(
            errorDescription("PyInterpreter.fabio: Filename '" + filename
                             + "' cannot be converted to Python unicode string"));
    }

    // get the result of the call `fabio.open(<filename>)`
    PyObject* pResult_open = PyObject_CallFunctionObjArgs(pFunc, pFilename, NULL);
    Py_DecRef(pFunc);
    if (!pResult_open) {
        PyInterpreter::checkError();
        std::runtime_error(errorDescription("PyInterpreter.fabio: "
                                            "Invalid return value from calling the function "
                                            "'fabio.open(\""
                                            + filename + "\")'"));
    }

    // get `result.data` (must be a Numpy array)
    PyObject* npyArray_ptr = PyObject_GetAttrString(pResult_open, (char*)"data");
    Py_DecRef(pResult_open);
    if (!npyArray_ptr) {
        PyInterpreter::checkError();
        std::runtime_error(errorDescription("PyInterpreter.fabio: Invalid return value from "
                                            "calling the function 'fabio.open(\""
                                            + filename + "\")' (expected a Numpy array)"));
    }

    // returns a _new_ reference; ie. caller is responsible for the ref-count
    return {npyArray_ptr};
}
