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

import io.agentscope.core.message.TextBlock;
import io.agentscope.core.message.ToolResultBlock;
import io.agentscope.core.tool.ToolCallParam;
import org.apache.commons.compress.archivers.tar.TarArchiveEntry;
import org.apache.commons.compress.archivers.tar.TarArchiveOutputStream;
import org.apache.commons.compress.archivers.zip.ZipArchiveEntry;
import org.apache.commons.compress.archivers.zip.ZipArchiveOutputStream;
import org.junit.Assert;
import org.junit.Test;

import java.io.IOException;
import java.io.OutputStream;
import java.nio.charset.StandardCharsets;
import java.nio.file.Files;
import java.nio.file.Path;
import java.time.Duration;
import java.util.Map;
import java.util.List;
import java.util.concurrent.TimeUnit;

/**
 * 测试 Java 安全归档执行器的创建、展开与恶意条目边界。
 */
public class SafeArchiveCommandExecutorTest {

    @Test
    public void shouldCreateAndExtractGzipZipAndTarArchives() throws IOException {
        Fixture fixture = fixture(new WorkspaceQuotaLimits(16 * 1024 * 1024L,
                8 * 1024 * 1024L, 1000, 1024 * 1024));
        Files.createDirectories(fixture.root().resolve("source/nested"));
        Files.writeString(fixture.root().resolve("source/nested/value.txt"), "archive-value\n");
        Files.writeString(fixture.root().resolve("plain.txt"), "plain-value\n");
        Files.writeString(fixture.root().resolve("multi-a.txt"), "a\n");
        Files.writeString(fixture.root().resolve("multi-b.txt"), "b\n");

        assertSuccess(fixture, "gzip -k plain.txt");
        Files.delete(fixture.root().resolve("plain.txt"));
        assertSuccess(fixture, "gunzip -k plain.txt.gz");
        Assert.assertEquals("plain-value\n", Files.readString(fixture.root().resolve("plain.txt")));
        assertSuccess(fixture, "gzip multi-a.txt multi-b.txt");
        Assert.assertFalse(Files.exists(fixture.root().resolve("multi-a.txt")));
        Assert.assertFalse(Files.exists(fixture.root().resolve("multi-b.txt")));
        assertSuccess(fixture, "gunzip multi-a.txt.gz multi-b.txt.gz");
        Assert.assertEquals("a\n", Files.readString(fixture.root().resolve("multi-a.txt")));
        Assert.assertEquals("b\n", Files.readString(fixture.root().resolve("multi-b.txt")));

        assertSuccess(fixture, "zip -q -r bundle.zip source");
        assertSuccess(fixture, "unzip -q bundle.zip -d zip-output");
        Assert.assertEquals("archive-value\n",
                Files.readString(fixture.root().resolve("zip-output/source/nested/value.txt")));

        assertSuccess(fixture, "tar -czf bundle.tar.gz source");
        String listed = execute(fixture, "tar -tzf bundle.tar.gz");
        Assert.assertTrue(listed, listed.contains("source/nested/value.txt"));
        Assert.assertFalse(listed.contains(fixture.root().toString()));
        assertSuccess(fixture, "tar -xzf bundle.tar.gz -C tar-output");
        Assert.assertEquals("archive-value\n",
                Files.readString(fixture.root().resolve("tar-output/source/nested/value.txt")));
    }

