package com.easyagents.agent.runtime.tool.operate;

import com.easyagents.agent.runtime.hitl.AgentToolApprovalEvaluation;
import io.agentscope.core.message.TextBlock;
import io.agentscope.core.message.ToolResultBlock;
import io.agentscope.core.tool.ToolCallParam;
import org.junit.Assert;
import org.junit.Assume;
import org.junit.Test;

import java.io.IOException;
import java.nio.file.Files;
import java.nio.file.Path;
import java.time.Duration;
import java.util.Map;
import java.util.Optional;
import java.util.Set;
import java.util.concurrent.TimeUnit;

/**
 * 测试受控 Shell 策略与执行边界。
 */
public class ControlledShellToolTest {

    @Test
    public void shouldExecuteAllowlistedCommandWithoutLeakingWorkspaceRoot() throws IOException {
        Fixture fixture = fixture();

        String result = execute(fixture.tool(), Map.of("command", "pwd"));

        Assert.assertTrue(result, result.contains("<returncode>0</returncode>"));
        Assert.assertFalse(result.contains(fixture.root().toString()));
        Assert.assertTrue(result, result.contains("<stdout truncated=\"false\">.\n</stdout>"));
    }

    @Test
    public void shouldRejectOperatorsExpansionAbsolutePathsAndUnknownCommands() throws IOException {
        Fixture fixture = fixture();

        assertRejected(fixture.tool(), "pwd | cat");
        assertRejected(fixture.tool(), "cat $HOME/secret");
        assertRejected(fixture.tool(), "cat /etc/passwd");
        assertRejected(fixture.tool(), "curl https://example.com");
        assertRejected(fixture.tool(), "pwd\ncat secret");
    }

    @Test
    public void shouldRestrictScriptEntrypointsAndDangerousRemove() throws IOException {
        Fixture fixture = fixture();
        Files.writeString(fixture.root().resolve("safe.py"), "print('ok')\n");

        assertRejected(fixture.tool(), "python3 -c 'print(1)'");
        assertRejected(fixture.tool(), "python3 -m http.server");
        assertRejected(fixture.tool(), "node --eval '1+1'");
        assertRejected(fixture.tool(), "rm -rf .");
        assertRejected(fixture.tool(), "rm -r -f output");
        assertRejected(fixture.tool(), "rm --recursive --force output");
        Assert.assertTrue(execute(fixture.tool(), Map.of("command", "python3 safe.py"))
                .contains("<stdout truncated=\"false\">ok"));
    }

    @Test
    public void shouldExposeOnlyFixedArchiveCommandSet() throws IOException {
        Fixture fixture = fixture();

        assertRejected(fixture.tool(), "tar --checkpoint-action=exec=sh -cf archive.tar input.txt");
        assertRejected(fixture.tool(), "zip -TT sh archive.zip input.txt");
        Assert.assertTrue(ControlledShellTool.DEFAULT_ALLOWED_COMMANDS.containsAll(
                Set.of("gzip", "gunzip", "zip", "unzip", "tar")));
    }

