package com.easyagents.skill.codec;

import com.easyagents.skill.exception.SkillPackageException;
import com.easyagents.skill.model.Skill;
import com.easyagents.skill.store.SkillContentStore;
import com.easyagents.skill.store.memory.InMemorySkillContentStore;
import org.junit.Assert;
import org.junit.Test;

import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.InputStream;
import java.nio.charset.StandardCharsets;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
import java.util.zip.ZipEntry;
import java.util.zip.ZipOutputStream;

/**
 * ZipSkillPackageCodec 单元测试。
 */
public class ZipSkillPackageCodecTest {

    /**
     * 导入包含一个 Skill 文件夹的 zip。
     */
    @Test
    public void importOneSkillFolder() {
        List<Skill> skills = importZip(files(
                "skill-a/SKILL.md", skillMd("Skill A", "Desc A"),
                "skill-a/references/rules/a.md", "# Rule",
                "skill-a/scripts/run.py", "print('ok')",
                "skill-a/assets/images/logo.png", "png-data"
        ));

        Assert.assertEquals(1, skills.size());
        Skill skill = skills.get(0);
        Assert.assertEquals("skill-a", skill.getId());
        Assert.assertEquals("Skill A", skill.getName());
        Assert.assertEquals("Desc A", skill.getDescription());
        Assert.assertEquals(1, skill.getReferences().size());
        Assert.assertEquals("references/rules/a.md", skill.getReferences().get(0).getPath());
        Assert.assertEquals(1, skill.getScripts().size());
        Assert.assertEquals(1, skill.getAssets().size());
        Assert.assertEquals("assets/images/logo.png", skill.getAssets().get(0).getPath());
    }

    /**
     * 导入包含多个 Skill 文件夹的 zip。
     */
    @Test
    public void importMultipleSkillFolders() {
        List<Skill> skills = importZip(files(
                "skill-a/SKILL.md", skillMd("Skill A", "Desc A"),
                "skill-b/SKILL.md", skillMd("Skill B", "Desc B")
        ));

        Assert.assertEquals(2, skills.size());
        Assert.assertEquals("skill-a", skills.get(0).getId());
        Assert.assertEquals("skill-b", skills.get(1).getId());
    }

    /**
     * 拒绝 zip 根目录直接包含 SKILL.md。
     */
    @Test(expected = SkillPackageException.class)
    public void rejectRootSkillMd() {
        importZip(files("SKILL.md", skillMd("Skill", "Desc")));
    }

    /**
     * 拒绝未知顶级目录。
     */
    @Test(expected = SkillPackageException.class)
    public void rejectUnknownTopLevelDirectory() {
        importZip(files(
                "skill-a/SKILL.md", skillMd("Skill A", "Desc A"),
                "skill-a/unknown/a.md", "# Unknown"
        ));
    }

    /**
     * 拒绝 references 下的非 md 文件。
     */
    @Test(expected = SkillPackageException.class)
    public void rejectNonMarkdownReference() {
        importZip(files(
                "skill-a/SKILL.md", skillMd("Skill A", "Desc A"),
                "skill-a/references/a.txt", "text"
        ));
    }

    /**
     * 拒绝不支持的脚本扩展名。
     */
    @Test(expected = SkillPackageException.class)
    public void rejectUnsupportedScript() {
        importZip(files(
                "skill-a/SKILL.md", skillMd("Skill A", "Desc A"),
                "skill-a/scripts/run.rb", "puts 'ok'"
        ));
    }

    /**
     * assets 允许任意扩展名。
     */
    @Test
    public void allowArbitraryAssetExtension() {
        List<Skill> skills = importZip(files(
                "skill-a/SKILL.md", skillMd("Skill A", "Desc A"),
                "skill-a/assets/model.custom", "asset"
        ));

        Assert.assertEquals(1, skills.get(0).getAssets().size());
        Assert.assertEquals("application/octet-stream", skills.get(0).getAssets().get(0).getMediaType());
    }

    /**
     * 忽略 macOS 系统文件和空目录。
     */
    @Test
    public void ignoreSystemFilesAndEmptyDirectories() {
        Map<String, String> files = files(
                "__MACOSX/._x", "ignored",
                "skill-a/.DS_Store", "ignored",
                "skill-a/SKILL.md", skillMd("Skill A", "Desc A")
        );
        List<Skill> skills = new ZipSkillPackageCodec().importZip(new ByteArrayInputStream(zip(files, "skill-a/references/")));

        Assert.assertEquals(1, skills.size());
        Assert.assertEquals("skill-a", skills.get(0).getId());
    }