    @Test
    public void shouldRejectZipSlipDuplicateAndTargetConflictWithoutPartialOutput() throws IOException {
        Fixture fixture = fixture(new WorkspaceQuotaLimits(8 * 1024 * 1024L,
                4 * 1024 * 1024L, 1000, 1024 * 1024));
        createZip(fixture.root().resolve("slip.zip"),
                new ZipContent("../escape.txt", "escape"));
        createZip(fixture.root().resolve("duplicate.zip"),
                new ZipContent("same.txt", "first"), new ZipContent("same.txt", "second"));
        createZip(fixture.root().resolve("backslash.zip"),
                new ZipContent("..\\escape.txt", "escape"));
        createZip(fixture.root().resolve("conflict.zip"),
                new ZipContent("first.txt", "first"), new ZipContent("second.txt", "second"));
        Files.createDirectories(fixture.root().resolve("output"));
        Files.writeString(fixture.root().resolve("output/second.txt"), "existing");

        Assert.assertTrue(execute(fixture, "unzip slip.zip -d slip-output")
                .contains("ARCHIVE_ENTRY_DENIED"));
        Assert.assertFalse(Files.exists(fixture.root().getParent().resolve("escape.txt")));
        Assert.assertTrue(execute(fixture, "unzip duplicate.zip -d duplicate-output")
                .contains("ARCHIVE_DUPLICATE_ENTRY"));
        Assert.assertTrue(execute(fixture, "unzip backslash.zip -d backslash-output")
                .contains("ARCHIVE_ENTRY_DENIED"));
        Assert.assertTrue(execute(fixture, "unzip conflict.zip -d output")
                .contains("ARCHIVE_TARGET_CONFLICT"));
        Assert.assertFalse(Files.exists(fixture.root().resolve("output/first.txt")));
        Assert.assertEquals("existing", Files.readString(fixture.root().resolve("output/second.txt")));
    }

    @Test
    public void shouldRejectTarLinksDevicesAndFifoEntries() throws IOException {
        Fixture fixture = fixture(new WorkspaceQuotaLimits(8 * 1024 * 1024L,
                4 * 1024 * 1024L, 1000, 1024 * 1024));
        createTarSpecial(fixture.root().resolve("link.tar"), "link", (byte) '2');
        createTarSpecial(fixture.root().resolve("hard-link.tar"), "hard-link", (byte) '1');
        createTarSpecial(fixture.root().resolve("device.tar"), "device", (byte) '3');
        createTarSpecial(fixture.root().resolve("fifo.tar"), "fifo", (byte) '6');

        Assert.assertTrue(execute(fixture, "tar -xf link.tar").contains("ARCHIVE_ENTRY_DENIED"));
        Assert.assertTrue(execute(fixture, "tar -xf hard-link.tar").contains("ARCHIVE_ENTRY_DENIED"));
        Assert.assertTrue(execute(fixture, "tar -xf device.tar").contains("ARCHIVE_ENTRY_DENIED"));
        Assert.assertTrue(execute(fixture, "tar -xf fifo.tar").contains("ARCHIVE_ENTRY_DENIED"));
        Assert.assertFalse(Files.exists(fixture.root().resolve("link")));
        Assert.assertFalse(Files.exists(fixture.root().resolve("device")));
        Assert.assertFalse(Files.exists(fixture.root().resolve("fifo")));
    }

    @Test
    public void shouldRejectExpandedArchiveBeyondQuota() throws IOException {
        Fixture fixture = fixture(new WorkspaceQuotaLimits(512, 8 * 1024, 100, 1024));
        createZip(fixture.root().resolve("bomb.zip"),
                new ZipContent("expanded.txt", "0".repeat(4096)));

        String output = execute(fixture, "unzip bomb.zip -d expanded");

        Assert.assertTrue(output, output.contains("WORKSPACE_QUOTA_EXCEEDED"));
        Assert.assertFalse(Files.exists(fixture.root().resolve("expanded")));
    }

    @Test
    public void shouldRejectUnknownArchiveOptions() throws IOException {
        Fixture fixture = fixture(new WorkspaceQuotaLimits(1024 * 1024, 1024 * 1024, 100, 1024));
        Files.writeString(fixture.root().resolve("input.txt"), "input");

        for (String command : new String[]{
                "gzip -c input.txt",
                "zip -T bundle.zip input.txt",
                "unzip -o bundle.zip",
                "tar --checkpoint-action=exec=id -cf bundle.tar input.txt",
                "tar -xf bundle.tar member.txt"}) {
            String output = execute(fixture, command);
            Assert.assertTrue(command + " => " + output, output.contains("SHELL_COMMAND_DENIED"));
        }
    }

