/**
 * @file python_execute.cpp
 * @brief OpenManus Python执行工具实现
 * 
 * 这个文件实现了OpenManus的Python执行工具，使用Python.h直接调用Python解释器。
 */

#include "mcp/include/mcp_server.h"
#include "mcp/include/mcp_tool.h"
#include "mcp/include/mcp_resource.h"

#include <iostream>
#include <string>
#include <memory>
#include <stdexcept>
#include <mutex>

// 检查是否找到Python
#ifdef PYTHON_FOUND
#include <Python.h>
#endif

/**
 * @class python_interpreter
 * @brief Python解释器类，用于执行Python代码
 */
class python_interpreter {
private:
    // 互斥锁，确保Python解释器的线程安全
    mutable std::mutex py_mutex;
    bool is_initialized;

public:
    /**
     * @brief 构造函数，初始化Python解释器
     */
    python_interpreter() : is_initialized(false) {
#ifdef PYTHON_FOUND
        try {
            Py_Initialize();
            if (Py_IsInitialized()) {
                is_initialized = true;
                // 初始化线程支持
                PyEval_InitThreads();
                // 释放GIL，允许其他线程获取
                PyThreadState *_save = PyEval_SaveThread();
            } else {
                std::cerr << "Python解释器初始化失败" << std::endl;
            }
        } catch (const std::exception& e) {
            std::cerr << "Python解释器初始化异常: " << e.what() << std::endl;
        }
#endif
    }

    /**
     * @brief 析构函数，释放Python解释器
     */
    ~python_interpreter() {
#ifdef PYTHON_FOUND
        if (is_initialized) {
            std::lock_guard<std::mutex> lock(py_mutex);
            Py_Finalize();
            is_initialized = false;
        }
#endif
    }