    @Test
    public void shouldClassifyReadOnlyWriteConversionAndScriptApproval() throws IOException {
        Fixture fixture = fixture();
        Files.writeString(fixture.root().resolve("safe.py"), "print('one')\n");
        Files.writeString(fixture.root().resolve("input.pdf"), "%PDF-1.4\n");
        Files.writeString(fixture.root().resolve("input.md"), "# title\n");
        Files.createDirectory(fixture.root().resolve("output"));

        assertApproval(fixture.tool(), "rg --files .", false, false, null);
        assertApproval(fixture.tool(), "tree .", false, false, null);
        assertApproval(fixture.tool(), "pdftotext input.pdf -", false, false, null);
        assertApproval(fixture.tool(), "mkdir generated", true, false, null);
        assertApproval(fixture.tool(), "rm generated.txt", true, true, null);
        assertApproval(fixture.tool(), "pandoc input.md -o output/report.docx", true, false, null);
        assertApproval(fixture.tool(), "soffice --convert-to pdf --outdir output input.md",
                true, false, null);
        assertApproval(fixture.tool(), "pdftotext input.pdf output/content.txt", true, false, null);

        AgentToolApprovalEvaluation first = fixture.tool().approvalEvaluation(
                Map.of("command", "python3 safe.py --mode report"));
        Assert.assertTrue(first.valid());
        Assert.assertTrue(first.approvalRequired());
        Assert.assertFalse(first.forced());
        Assert.assertTrue(first.reusableScope().startsWith("SHELL_SCRIPT:"));

        Files.writeString(fixture.root().resolve("safe.py"), "print('two')\n");
        AgentToolApprovalEvaluation changed = fixture.tool().approvalEvaluation(
                Map.of("command", "python3 safe.py --mode report"));
        Assert.assertNotEquals(first.reusableScope(), changed.reusableScope());
    }

    @Test
    public void shouldRejectUnsafeProductivityCommandOptionsBeforeApproval() throws IOException {
        Fixture fixture = fixture();
        Files.writeString(fixture.root().resolve("input.md"), "# title\n");
        Files.writeString(fixture.root().resolve("input.pdf"), "%PDF-1.4\n");
        Files.writeString(fixture.root().resolve("paths.txt"), "input.md\n");
        Files.writeString(fixture.root().resolve("args.txt"), "--empty\n");
        Files.createDirectory(fixture.root().resolve("output"));

        for (String command : new String[]{
                "find .",
                "tree -l .",
                "du --files0-from=paths.txt",
                "pandoc input.md --filter cat -o output/report.docx",
                "pandoc input.md -ooutput/report.docx",
                "pandoc https://example.com -o output/report.docx",
                "soffice --accept=socket --convert-to pdf --outdir output input.md",
                "pdftoppm /etc/passwd output/page",
                "qpdf @args.txt output/result.pdf"}) {
            AgentToolApprovalEvaluation evaluation = fixture.tool().approvalEvaluation(
                    Map.of("command", command));
            Assert.assertFalse(command, evaluation.valid());
            Assert.assertFalse(command, evaluation.approvalRequired());
        }
    }

    @Test
    public void shouldEnforceTimeoutAndOutputLimit() throws IOException {
        Fixture fixture = fixture();
        Files.writeString(fixture.root().resolve("slow.py"), "import time\ntime.sleep(5)\n");
        Files.writeString(fixture.root().resolve("large.py"), "print('x' * 10000)\n");

        String timeout = execute(fixture.tool(), Map.of("command", "python3 slow.py", "timeout", 1));
        String truncated = execute(fixture.tool(), Map.of("command", "python3 large.py"));

        Assert.assertTrue(timeout.contains("SHELL_TIMEOUT"));
        Assert.assertTrue(truncated.contains("truncated=\"true\""));
        Assert.assertTrue(truncated.contains("OUTPUT_TRUNCATED"));
    }