    @Test
    public void shouldCompensateOutputsWhenCommitFailsAfterFirstMove() throws IOException {
        Path root = Files.createTempDirectory("safe-archive-compensation-").toAbsolutePath();
        createZip(root.resolve("two-files.zip"),
                new ZipContent("first.txt", "first"), new ZipContent("second.txt", "second"));
        WorkspacePathGuard pathGuard = new WorkspacePathGuard(root);
        WorkspaceQuotaGuard quotaGuard = new WorkspaceQuotaGuard(
                pathGuard,
                new WorkspaceQuotaLimits(1024 * 1024, 1024 * 1024, 100, 1024),
                WorkspaceQuotaHook.noop());
        SafeArchiveCommandExecutor executor = new SafeArchiveCommandExecutor(
                pathGuard, quotaGuard, 1024, committedCount -> {
                    if (committedCount == 1) {
                        throw new IllegalStateException("injected commit failure");
                    }
                });

        try {
            executor.execute(List.of("unzip", "two-files.zip", "-d", "output"),
                    System.nanoTime() + TimeUnit.SECONDS.toNanos(5));
            Assert.fail("Expected archive commit failure");
        } catch (WorkspaceToolException expected) {
            Assert.assertEquals("ARCHIVE_COMMIT_FAILED", expected.code());
        }

        Assert.assertFalse(Files.exists(root.resolve("output/first.txt")));
        Assert.assertFalse(Files.exists(root.resolve("output/second.txt")));
        Assert.assertFalse(Files.exists(root.resolve("output")));
    }

    private Fixture fixture(WorkspaceQuotaLimits limits) throws IOException {
        Path root = Files.createTempDirectory("safe-archive-").toAbsolutePath();
        WorkspacePathGuard pathGuard = new WorkspacePathGuard(root);
        WorkspaceQuotaGuard quotaGuard = new WorkspaceQuotaGuard(pathGuard, limits, WorkspaceQuotaHook.noop());
        AgentOperateToolSpec spec = new AgentOperateToolSpec();
        spec.setShellDefaultTimeout(Duration.ofSeconds(5));
        spec.setShellMaxTimeout(Duration.ofSeconds(10));
        return new Fixture(root, new ControlledShellTool(pathGuard, quotaGuard, spec));
    }

    private void assertSuccess(Fixture fixture, String command) {
        String output = execute(fixture, command);
        Assert.assertTrue(command + " => " + output, output.contains("<returncode>0</returncode>"));
        Assert.assertFalse(output.contains(fixture.root().toString()));
    }

    private String execute(Fixture fixture, String command) {
        ToolResultBlock result = fixture.tool().callAsync(ToolCallParam.builder()
                .input(Map.of("command", command)).build()).block();
        return ((TextBlock) result.getOutput().get(0)).getText();
    }

    private void createZip(Path target, ZipContent... contents) throws IOException {
        try (ZipArchiveOutputStream zip = new ZipArchiveOutputStream(target)) {
            for (ZipContent content : contents) {
                byte[] bytes = content.content().getBytes(StandardCharsets.UTF_8);
                ZipArchiveEntry entry = new ZipArchiveEntry(content.name());
                entry.setSize(bytes.length);
                zip.putArchiveEntry(entry);
                zip.write(bytes);
                zip.closeArchiveEntry();
            }
            zip.finish();
        }
    }

    private void createTarSpecial(Path target, String name, byte linkFlag) throws IOException {
        try (OutputStream raw = Files.newOutputStream(target);
             TarArchiveOutputStream tar = new TarArchiveOutputStream(raw)) {
            TarArchiveEntry entry = new TarArchiveEntry(name, linkFlag);
            entry.setSize(0);
            if (linkFlag == '2') {
                entry.setLinkName("../outside");
            }
            tar.putArchiveEntry(entry);
            tar.closeArchiveEntry();
            tar.finish();
        }
    }

    private record ZipContent(String name, String content) {
    }

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