package com.gzzm.lobster.sandbox;

import com.gzzm.lobster.config.LobsterConfig;
import com.gzzm.platform.commons.Tools;
import net.cyan.arachne.annotation.Service;

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStream;
import java.io.InputStreamReader;
import java.nio.charset.StandardCharsets;
import java.util.ArrayList;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.TimeUnit;
import java.util.function.BooleanSupplier;
import java.util.function.Consumer;

/**
 * DockerRunner —— 封装 docker 子进程调用 / Thin wrapper over docker CLI.
 *
 * <p>用 {@link ProcessBuilder} 直接起 docker 子进程；不引入 docker-java 依赖，
 * 避免镜像仓库 / SPI 级别的耦合. 在 Windows 以及 Linux 下均可用，只要 docker CLI 在 PATH.
 *
 * <p><b>生命周期：三段式 {@code docker create → 登记真实 container id → docker start -a → wait}</b>：
 * <ul>
 *   <li>{@code docker run} 一条命令把 create 和 start 合在一起，但登记用的是调用方
 *       猜的 {@code --name}，"put 进 map" 和 "docker daemon 真的创建好容器" 之间还是
 *       有几十到几百毫秒的空窗—— {@code killByRun} 在这段时间 {@code docker kill} 会
 *       返 "No such container"</li>
 *   <li>拆成 create + start 后，{@code docker create} 同步返回的已经是 **已创建**
 *       的容器 id，通过 {@link RunSpec#onContainerCreated} 回调登记给 SandboxService；
 *       之后任何 {@code kill/rm} 都能命中</li>
 * </ul>
 */
@Service
public class DockerRunner {

    public static final int EXIT_TIMEOUT = -1;
    public static final int EXIT_OOM = -2;

    public static final class RunSpec {
        public String image;
        public String containerName;
        /** 非 root uid:gid，形如 "10001:10001" */
        public String user;
        public int memoryMb;
        public double cpus;
        public int pidsLimit;
        public long walltimeSec;
        /** key=hostPath value=containerPath[:ro] */
        public final List<Mount> mounts = new ArrayList<>();
        /** tmpfs mounts，形如 /tmp:size=64m */
        public final List<String> tmpfs = new ArrayList<>();
        /** 传给 image entrypoint 的参数 */
        public final List<String> args = new ArrayList<>();
        /**
         * 可选的 entrypoint override. null 使用镜像 Dockerfile 的 ENTRYPOINT.
         * 场景：Dockerfile 默认 `tini -- python`（跑 .py 用）；language=javascript 时
         * SandboxService 设 entrypoint="/usr/bin/tini"，args 前置 "--","node","/work/entry.js".
         */
        public String entrypoint;
        /**
         * 外部取消检查 / External cancel poll hook. 每秒检查一次；返 true 则立刻 docker kill.
         * 与 {@link #onContainerCreated} 配合消除 "登记 → 容器就绪" 的窗口期.
         */
        public BooleanSupplier cancelCheck;
        /**
         * 容器创建成功后回调 / Invoked once docker create returns a container id.
         * 调用方此时应把 id 登记到取消表，后续 {@code kill} 都用这个 id（而不是 name）.
         */
        public Consumer<String> onContainerCreated;
    }

    public static final class Mount {
        public final String host;
        public final String container;
        public final boolean readOnly;
        public Mount(String host, String container, boolean readOnly) {
            this.host = host; this.container = container; this.readOnly = readOnly;
        }
    }

    public static final class RunOutcome {
        public int exitCode;
        public String stdout = "";
        public String stderr = "";
        public boolean timedOut;
        public boolean oomKilled;
        public boolean cancelled;
    }

    /** 阻塞执行：create → register → start -a → poll wait → 收集结果 → 保证回收. */
    public RunOutcome run(RunSpec spec) throws SandboxException {
        RunOutcome out = new RunOutcome();
        String containerId = create(spec);
        try {
            if (spec.onContainerCreated != null) {
                try { spec.onContainerCreated.accept(containerId); }
                catch (Throwable ignore) { /* 不阻塞 start */ }
            }
            // create 完成到 start 之间的窗口也要判一次 cancel
            if (spec.cancelCheck != null) {
                try {
                    if (spec.cancelCheck.getAsBoolean()) {
                        out.cancelled = true;
                        out.exitCode = EXIT_TIMEOUT;
                        return out;
                    }
                } catch (Throwable ignore) { /* ignore */ }
            }
            return startCreated(spec, containerId);
        } finally {
            // 本类不给 docker create 加 --rm（版本兼容性问题），所以每次都显式 rm -f 兜底.
            // 容器已退出 / 正在删 → rm -f 返 "No such container" 也无所谓.
            removeContainerQuiet(containerId);
        }
    }

