package com.axis.innovators.box.util;

import org.python.core.PyException;
import org.python.core.PySyntaxError;
import org.python.util.PythonInterpreter;
import java.io.File;
import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.nio.file.*;
import java.util.*;

/**
 * Python解释器工具类
 * @author tzdwindows 7
 */
public class RunPython {
    private static final String AST_PROCESSOR =
            "def __process_code(raw_code):\n" +
                    "    try:\n" +
                    "        import ast, traceback\n" +
                    "        tree = ast.parse(raw_code)\n" +
                    "        for node in ast.walk(tree):\n" +
                    "            if isinstance(node, ast.FunctionDef):\n" +
                    "                node.returns = None\n" +
                    "                for arg in node.args.args:\n" +
                    "                    arg.annotation = None\n" +
                    "        return compile(tree, '<string>', 'exec')\n" +
                    "    except Exception:\n" +
                    "        return compile(raw_code, '<string>', 'exec')\n";

    // 线程局部存储优化
    private static final ThreadLocal<PythonInterpreter> PY_THREAD_LOCAL =
            ThreadLocal.withInitial(() -> {
                PythonInterpreter pi = new PythonInterpreter();
                pi.exec("from __future__ import print_function");
                pi.exec("import sys");
                pi.exec(AST_PROCESSOR);
                return pi;
            });

    public static PythonResult executeWithJython(String scriptPath) {
        return executeWithJython(scriptPath, Collections.emptyList());
    }

    public static Object callPythonFunction(PythonResult result,
                                            String functionName,
                                            Object... params) {
        if (result == null || !result.isValid()) {
            return "INVALID_INTERPRETER";
        }
        try {
            return result.callFunction(functionName, params);
        } catch (Exception e) {
            return "CALL_ERROR: " + e.getMessage();
        }
    }

    public static PythonResult executeWithJython(String scriptPath,
                                                 List<String> dependencyPaths) {
        PythonInterpreter pi = PY_THREAD_LOCAL.get();
        try {
            pi.cleanup();
            pi.exec("sys.path = []");

            pi.exec("import sys");
            pi.exec("reload(sys)");
            pi.exec("sys.setdefaultencoding('utf-8')");

            dependencyPaths.stream()
                    .map(dep -> new File(dep).getAbsolutePath().replace("\\", "/"))
                    .forEach(path -> pi.exec("sys.path.append(r'" + path + "')"));

            String code = Files.readString(Path.of(scriptPath), StandardCharsets.UTF_8);
            pi.set("raw_code", code);
            pi.exec("exec(__process_code(raw_code))");

            return PythonResult.of(
                    Optional.ofNullable(pi.get("result", Integer.class)).orElse(0),
                    Optional.ofNullable(pi.get("errorResult", String.class)).orElse(""),
                    pi
            );
        } catch (NoSuchFileException e) {
            return PythonResult.error(-2, "FILE_NOT_FOUND: " + scriptPath);
        } catch (IOException e) {
            return PythonResult.error(-3, "IO_ERROR: " + e.getMessage());
        } catch (PySyntaxError e) {
            String errorMsg = "SYNTAX_ERROR: " + e.value.toString();
            return PythonResult.error(-4, errorMsg);
        } catch (PyException e) {
            return PythonResult.error(-5, "RUNTIME_ERROR: " + e.getMessage());
        }
    }
}