    /**
     * @brief 执行Python代码
     * @param input 包含Python代码的JSON对象
     * @return 执行结果的JSON对象
     */
    mcp::json forward(const mcp::json& input) const {
#ifdef PYTHON_FOUND
        if (!is_initialized) {
            return mcp::json{{"error", "Python解释器未正确初始化"}};
        }

        // 获取GIL锁
        std::lock_guard<std::mutex> lock(py_mutex);
        PyGILState_STATE gstate = PyGILState_Ensure();

        mcp::json result_json;
        
        try {
            if (input.contains("code") && input["code"].is_string()) {
                std::string code = input["code"].get<std::string>();

                // 获取主模块和字典
                PyObject *main_module = PyImport_AddModule("__main__");
                if (!main_module) {
                    PyGILState_Release(gstate);
                    return mcp::json{{"error", "无法获取Python主模块"}};
                }

                PyObject *main_dict = PyModule_GetDict(main_module);
                if (!main_dict) {
                    PyGILState_Release(gstate);
                    return mcp::json{{"error", "无法获取Python主模块字典"}};
                }

                // 导入sys和io模块
                PyObject *sys_module = PyImport_ImportModule("sys");
                if (!sys_module) {
                    PyErr_Print();
                    PyGILState_Release(gstate);
                    return mcp::json{{"error", "无法导入sys模块"}};
                }

                PyObject *io_module = PyImport_ImportModule("io");
                if (!io_module) {
                    Py_DECREF(sys_module);
                    PyErr_Print();
                    PyGILState_Release(gstate);
                    return mcp::json{{"error", "无法导入io模块"}};
                }

                // 获取StringIO类
                PyObject *string_io = PyObject_GetAttrString(io_module, "StringIO");
                if (!string_io) {
                    Py_DECREF(io_module);
                    Py_DECREF(sys_module);
                    PyErr_Print();
                    PyGILState_Release(gstate);
                    return mcp::json{{"error", "无法获取StringIO类"}};
                }

                // 创建StringIO对象
                PyObject *sys_stdout = PyObject_CallObject(string_io, nullptr);
                if (!sys_stdout) {
                    Py_DECREF(string_io);
                    Py_DECREF(io_module);
                    Py_DECREF(sys_module);
                    PyErr_Print();
                    PyGILState_Release(gstate);
                    return mcp::json{{"error", "无法创建stdout StringIO对象"}};
                }

                PyObject *sys_stderr = PyObject_CallObject(string_io, nullptr);
                if (!sys_stderr) {
                    Py_DECREF(sys_stdout);
                    Py_DECREF(string_io);
                    Py_DECREF(io_module);
                    Py_DECREF(sys_module);
                    PyErr_Print();
                    PyGILState_Release(gstate);
                    return mcp::json{{"error", "无法创建stderr StringIO对象"}};
                }

                // 保存原始的stdout和stderr
                PyObject *old_stdout = PySys_GetObject("stdout");
                PyObject *old_stderr = PySys_GetObject("stderr");
                
                if (old_stdout) Py_INCREF(old_stdout);
                if (old_stderr) Py_INCREF(old_stderr);

                // 替换sys.stdout和sys.stderr
                if (PySys_SetObject("stdout", sys_stdout) != 0 || 
                    PySys_SetObject("stderr", sys_stderr) != 0) {
                    Py_DECREF(sys_stderr);
                    Py_DECREF(sys_stdout);
                    Py_DECREF(string_io);
                    Py_DECREF(io_module);
                    Py_DECREF(sys_module);
                    PyErr_Print();
                    PyGILState_Release(gstate);
                    return mcp::json{{"error", "无法设置stdout/stderr重定向"}};
                }

                // 执行Python代码
                PyObject *result = PyRun_String(code.c_str(), Py_file_input, main_dict, main_dict);
                if (!result) {
                    PyErr_Print();
                }
                Py_XDECREF(result);

                // 获取输出和错误
                PyObject *out_value = PyObject_CallMethod(sys_stdout, "getvalue", nullptr);
                PyObject *err_value = PyObject_CallMethod(sys_stderr, "getvalue", nullptr);

                std::string output, error;
                
                // 安全地转换Python字符串到C++字符串
                if (out_value && PyUnicode_Check(out_value)) {
                    output = PyUnicode_AsUTF8(out_value);
                }
                
                if (err_value && PyUnicode_Check(err_value)) {
                    error = PyUnicode_AsUTF8(err_value);
                }

                // 恢复原始的stdout和stderr
                if (old_stdout) {
                    PySys_SetObject("stdout", old_stdout);
                    Py_DECREF(old_stdout);
                }
                
                if (old_stderr) {
                    PySys_SetObject("stderr", old_stderr);
                    Py_DECREF(old_stderr);
                }

                // 清理
                Py_XDECREF(out_value);
                Py_XDECREF(err_value);
                Py_DECREF(sys_stdout);
                Py_DECREF(sys_stderr);
                Py_DECREF(string_io);
                Py_DECREF(io_module);
                Py_DECREF(sys_module);

                // 准备JSON输出
                if (!output.empty()) {
                    result_json["output"] = output;
                }
                if (!error.empty()) {
                    result_json["error"] = error;
                }

                if (result_json.empty()) {
                    std::string last_line;
                    std::istringstream code_stream(code);
                    while (std::getline(code_stream, last_line, '\n')) {}
                    size_t pos = last_line.find_last_of(';') + 1;
                    pos = last_line.find("=") + 1;
                    while (pos < last_line.size() && isblank(last_line[pos])) {
                        pos++;
                    }
                    if (pos != std::string::npos) {
                        last_line = last_line.substr(pos);
                    }

                    result_json["warning"] = "No output. Maybe try with print(" + last_line + ")?";
                }
            } else {
                result_json["error"] = "Invalid parameters or code not provided";
            }
        } catch (const std::exception& e) {
            result_json["error"] = std::string("Python执行异常: ") + e.what();
        }

        // 释放GIL
        PyGILState_Release(gstate);
        return result_json;
#else
        return mcp::json{{"error", "Python interpreter not available"}};
#endif
    }
};

// 全局Python解释器实例
static python_interpreter interpreter;

// Python执行工具处理函数
mcp::json python_execute_handler(const mcp::json& args) {
    if (!args.contains("code")) {
        throw mcp::mcp_exception(mcp::error_code::invalid_params, "缺少'code'参数");
    }
    
    try {
        // 使用Python解释器执行代码
        mcp::json result = interpreter.forward(args);
        
        return {{
            {"type", "text"},
            {"text", result.dump(2)}
        }};
    } catch (const std::exception& e) {
        throw mcp::mcp_exception(mcp::error_code::internal_error, 
                                "执行Python代码失败: " + std::string(e.what()));
    }
}

// Register the PythonExecute tool
void register_python_execute_tool(mcp::server& server) {
    mcp::tool python_tool = mcp::tool_builder("python_execute")
        .with_description("Execute Python code and return the result")
        .with_string_param("code", "The Python code to execute", true)
        .build();
    
    server.register_tool(python_tool, python_execute_handler);
}