    public RunOutcome startCreated(RunSpec spec, String containerId) throws SandboxException {
        RunOutcome out = new RunOutcome();
        dockerStartAttach(spec, containerId, out);
        // OOM 检测：docker 137 + stderr "OOMKilled"
        if (!out.cancelled && !out.timedOut
                && (out.exitCode == 137
                    || (out.stderr != null && out.stderr.contains("OOMKilled")))) {
            out.oomKilled = true;
            out.exitCode = EXIT_OOM;
        }
        return out;
    }

    public String create(RunSpec spec) throws SandboxException {
        List<String> cmd = new ArrayList<>();
        cmd.add(LobsterConfig.getSandboxDockerBin());
        cmd.add("create");
        // 注意：docker create 不加 --rm。--rm 在部分 docker 版本是 `docker run` 专属，
        // 在 create 命令上会直接报 "unknown flag: --rm"；即便较新版本支持，也没必要 ——
        // 本类 run() 的 finally 会 `docker rm -f <cid>` 兜底清理，语义等价且跨版本稳.
        cmd.add("--read-only");
        cmd.add("--network"); cmd.add("none");
        cmd.add("--cap-drop"); cmd.add("ALL");
        cmd.add("--security-opt"); cmd.add("no-new-privileges");
        if (spec.user != null && !spec.user.isEmpty()) {
            cmd.add("--user"); cmd.add(spec.user);
        }
        if (spec.containerName != null && !spec.containerName.isEmpty()) {
            cmd.add("--name"); cmd.add(spec.containerName);
        }
        if (spec.entrypoint != null && !spec.entrypoint.isEmpty()) {
            cmd.add("--entrypoint"); cmd.add(spec.entrypoint);
        }
        cmd.add("--memory"); cmd.add(spec.memoryMb + "m");
        cmd.add("--cpus"); cmd.add(String.valueOf(spec.cpus));
        cmd.add("--pids-limit"); cmd.add(String.valueOf(spec.pidsLimit));
        cmd.add("--ulimit"); cmd.add("nofile=256:256");
        for (Mount m : spec.mounts) {
            cmd.add("-v");
            cmd.add(m.host + ":" + m.container + (m.readOnly ? ":ro" : ""));
        }
        for (String tmp : spec.tmpfs) {
            cmd.add("--tmpfs");
            cmd.add(tmp);
        }
        cmd.add(spec.image);
        cmd.addAll(spec.args);

        Process p;
        try {
            p = new ProcessBuilder(cmd).redirectErrorStream(false).start();
        } catch (IOException e) {
            throw new SandboxException("sandbox.docker_create",
                    "failed to run docker create: " + e.getMessage(), e);
        }
        // docker create 输出 container id 一行到 stdout；stderr 单独抓用于错误信息
        StringBuilder idBuf = new StringBuilder();
        StreamDrainer errDrain = new StreamDrainer(p.getErrorStream());
        Thread tErr = new Thread(errDrain, "sandbox-create-err");
        tErr.setDaemon(true); tErr.start();
        try (BufferedReader r = new BufferedReader(new InputStreamReader(p.getInputStream(), StandardCharsets.UTF_8))) {
            String line;
            while ((line = r.readLine()) != null) idBuf.append(line);
        } catch (IOException ignore) { /* proc may end */ }
        boolean done;
        try {
            done = p.waitFor(30, TimeUnit.SECONDS);
        } catch (InterruptedException ie) {
            p.destroyForcibly();
            Thread.currentThread().interrupt();
            throw new SandboxException("sandbox.docker_create", "create interrupted");
        }
        try { tErr.join(1000); } catch (InterruptedException ignore) { /* ignore */ }
        if (!done) {
            p.destroyForcibly();
            throw new SandboxException("sandbox.docker_create", "create timed out");
        }
        if (p.exitValue() != 0) {
            throw new SandboxException("sandbox.docker_create",
                    "create failed (exit=" + p.exitValue() + "): " + errDrain.getText());
        }
        String id = idBuf.toString().trim();
        if (id.isEmpty()) {
            throw new SandboxException("sandbox.docker_create", "empty container id from docker create");
        }
        return id;
    }

