package tech.easyflow.ai.easyagentsflow.code;

import com.alibaba.fastjson2.JSON;
import com.alibaba.fastjson2.JSONObject;
import com.easyagents.flow.core.chain.Chain;
import com.easyagents.flow.core.code.CodeRuntimeEngine;
import com.easyagents.flow.core.node.CodeNode;
import com.easyagents.flow.core.util.StringUtil;

import java.io.ByteArrayOutputStream;
import java.io.File;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.nio.charset.StandardCharsets;
import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.TimeUnit;

public class PythonRuntimeEngine implements CodeRuntimeEngine {

    private static final String PYTHON_HELPER_SCRIPT = String.join("\n",
            "import contextlib",
            "import io",
            "import json",
            "import sys",
            "import traceback",
            "",
            "def _read_payload():",
            "    raw = sys.stdin.read()",
            "    if not raw:",
            "        return {}",
            "    return json.loads(raw)",
            "",
            "def _main():",
            "    payload = _read_payload()",
            "    code = payload.get('code') or ''",
            "    context = payload.get('context')",
            "    if not isinstance(context, dict):",
            "        context = {}",
            "",
            "    local_ctx = dict(context)",
            "    local_ctx['_result'] = {}",
            "",
            "    captured_stdout = io.StringIO()",
            "    captured_stderr = io.StringIO()",
            "",
            "    try:",
            "        with contextlib.redirect_stdout(captured_stdout), contextlib.redirect_stderr(captured_stderr):",
            "            exec(code, {'__builtins__': __builtins__}, local_ctx)",
            "    except Exception:",
            "        sys.stderr.write(traceback.format_exc())",
            "        sys.exit(2)",
            "",
            "    result = local_ctx.get('_result')",
            "    if not isinstance(result, dict):",
            "        sys.stderr.write('Python 脚本执行结束后，_result 必须是 dict\\n')",
            "        sys.exit(3)",
            "",
            "    if captured_stdout.getvalue():",
            "        sys.stderr.write(captured_stdout.getvalue())",
            "    if captured_stderr.getvalue():",
            "        sys.stderr.write(captured_stderr.getvalue())",
            "",
            "    try:",
            "        sys.stdout.write(json.dumps(result, ensure_ascii=False))",
            "    except Exception:",
            "        sys.stderr.write(traceback.format_exc())",
            "        sys.exit(4)",
            "",
            "if __name__ == '__main__':",
            "    _main()"
    );

    private final String command;
    private final long timeoutMs;
    private final int maxOutputBytes;
    private final String workingDir;

    public PythonRuntimeEngine(String command, long timeoutMs, int maxOutputBytes, String workingDir) {
        this.command = command;
        this.timeoutMs = timeoutMs;
        this.maxOutputBytes = maxOutputBytes;
        this.workingDir = workingDir;
    }

    public static ProbeResult probe(String command, long timeoutMs) {
        if (StringUtil.noText(command)) {
            return new ProbeResult(false, "未配置 python 命令");
        }

        Process process = null;
        StreamCollector outputCollector = null;
        Thread outputThread = null;
        try {
            process = new ProcessBuilder(command, "--version")
                    .redirectErrorStream(true)
                    .start();

            outputCollector = new StreamCollector(process.getInputStream(), 4096);
            outputThread = new Thread(outputCollector, "python-probe-output");
            outputThread.setDaemon(true);
            outputThread.start();

            boolean finished = process.waitFor(Math.max(timeoutMs, 1000), TimeUnit.MILLISECONDS);
            waitCollector(outputThread);

            String output = outputCollector == null ? "" : outputCollector.getText().trim();
            if (!finished) {
                process.destroyForcibly();
                return new ProbeResult(false, "执行 python --version 超时");
            }

            if (process.exitValue() != 0) {
                return new ProbeResult(false, StringUtil.hasText(output) ? output : "python 命令返回非 0");
            }

            if (outputCollector != null && outputCollector.isTruncated()) {
                return new ProbeResult(false, "python --version 输出过长");
            }

            return new ProbeResult(true, StringUtil.hasText(output) ? output : "Python 可用");
        } catch (Exception e) {
            return new ProbeResult(false, e.getMessage());
        } finally {
            if (process != null) {
                process.destroy();
            }
        }
    }