    /**
     * 嵌套 references 和 assets 路径正常导入。
     */
    @Test
    public void allowNestedManagedDirectories() {
        List<Skill> skills = importZip(files(
                "skill-a/SKILL.md", skillMd("Skill A", "Desc A"),
                "skill-a/references/rules/a.md", "# Nested",
                "skill-a/assets/images/logo.png", "logo"
        ));

        Assert.assertEquals("references/rules/a.md", skills.get(0).getReferences().get(0).getPath());
        Assert.assertEquals("assets/images/logo.png", skills.get(0).getAssets().get(0).getPath());
    }

    /**
     * 资产字节写入内容存储。
     */
    @Test
    public void storeAssetContent() {
        InMemorySkillContentStore store = new InMemorySkillContentStore();
        List<Skill> skills = new ZipSkillPackageCodec(store).importZip(new ByteArrayInputStream(zip(files(
                "skill-a/SKILL.md", skillMd("Skill A", "Desc A"),
                "skill-a/assets/data.bin", "abc"
        ))));

        String contentRef = skills.get(0).getAssets().get(0).getContentRef();
        Assert.assertTrue(store.exists(contentRef));
        Assert.assertArrayEquals("abc".getBytes(StandardCharsets.UTF_8), store.readAllBytes(contentRef));
    }

    /**
     * 导入失败时不写入资产内容。
     */
    @Test
    public void failedImportDoesNotWriteAssetContent() {
        CountingContentStore store = new CountingContentStore();

        try {
            new ZipSkillPackageCodec(store).importZip(new ByteArrayInputStream(zip(files(
                    "skill-a/SKILL.md", skillMd("Skill A", "Desc A"),
                    "skill-a/assets/data.bin", "abc",
                    "skill-a/unknown/a.md", "# Unknown"
            ))));
            Assert.fail("Import should fail.");
        } catch (SkillPackageException expected) {
            Assert.assertEquals(0, store.putCount);
        }
    }

    /**
     * 多 Skill 包中后续 Skill 失败时，不提前写入前面 Skill 的资产。
     */
    @Test
    public void failedMultiSkillImportDoesNotWriteEarlierAssetContent() {
        CountingContentStore store = new CountingContentStore();

        try {
            new ZipSkillPackageCodec(store).importZip(new ByteArrayInputStream(zip(files(
                    "skill-a/SKILL.md", skillMd("Skill A", "Desc A"),
                    "skill-a/assets/data.bin", "abc",
                    "skill-b/references/a.md", "# Missing SKILL.md"
            ))));
            Assert.fail("Import should fail.");
        } catch (SkillPackageException expected) {
            Assert.assertEquals(0, store.putCount);
        }
    }

    /**
     * 拒绝嵌套 frontmatter。
     */
    @Test(expected = SkillPackageException.class)
    public void rejectNestedFrontmatter() {
        importZip(files(
                "skill-a/SKILL.md", "---\nname: Skill A\ndescription: Desc A\nconfig:\n  mode: strict\n---\n# Skill A\n"
        ));
    }

    private static List<Skill> importZip(Map<String, String> files) {
        return new ZipSkillPackageCodec().importZip(new ByteArrayInputStream(zip(files)));
    }

    private static Map<String, String> files(String... keyValues) {
        Map<String, String> files = new LinkedHashMap<>();
        for (int i = 0; i < keyValues.length; i += 2) {
            files.put(keyValues[i], keyValues[i + 1]);
        }
        return files;
    }

    private static byte[] zip(Map<String, String> files, String... directories) {
        try {
            ByteArrayOutputStream bytes = new ByteArrayOutputStream();
            try (ZipOutputStream zipOutputStream = new ZipOutputStream(bytes)) {
                for (String directory : directories) {
                    zipOutputStream.putNextEntry(new ZipEntry(directory));
                    zipOutputStream.closeEntry();
                }
                for (Map.Entry<String, String> file : files.entrySet()) {
                    zipOutputStream.putNextEntry(new ZipEntry(file.getKey()));
                    zipOutputStream.write(file.getValue().getBytes(StandardCharsets.UTF_8));
                    zipOutputStream.closeEntry();
                }
            }
            return bytes.toByteArray();
        } catch (Exception e) {
            throw new IllegalStateException(e);
        }
    }

    private static String skillMd(String name, String description) {
        return "---\nname: " + name + "\ndescription: " + description + "\n---\n# " + name + "\n";
    }

    private static final class CountingContentStore implements SkillContentStore {

        private int putCount;

        @Override
        public String put(byte[] bytes) {
            putCount++;
            return "sha256:" + com.easyagents.skill.util.SkillHashes.sha256Hex(bytes);
        }

        @Override
        public InputStream open(String contentRef) {
            throw new UnsupportedOperationException();
        }

        @Override
        public byte[] readAllBytes(String contentRef) {
            throw new UnsupportedOperationException();
        }

        @Override
        public boolean exists(String contentRef) {
            return false;
        }
    }
}