    private void dockerStartAttach(RunSpec spec, String containerId, RunOutcome out) throws SandboxException {
        List<String> cmd = new ArrayList<>();
        cmd.add(LobsterConfig.getSandboxDockerBin());
        cmd.add("start"); cmd.add("-a"); cmd.add(containerId);

        Process proc;
        try {
            proc = new ProcessBuilder(cmd).redirectErrorStream(false).start();
        } catch (IOException e) {
            throw new SandboxException("sandbox.docker_start",
                    "failed to start container: " + e.getMessage(), e);
        }

        StreamDrainer stdoutDrainer = new StreamDrainer(proc.getInputStream());
        StreamDrainer stderrDrainer = new StreamDrainer(proc.getErrorStream());
        Thread tOut = new Thread(stdoutDrainer, "sandbox-stdout-" + shortId(containerId));
        Thread tErr = new Thread(stderrDrainer, "sandbox-stderr-" + shortId(containerId));
        tOut.setDaemon(true); tErr.setDaemon(true);
        tOut.start(); tErr.start();

        long deadlineMs = System.currentTimeMillis() + spec.walltimeSec * 1000L;
        boolean finished = false, cancelled = false;
        try {
            while (System.currentTimeMillis() < deadlineMs) {
                if (spec.cancelCheck != null) {
                    try {
                        if (spec.cancelCheck.getAsBoolean()) {
                            cancelled = true;
                            break;
                        }
                    } catch (Throwable ignore) { /* ignore */ }
                }
                if (proc.waitFor(1, TimeUnit.SECONDS)) {
                    finished = true;
                    break;
                }
            }
        } catch (InterruptedException ie) {
            kill(containerId);
            proc.destroyForcibly();
            Thread.currentThread().interrupt();
            out.exitCode = EXIT_TIMEOUT;
            out.timedOut = true;
            return;
        }
        if (cancelled) {
            kill(containerId);
            try { proc.waitFor(2, TimeUnit.SECONDS); } catch (InterruptedException ignore) { Thread.currentThread().interrupt(); }
            proc.destroyForcibly();
            out.cancelled = true;
            out.exitCode = EXIT_TIMEOUT;
        } else if (!finished) {
            kill(containerId);
            proc.destroyForcibly();
            out.exitCode = EXIT_TIMEOUT;
            out.timedOut = true;
        } else {
            // `docker start -a` 的 CLI 退出码在多数版本等于容器内进程退出码，但并非严格保证
            // （网络断连 / attach 失败时会用 CLI 自己的退码）. 用 `docker wait` 拿容器权威退码.
            int containerExit = dockerWaitExitCode(containerId);
            out.exitCode = containerExit >= 0 ? containerExit : proc.exitValue();
        }
        try { tOut.join(2000); } catch (InterruptedException ignore) { Thread.currentThread().interrupt(); }
        try { tErr.join(2000); } catch (InterruptedException ignore) { Thread.currentThread().interrupt(); }
        out.stdout = stdoutDrainer.getText();
        out.stderr = stderrDrainer.getText();
    }

    /**
     * 调 {@code docker wait <id>} 拿容器退出码（stdout 是一个整数一行）.
     * 失败或容器已被 rm 返 -1，让调用方 fallback.
     */
    private int dockerWaitExitCode(String containerId) {
        try {
            Process p = new ProcessBuilder(LobsterConfig.getSandboxDockerBin(), "wait", containerId)
                    .redirectErrorStream(false).start();
            StringBuilder buf = new StringBuilder();
            try (BufferedReader r = new BufferedReader(new InputStreamReader(p.getInputStream(), StandardCharsets.UTF_8))) {
                String line;
                while ((line = r.readLine()) != null) buf.append(line);
            }
            if (!p.waitFor(5, TimeUnit.SECONDS)) {
                p.destroyForcibly();
                return -1;
            }
            if (p.exitValue() != 0) return -1;
            try { return Integer.parseInt(buf.toString().trim()); }
            catch (NumberFormatException ignore) { return -1; }
        } catch (Throwable t) {
            return -1;
        }
    }