    @Test
    public void shouldRejectSecondaryExecutionAndIndirectFileOptions() throws IOException {
        Fixture fixture = fixture();
        Files.writeString(fixture.root().resolve("input.txt"), "alpha\n");
        Files.writeString(fixture.root().resolve("input.json"), "{\"value\":1}\n");
        Files.writeString(fixture.root().resolve("paths.txt"), "/etc/passwd\n");

        for (String command : new String[]{
                "awk 'BEGIN { system(\"id\") }' input.txt",
                "awk '{ getline value }' input.txt",
                "awk '{ print ENVIRON }' input.txt",
                "awk -fprogram.awk input.txt",
                "awk --profile=profile.txt '{ print }' input.txt",
                "sed -e 'e id' input.txt",
                "sed '1r input.txt' input.txt",
                "sed 's/alpha/beta/w stolen.txt' input.txt",
                "sed -i 's/alpha/beta/' input.txt",
                "rg --pre cat alpha .",
                "rg --pre-glob '*.txt' alpha .",
                "rg --hostname-bin pwd alpha .",
                "rg -z alpha .",
                "rg -L alpha .",
                "rg --follow alpha .",
                "ls -RL .",
                "grep -Rfpatterns.txt alpha .",
                "jq 'env' input.json",
                "jq '$ENV' input.json",
                "jq -Lmodules '.' input.json",
                "sort -ooutput.txt input.txt",
                "sort --compress-program=cat input.txt",
                "uniq input.txt output.txt",
                "file -fpaths.txt",
                "sha256sum --check paths.txt",
                "wc --files0-from=paths.txt",
                "tail --follow=name input.txt",
                "cp -L input.txt copied.txt",
                "cp --symbolic-link input.txt copied.txt",
                "cp -l input.txt copied.txt"}) {
            assertRejected(fixture.tool(), command);
        }
    }

    @Test
    public void shouldBestEffortTerminateObservedChildAfterSuccessfulScript() throws Exception {
        Fixture fixture = fixture();
        Files.writeString(fixture.root().resolve("spawn.py"), """
                import pathlib
                import subprocess
                import sys
                import time
                child = subprocess.Popen(
                    [sys.executable, '-c', 'import time; time.sleep(30)'],
                    stdout=subprocess.DEVNULL,
                    stderr=subprocess.DEVNULL,
                    start_new_session=True)
                pathlib.Path('child.pid').write_text(str(child.pid), encoding='utf-8')
                time.sleep(0.3)
                """);
        ProcessHandle child = null;
        try {
            String result = execute(fixture.tool(), Map.of("command", "python3 spawn.py"));
            Assert.assertTrue(result, result.contains("<returncode>0</returncode>"));
            long pid = Long.parseLong(Files.readString(fixture.root().resolve("child.pid")).trim());
            Optional<ProcessHandle> handle = ProcessHandle.of(pid);
            child = handle.orElse(null);
            long deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(2);
            while (child != null && child.isAlive() && System.nanoTime() < deadline) {
                Thread.sleep(20);
            }
            Assert.assertTrue("Observed child process must be terminated", child == null || !child.isAlive());
        } finally {
            if (child != null && child.isAlive()) {
                child.destroyForcibly();
            }
        }
    }

    @Test
    public void shouldTerminateLinuxProcessGroupAfterFastParentExit() throws Exception {
        assumeLinuxWithPython();
        Fixture fixture = fixture();
        Files.writeString(fixture.root().resolve("fast-spawn.py"), """
                import pathlib
                import subprocess
                import sys
                child = subprocess.Popen(
                    [sys.executable, '-c', 'import time; time.sleep(30)'],
                    stdout=subprocess.DEVNULL,
                    stderr=subprocess.DEVNULL)
                pathlib.Path('fast-child.pid').write_text(str(child.pid), encoding='utf-8')
                """);
        ProcessHandle child = null;
        try {
            String result = execute(fixture.tool(), Map.of("command", "python3 fast-spawn.py"));
            Assert.assertTrue(result, result.contains("<returncode>0</returncode>"));
            child = ProcessHandle.of(readPid(fixture.root().resolve("fast-child.pid"))).orElse(null);
            assertTerminates(child, "Linux process-group child must be terminated after parent exit");
        } finally {
            destroyIfAlive(child);
        }
    }