    @Override
    public Map<String, Object> execute(String code, CodeNode node, Chain chain) {
        if (StringUtil.noText(command)) {
            throw new RuntimeException("python 执行命令为空，请检查配置 node.code-engine.python.command");
        }

        Process process = null;
        StreamCollector stdoutCollector = null;
        StreamCollector stderrCollector = null;
        Thread stdoutThread = null;
        Thread stderrThread = null;

        try {
            ProcessBuilder builder = new ProcessBuilder(command, "-c", PYTHON_HELPER_SCRIPT);
            if (StringUtil.hasText(workingDir)) {
                builder.directory(new File(workingDir));
            }
            process = builder.start();

            stdoutCollector = new StreamCollector(process.getInputStream(), maxOutputBytes);
            stderrCollector = new StreamCollector(process.getErrorStream(), maxOutputBytes);
            stdoutThread = new Thread(stdoutCollector, "python-runtime-stdout");
            stderrThread = new Thread(stderrCollector, "python-runtime-stderr");
            stdoutThread.setDaemon(true);
            stderrThread.setDaemon(true);
            stdoutThread.start();
            stderrThread.start();

            Map<String, Object> context = buildContext(chain, node);
            Map<String, Object> payload = new HashMap<>();
            payload.put("code", code);
            payload.put("context", context);

            String payloadText = JSON.toJSONString(payload);
            try (OutputStream stdin = process.getOutputStream()) {
                stdin.write(payloadText.getBytes(StandardCharsets.UTF_8));
                stdin.flush();
            }

            boolean finished = process.waitFor(Math.max(timeoutMs, 1), TimeUnit.MILLISECONDS);
            if (!finished) {
                process.destroyForcibly();
                process.waitFor(1, TimeUnit.SECONDS);
            }

            waitCollector(stdoutThread);
            waitCollector(stderrThread);

            String stdout = stdoutCollector == null ? "" : stdoutCollector.getText().trim();
            String stderr = stderrCollector == null ? "" : stderrCollector.getText().trim();

            if (stdoutCollector != null && stdoutCollector.isTruncated()) {
                throw new RuntimeException("Python 执行输出超出限制(" + maxOutputBytes + " bytes)");
            }
            if (stderrCollector != null && stderrCollector.isTruncated()) {
                stderr = stderr + "\n...(stderr too large, truncated)";
            }

            if (!finished) {
                throw new RuntimeException("Python 脚本执行超时(" + timeoutMs + "ms)" + (StringUtil.hasText(stderr) ? "，stderr: " + stderr : ""));
            }

            int exitCode = process.exitValue();
            if (exitCode != 0) {
                String errorMsg = StringUtil.hasText(stderr) ? stderr : (StringUtil.hasText(stdout) ? stdout : "无错误输出");
                throw new RuntimeException("Python 脚本执行失败(exit=" + exitCode + "): " + errorMsg);
            }

            if (StringUtil.noText(stdout)) {
                throw new RuntimeException("Python 脚本执行失败：stdout 为空，未返回 _result JSON");
            }

            Object parsed = JSON.parse(stdout);
            if (!(parsed instanceof JSONObject)) {
                throw new RuntimeException("Python 脚本执行失败：输出不是 JSON Object");
            }

            return ((JSONObject) parsed);
        } catch (Exception e) {
            throw new RuntimeException("Python 脚本执行失败: " + e.getMessage(), e);
        } finally {
            if (process != null) {
                process.destroy();
            }
        }
    }

    private Map<String, Object> buildContext(Chain chain, CodeNode node) {
        Map<String, Object> context = new HashMap<>();

        Map<String, Object> all = chain.getState().getMemory();
        all.forEach((key, value) -> {
            if (!key.contains(".")) {
                context.put(key, value);
            }
        });

        Map<String, Object> parameterValues = chain.getState().resolveParameters(node);
        if (parameterValues != null && !parameterValues.isEmpty()) {
            context.putAll(parameterValues);
        }

        context.put("_env", chain.getState().getEnvMap());
        return context;
    }

    private static void waitCollector(Thread collectorThread) {
        if (collectorThread == null) {
            return;
        }
        try {
            collectorThread.join(2000L);
        } catch (InterruptedException e) {
            Thread.currentThread().interrupt();
        }
    }

    public static class ProbeResult {
        private final boolean available;
        private final String message;

        public ProbeResult(boolean available, String message) {
            this.available = available;
            this.message = message;
        }

        public boolean isAvailable() {
            return available;
        }

        public String getMessage() {
            return message;
        }
    }

    private static class StreamCollector implements Runnable {
        private final InputStream inputStream;
        private final int maxBytes;
        private final ByteArrayOutputStream output = new ByteArrayOutputStream();
        private volatile boolean truncated;

        private StreamCollector(InputStream inputStream, int maxBytes) {
            this.inputStream = inputStream;
            this.maxBytes = maxBytes;
        }

        @Override
        public void run() {
            byte[] buffer = new byte[4096];
            int n;
            int written = 0;
            try {
                while ((n = inputStream.read(buffer)) != -1) {
                    if (written < maxBytes) {
                        int canWrite = Math.min(maxBytes - written, n);
                        output.write(buffer, 0, canWrite);
                        written += canWrite;
                        if (canWrite < n) {
                            truncated = true;
                        }
                    } else {
                        truncated = true;
                    }
                }
            } catch (IOException ignored) {
                // 忽略流读异常，最终由进程退出码和输出判断
            }
        }

        public boolean isTruncated() {
            return truncated;
        }

        public String getText() {
            return output.toString(StandardCharsets.UTF_8);
        }
    }
}