    /**
     * 通过 id（或 name） kill 容器；best-effort. 命令返非零（已退出 / 不存在）不记错.
     */
    public void kill(String containerIdOrName) {
        if (containerIdOrName == null || containerIdOrName.isEmpty()) return;
        try {
            Process p = new ProcessBuilder(LobsterConfig.getSandboxDockerBin(), "kill", containerIdOrName)
                    .redirectErrorStream(true).start();
            if (!p.waitFor(5, TimeUnit.SECONDS)) {
                p.destroyForcibly();
            }
        } catch (Throwable t) {
            try { Tools.log("[DockerRunner] kill failed: " + containerIdOrName, t); } catch (Throwable ignore) { /* ignore */ }
        }
    }

    /** docker rm -f，兜底清理；--rm 已生效的正常 run 返 "No such container" 无妨. */
    public void removeContainerQuiet(String containerId) {
        if (containerId == null || containerId.isEmpty()) return;
        try {
            Process p = new ProcessBuilder(LobsterConfig.getSandboxDockerBin(), "rm", "-f", containerId)
                    .redirectErrorStream(true).start();
            if (!p.waitFor(5, TimeUnit.SECONDS)) {
                p.destroyForcibly();
            }
        } catch (Throwable ignore) { /* best effort */ }
    }

    /**
     * Returns Docker's current container state, or null when the container is
     * missing/uninspectable. Pool reuse only accepts a clean never-started
     * container; anything else is safer to discard and recreate.
     */
    public String inspectContainerStatus(String containerIdOrName) {
        if (containerIdOrName == null || containerIdOrName.isEmpty()) return null;
        try {
            List<String> cmd = new ArrayList<>();
            cmd.add(LobsterConfig.getSandboxDockerBin());
            cmd.add("inspect");
            cmd.add("--format");
            cmd.add("{{.State.Status}}");
            cmd.add(containerIdOrName);
            List<String> lines = runDockerLines(cmd, 5, "sandbox.docker_inspect");
            String status = lines.isEmpty() ? "" : lines.get(0).trim();
            return status.isEmpty() ? null : status;
        } catch (Throwable t) {
            return null;
        }
    }

    public boolean isReusableCreatedContainer(String containerIdOrName) {
        String status = inspectContainerStatus(containerIdOrName);
        return "created".equals(status);
    }

    public List<Map<String, Object>> listSandboxContainers() throws SandboxException {
        List<String> cmd = new ArrayList<>();
        cmd.add(LobsterConfig.getSandboxDockerBin());
        cmd.add("ps");
        cmd.add("-a");
        cmd.add("--filter");
        cmd.add("name=lobster-sbx");
        cmd.add("--format");
        cmd.add("{{.ID}}\t{{.Names}}\t{{.Image}}\t{{.State}}\t{{.Status}}\t{{.CreatedAt}}");
        List<String> lines = runDockerLines(cmd, 10, "sandbox.docker_ps");
        List<Map<String, Object>> out = new ArrayList<>();
        for (String line : lines) {
            String[] p = line.split("\t", -1);
            if (p.length < 2) continue;
            Map<String, Object> row = new LinkedHashMap<>();
            row.put("id", p[0]);
            row.put("name", p[1]);
            row.put("image", p.length > 2 ? p[2] : "");
            row.put("state", p.length > 3 ? p[3] : "");
            row.put("status", p.length > 4 ? p[4] : "");
            row.put("createdAt", p.length > 5 ? p[5] : "");
            row.put("pool", p[1] != null && p[1].startsWith("lobster-sbx-pool-"));
            out.add(row);
        }
        return out;
    }