    @Test
    public void shouldTerminateLinuxProcessGroupOnTimeout() throws Exception {
        assumeLinuxWithPython();
        Fixture fixture = fixture();
        Files.writeString(fixture.root().resolve("timeout-spawn.py"), """
                import pathlib
                import subprocess
                import sys
                import time
                child = subprocess.Popen(
                    [sys.executable, '-c', 'import time; time.sleep(30)'],
                    stdout=subprocess.DEVNULL,
                    stderr=subprocess.DEVNULL)
                pathlib.Path('timeout-child.pid').write_text(str(child.pid), encoding='utf-8')
                time.sleep(30)
                """);
        ProcessHandle child = null;
        try {
            String result = execute(fixture.tool(), Map.of(
                    "command", "python3 timeout-spawn.py", "timeout", 1));
            Assert.assertTrue(result, result.contains("SHELL_TIMEOUT"));
            child = ProcessHandle.of(readPid(fixture.root().resolve("timeout-child.pid"))).orElse(null);
            assertTerminates(child, "Linux process-group child must be terminated on timeout");
        } finally {
            destroyIfAlive(child);
        }
    }

    private void assumeLinuxWithPython() {
        Assume.assumeTrue(System.getProperty("os.name", "").toLowerCase().contains("linux"));
        Assume.assumeTrue(Files.isExecutable(Path.of("/usr/bin/python3"))
                || Files.isExecutable(Path.of("/usr/local/bin/python3")));
    }

    private long readPid(Path path) throws Exception {
        long deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(2);
        while (!Files.exists(path) && System.nanoTime() < deadline) {
            Thread.sleep(10);
        }
        Assert.assertTrue("Child PID file must be created", Files.exists(path));
        return Long.parseLong(Files.readString(path).trim());
    }

    private void assertTerminates(ProcessHandle child, String message) throws InterruptedException {
        long deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(2);
        while (child != null && child.isAlive() && System.nanoTime() < deadline) {
            Thread.sleep(20);
        }
        Assert.assertTrue(message, child == null || !child.isAlive());
    }

    private void destroyIfAlive(ProcessHandle process) {
        if (process != null && process.isAlive()) {
            process.destroyForcibly();
        }
    }

    private Fixture fixture() throws IOException {
        Path root = Files.createTempDirectory("controlled-shell-").toAbsolutePath();
        WorkspacePathGuard pathGuard = new WorkspacePathGuard(root);
        WorkspaceQuotaGuard quotaGuard = new WorkspaceQuotaGuard(
                pathGuard,
                new WorkspaceQuotaLimits(1024 * 1024, 1024 * 1024, 100, 1024 * 1024),
                WorkspaceQuotaHook.noop());
        AgentOperateToolSpec spec = new AgentOperateToolSpec();
        spec.setShellDefaultTimeout(Duration.ofSeconds(2));
        spec.setShellMaxTimeout(Duration.ofSeconds(3));
        spec.setShellMaxOutputSize(256);
        return new Fixture(root, new ControlledShellTool(pathGuard, quotaGuard, spec));
    }

    private String execute(ControlledShellTool tool, Map<String, Object> input) {
        ToolResultBlock result = tool.callAsync(ToolCallParam.builder().input(input).build()).block();
        return ((TextBlock) result.getOutput().get(0)).getText();
    }

    private void assertRejected(ControlledShellTool tool, String command) {
        String result = execute(tool, Map.of("command", command));
        Assert.assertTrue(result, result.contains("SHELL_COMMAND_DENIED")
                || result.contains("WORKSPACE_PATH_INVALID"));
    }

    private void assertApproval(ControlledShellTool tool,
                                String command,
                                boolean required,
                                boolean forced,
                                String scope) {
        AgentToolApprovalEvaluation evaluation = tool.approvalEvaluation(Map.of("command", command));
        Assert.assertTrue(command, evaluation.valid());
        Assert.assertEquals(command, required, evaluation.approvalRequired());
        Assert.assertEquals(command, forced, evaluation.forced());
        Assert.assertEquals(command, scope, evaluation.reusableScope());
    }

    /**
     * Shell 测试夹具。
     *
     * @param root 工作区根
     * @param tool Shell 工具
     */
    private record Fixture(Path root, ControlledShellTool tool) {
    }
}