    public Map<String, Map<String, Object>> statsByContainerName(List<String> names) throws SandboxException {
        Map<String, Map<String, Object>> out = new LinkedHashMap<>();
        if (names == null || names.isEmpty()) return out;
        List<String> cmd = new ArrayList<>();
        cmd.add(LobsterConfig.getSandboxDockerBin());
        cmd.add("stats");
        cmd.add("--no-stream");
        cmd.add("--format");
        cmd.add("{{.Name}}\t{{.CPUPerc}}\t{{.MemUsage}}\t{{.MemPerc}}\t{{.NetIO}}\t{{.BlockIO}}\t{{.PIDs}}");
        cmd.addAll(names);
        List<String> lines = runDockerLines(cmd, 15, "sandbox.docker_stats");
        for (String line : lines) {
            String[] p = line.split("\t", -1);
            if (p.length < 1 || p[0].isEmpty()) continue;
            Map<String, Object> row = new LinkedHashMap<>();
            row.put("name", p[0]);
            row.put("cpuPerc", p.length > 1 ? p[1] : "");
            row.put("memUsage", p.length > 2 ? p[2] : "");
            row.put("memPerc", p.length > 3 ? p[3] : "");
            row.put("netIO", p.length > 4 ? p[4] : "");
            row.put("blockIO", p.length > 5 ? p[5] : "");
            row.put("pids", p.length > 6 ? p[6] : "");
            out.put(p[0], row);
        }
        return out;
    }

    private List<String> runDockerLines(List<String> cmd, long timeoutSec, String code) throws SandboxException {
        Process p;
        try {
            p = new ProcessBuilder(cmd).redirectErrorStream(false).start();
        } catch (IOException e) {
            throw new SandboxException(code, "failed to run docker: " + e.getMessage(), e);
        }
        StreamDrainer errDrain = new StreamDrainer(p.getErrorStream());
        StreamDrainer outDrain = new StreamDrainer(p.getInputStream());
        Thread tOut = new Thread(outDrain, "docker-query-out");
        Thread tErr = new Thread(errDrain, "docker-query-err");
        tOut.setDaemon(true);
        tErr.setDaemon(true);
        tOut.start();
        tErr.start();
        boolean done;
        try {
            done = p.waitFor(timeoutSec, TimeUnit.SECONDS);
        } catch (InterruptedException e) {
            p.destroyForcibly();
            Thread.currentThread().interrupt();
            throw new SandboxException(code, "docker query interrupted", e);
        }
        if (!done) {
            p.destroyForcibly();
            try { tOut.join(1000); } catch (InterruptedException ignore) { Thread.currentThread().interrupt(); }
            try { tErr.join(1000); } catch (InterruptedException ignore) { Thread.currentThread().interrupt(); }
            throw new SandboxException(code, "docker query timed out");
        }
        try { tOut.join(1000); } catch (InterruptedException ignore) { Thread.currentThread().interrupt(); }
        try { tErr.join(1000); } catch (InterruptedException ignore) { Thread.currentThread().interrupt(); }
        if (p.exitValue() != 0) {
            throw new SandboxException(code,
                    "docker query failed (exit=" + p.exitValue() + "): " + errDrain.getText());
        }
        List<String> lines = new ArrayList<>();
        String text = outDrain.getText();
        if (text != null && !text.isEmpty()) {
            try (BufferedReader r = new BufferedReader(new java.io.StringReader(text))) {
                String line;
                while ((line = r.readLine()) != null) lines.add(line);
            } catch (IOException ignore) { /* StringReader does not throw in practice */ }
        }
        return lines;
    }

    private static String shortId(String id) {
        return id == null || id.length() < 12 ? id : id.substring(0, 12);
    }

    /** 简单的流排空器 / Minimal stream drainer. */
    private static final class StreamDrainer implements Runnable {
        private static final int MAX_BUF = 256 * 1024; // 256KB 上限，防 stdout/stderr 打爆
        private final InputStream in;
        private final StringBuilder buf = new StringBuilder();
        private volatile boolean truncated;

        StreamDrainer(InputStream in) { this.in = in; }

        @Override
        public void run() {
            // 必须把管道读到 EOF，否则容器 stdout/stderr 写满后会 block 住主进程.
            // 达到 MAX_BUF 之后切换到"读即丢弃"模式，保证 docker 子进程能持续输出直到退出.
            try (BufferedReader r = new BufferedReader(new InputStreamReader(in, StandardCharsets.UTF_8))) {
                String line;
                while ((line = r.readLine()) != null) {
                    if (buf.length() < MAX_BUF) {
                        buf.append(line).append('\n');
                    } else {
                        truncated = true;
                    }
                }
            } catch (IOException ignore) { /* process may be killed mid-read */ }
        }

        String getText() {
            if (truncated) return buf.toString() + "\n... [truncated]";
            return buf.toString();
        }
    }
}