package com.easyagents.document.pdf.mineru;

import com.alibaba.fastjson2.JSONArray;
import com.alibaba.fastjson2.JSONObject;
import com.easyagents.document.core.mineru.MineruMapper;
import com.easyagents.document.core.mineru.MineruProperties;
import com.easyagents.document.core.mineru.MineruResultPayload;
import com.easyagents.document.core.exception.DocumentParseException;
import com.easyagents.document.core.entity.ParseRequest;
import com.easyagents.document.core.entity.ParseResponse;
import com.easyagents.document.core.entity.ParseResult;
import org.junit.Assert;
import org.junit.Test;

import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.util.List;
import java.util.Map;
import java.util.zip.ZipEntry;
import java.util.zip.ZipOutputStream;

/**
 * MinerU 结果映射测试。
 *
 * @author Codex
 * @since 2026-04-14
 */
public class MineruMapperTest {

    @Test
    public void shouldMapSyncResponse() {
        MineruMapper mapper = new MineruMapper(defaultProperties());
        MineruResultPayload payload = mapper.toResultPayload(syncPayload());

        ParseResponse response = mapper.toParseResponse(payload);
        Assert.assertEquals("vlm-http-client", response.getBackend());
        Assert.assertEquals(1, response.getResults().size());

        ParseResult result = response.getResults().get(0);
        Assert.assertEquals("demo", result.getFileName());
        Assert.assertEquals("# title", result.getMarkdown());
        Assert.assertEquals(1, result.getPages().size());
        Assert.assertFalse(result.getBlocks().isEmpty());
        Assert.assertEquals(1, result.getTables().size());
        Assert.assertEquals(2, result.getImages().size());
        Assert.assertNotNull(result.getImages().get(0).getContent());
        Assert.assertNotNull(result.getArtifacts().getMiddleJson());
        Assert.assertNotNull(result.getArtifacts().getContentList());
    }

    @Test
    public void shouldMapStringifiedSyncArtifacts() {
        MineruMapper mapper = new MineruMapper(defaultProperties());
        MineruResultPayload payload = mapper.toResultPayload(syncPayloadWithStringifiedArtifacts());

        ParseResponse response = mapper.toParseResponse(payload);

        Assert.assertEquals(1, response.getResults().size());
        ParseResult result = response.getResults().get(0);
        Assert.assertEquals("# title", result.getMarkdown());
        Assert.assertFalse(result.getBlocks().isEmpty());
        Assert.assertEquals(1, result.getTables().size());
        Assert.assertEquals(2, result.getImages().size());
        Assert.assertNotNull(result.getImages().get(0).getContent());
        Assert.assertTrue(result.getArtifacts().getMiddleJson() instanceof JSONObject);
        Assert.assertTrue(result.getArtifacts().getContentList() instanceof JSONArray);
        Assert.assertEquals("plain-model-output", result.getArtifacts().getModelOutput());
    }

    @Test
    public void shouldMapZipResponse() throws IOException {
        MineruMapper mapper = new MineruMapper(defaultProperties());
        ParseResponse response = mapper.fromZip(buildZip(true));

        Assert.assertEquals(1, response.getResults().size());
        ParseResult result = response.getResults().get(0);
        Assert.assertEquals("demo", result.getFileName());
        Assert.assertEquals("# title", result.getPlainText());
        Assert.assertEquals(1, result.getTables().size());
        Assert.assertEquals(2, result.getImages().size());
        Assert.assertNotNull(result.getImages().get(0).getContent());
        Assert.assertNotNull(result.getArtifacts().getExtraJsonArtifacts().get("contentListV2"));
    }

    @Test
    public void shouldAvoidDuplicatedVisualsWhenFallbackToMiddleJson() {
        MineruMapper mapper = new MineruMapper(defaultProperties());
        MineruResultPayload payload = mapper.toResultPayload(syncPayloadWithMiddleJsonFallback());

        ParseResponse response = mapper.toParseResponse(payload);

        Assert.assertEquals(1, response.getResults().size());
        ParseResult result = response.getResults().get(0);
        Assert.assertEquals(2, result.getBlocks().size());
        Assert.assertEquals(1, result.getTables().size());
        Assert.assertEquals(1, result.getImages().size());
    }

    @Test(expected = DocumentParseException.class)
    public void shouldRejectZipWithoutCriticalArtifacts() throws IOException {
        MineruMapper mapper = new MineruMapper(defaultProperties());
        mapper.fromZip(buildZip(false));
    }

    @Test
    public void shouldBuildAsyncFormWithFullArtifacts() {
        MineruMapper mapper = new MineruMapper(defaultProperties());
        ParseRequest request = new ParseRequest();
        request.setReturnMarkdown(false);
        request.setReturnMiddleJson(false);
        request.setReturnContentList(false);
        request.setReturnModelOutput(false);
        request.setReturnImages(false);

        Map<String, List<String>> fields = mapper.buildAsyncFormFields(request);

        Assert.assertEquals("true", fields.get("return_md").get(0));
        Assert.assertEquals("true", fields.get("return_middle_json").get(0));
        Assert.assertEquals("true", fields.get("return_content_list").get(0));
        Assert.assertEquals("true", fields.get("return_model_output").get(0));
        Assert.assertEquals("true", fields.get("return_images").get(0));
        Assert.assertEquals("true", fields.get("response_format_zip").get(0));
    }

    @Test
    public void shouldBuildRepeatedLangListFields() {
        MineruMapper mapper = new MineruMapper(defaultProperties());
        ParseRequest request = new ParseRequest();
        request.setLanguages(java.util.Arrays.asList("zh", "en"));

        Map<String, List<String>> fields = mapper.buildSyncFormFields(request);

        Assert.assertEquals(java.util.Arrays.asList("zh", "en"), fields.get("lang_list"));
    }

    @Test
    public void shouldFallbackVersionFromMiddleJsonWhenAsyncStatusVersionMissing() throws IOException {
        MineruMapper mapper = new MineruMapper(defaultProperties());
        ParseResponse response = mapper.fromZip(buildZip(true));

        mapper.enrichAsyncResponse(response, null, null);

        Assert.assertEquals("vlm", response.getBackend());
        Assert.assertEquals("3.0.9", response.getVersion());
    }

    @Test
    public void shouldMapNestedZipWhenModelArtifactIsArray() throws IOException {
        MineruMapper mapper = new MineruMapper(defaultProperties());

        ParseResponse response = mapper.fromZip(buildNestedDocxZipWithArrayModel());

        Assert.assertEquals(1, response.getResults().size());
        ParseResult result = response.getResults().get(0);
        Assert.assertEquals("demo", result.getFileName());
        Assert.assertEquals("# nested", result.getMarkdown());
        Assert.assertFalse(result.getBlocks().isEmpty());
        Assert.assertTrue(result.getArtifacts().getModelOutput() instanceof JSONArray);
        Assert.assertTrue(result.getArtifacts().getContentList() instanceof JSONArray);
    }

    private MineruProperties defaultProperties() {
        MineruProperties properties = new MineruProperties();
        properties.setBaseUrl("http://127.0.0.1:8000");
        return properties;
    }

    private JSONObject syncPayload() {
        JSONObject payload = new JSONObject();
        payload.put("backend", "vlm-http-client");
        payload.put("version", "3.0.9");

        JSONObject result = new JSONObject();
        result.put("md_content", "# title");
        result.put("middle_json", middleJson());
        result.put("content_list", contentList());
        result.put("model_output", new JSONObject());

        JSONObject images = new JSONObject();
        images.put("figure.png", "data:image/png;base64,ZmFrZQ==");
        result.put("images", images);

        JSONObject results = new JSONObject();
        results.put("demo", result);
        payload.put("results", results);
        return payload;
    }

    private JSONObject middleJson() {
        JSONObject middleJson = new JSONObject();
        middleJson.put("_backend", "vlm");
        middleJson.put("_version_name", "3.0.9");

        JSONObject page = new JSONObject();
        page.put("page_idx", 0);
        JSONArray pageSize = new JSONArray();
        pageSize.add(1000);
        pageSize.add(2000);
        page.put("page_size", pageSize);
        page.put("para_blocks", new JSONArray());

        JSONArray pdfInfo = new JSONArray();
        pdfInfo.add(page);
        middleJson.put("pdf_info", pdfInfo);
        return middleJson;
    }

    private JSONObject middleJsonForFallback() {
        JSONObject middleJson = new JSONObject();
        middleJson.put("_backend", "vlm");
        middleJson.put("_version_name", "3.0.9");

        JSONObject page = new JSONObject();
        page.put("page_idx", 0);
        page.put("page_size", bboxPageSize());

        JSONArray paraBlocks = new JSONArray();
        paraBlocks.add(middleBlock("table", "images/table.png"));
        paraBlocks.add(middleBlock("image", "images/figure.png"));
        page.put("para_blocks", paraBlocks);

        JSONArray tables = new JSONArray();
        tables.add(middleTable("images/table.png"));
        page.put("tables", tables);

        JSONArray images = new JSONArray();
        images.add(middleImage("images/figure.png"));
        page.put("images", images);

        JSONArray pdfInfo = new JSONArray();
        pdfInfo.add(page);
        middleJson.put("pdf_info", pdfInfo);
        return middleJson;
    }

    private JSONArray contentList() {
        JSONArray contentList = new JSONArray();

        JSONObject title = new JSONObject();
        title.put("type", "text");
        title.put("text", "title");
        title.put("text_level", 1);
        title.put("page_idx", 0);
        title.put("bbox", bbox());
        contentList.add(title);

        JSONObject image = new JSONObject();
        image.put("type", "image");
        image.put("img_path", "images/figure.png");
        image.put("image_caption", new JSONArray());
        image.put("image_footnote", new JSONArray());
        image.put("page_idx", 0);
        image.put("bbox", bbox());
        contentList.add(image);

        JSONObject table = new JSONObject();
        table.put("type", "table");
        table.put("img_path", "images/table.png");
        table.put("table_body", "<table></table>");
        table.put("table_caption", new JSONArray());
        table.put("table_footnote", new JSONArray());
        table.put("page_idx", 0);
        table.put("bbox", bbox());
        contentList.add(table);

        return contentList;
    }

    private JSONArray contentListV2() {
        JSONArray contentList = new JSONArray();
        JSONObject page = new JSONObject();
        page.put("page_idx", 0);
        contentList.add(page);
        return contentList;
    }

    private JSONArray bbox() {
        JSONArray bbox = new JSONArray();
        bbox.add(1);
        bbox.add(2);
        bbox.add(3);
        bbox.add(4);
        return bbox;
    }

    private JSONArray bboxPageSize() {
        JSONArray bbox = new JSONArray();
        bbox.add(1000);
        bbox.add(2000);
        return bbox;
    }

    private JSONObject syncPayloadWithMiddleJsonFallback() {
        JSONObject payload = new JSONObject();
        payload.put("backend", "vlm-http-client");
        payload.put("version", "3.0.9");

        JSONObject result = new JSONObject();
        result.put("md_content", "# fallback");
        result.put("middle_json", middleJsonForFallback());

        JSONObject images = new JSONObject();
        images.put("figure.png", "data:image/png;base64,ZmFrZQ==");
        images.put("table.png", "data:image/png;base64,ZmFrZQ==");
        result.put("images", images);

        JSONObject results = new JSONObject();
        results.put("demo", result);
        payload.put("results", results);
        return payload;
    }

    private JSONObject syncPayloadWithStringifiedArtifacts() {
        JSONObject payload = new JSONObject();
        payload.put("backend", "vlm-http-client");
        payload.put("version", "3.0.9");

        JSONObject result = new JSONObject();
        result.put("md_content", "# title");
        result.put("middle_json", middleJson().toJSONString());
        result.put("content_list", contentList().toJSONString());
        result.put("model_output", "plain-model-output");

        JSONObject images = new JSONObject();
        images.put("figure.png", "data:image/png;base64,ZmFrZQ==");
        images.put("table.png", "data:image/png;base64,ZmFrZQ==");
        result.put("images", images.toJSONString());

        JSONObject results = new JSONObject();
        results.put("demo", result.toJSONString());
        payload.put("results", results);
        return payload;
    }

    private JSONObject middleBlock(String type, String imagePath) {
        JSONObject block = new JSONObject();
        block.put("type", type);
        block.put("bbox", bbox());
        JSONArray blocks = new JSONArray();
        JSONObject childBlock = new JSONObject();
        JSONArray lines = new JSONArray();
        JSONObject line = new JSONObject();
        JSONArray spans = new JSONArray();

        JSONObject textSpan = new JSONObject();
        textSpan.put("content", type + "-text");
        spans.add(textSpan);

        JSONObject imageSpan = new JSONObject();
        imageSpan.put("img_path", imagePath);
        spans.add(imageSpan);

        line.put("spans", spans);
        lines.add(line);
        childBlock.put("lines", lines);
        blocks.add(childBlock);
        block.put("blocks", blocks);
        return block;
    }

    private JSONObject middleTable(String imagePath) {
        JSONObject table = new JSONObject();
        table.put("bbox", bbox());
        JSONArray blocks = new JSONArray();
        blocks.add(visualBlock("table_caption", null, "table-caption"));
        blocks.add(visualBlock("table_body", imagePath, null));
        table.put("blocks", blocks);
        return table;
    }

    private JSONObject middleImage(String imagePath) {
        JSONObject image = new JSONObject();
        image.put("bbox", bbox());
        JSONArray blocks = new JSONArray();
        blocks.add(visualBlock("image_caption", null, "image-caption"));
        blocks.add(visualBlock("image_body", imagePath, null));
        image.put("blocks", blocks);
        return image;
    }

    private JSONObject visualBlock(String type, String imagePath, String text) {
        JSONObject block = new JSONObject();
        block.put("type", type);
        JSONArray lines = new JSONArray();
        JSONObject line = new JSONObject();
        JSONArray spans = new JSONArray();
        JSONObject span = new JSONObject();
        if (imagePath != null) {
            span.put("img_path", imagePath);
        }
        if (text != null) {
            span.put("content", text);
        }
        spans.add(span);
        line.put("spans", spans);
        lines.add(line);
        block.put("lines", lines);
        return block;
    }

    private byte[] buildZip(boolean withArtifacts) throws IOException {
        ByteArrayOutputStream outputStream = new ByteArrayOutputStream();
        ZipOutputStream zipOutputStream = new ZipOutputStream(outputStream);
        if (withArtifacts) {
            addEntry(zipOutputStream, "demo/vlm/demo.md", "# title");
            addEntry(zipOutputStream, "demo/vlm/demo_middle.json", middleJson().toJSONString());
            addEntry(zipOutputStream, "demo/vlm/demo_content_list.json", contentList().toJSONString());
            addEntry(zipOutputStream, "demo/vlm/demo_content_list_v2.json", contentListV2().toJSONString());
            addEntry(zipOutputStream, "demo/vlm/demo_model.json", "{}");
        }
        addBinaryEntry(zipOutputStream, "demo/vlm/images/figure.png", "image".getBytes(StandardCharsets.UTF_8));
        addBinaryEntry(zipOutputStream, "demo/vlm/images/table.png", "image".getBytes(StandardCharsets.UTF_8));
        zipOutputStream.close();
        return outputStream.toByteArray();
    }

    private byte[] buildNestedDocxZipWithArrayModel() throws IOException {
        ByteArrayOutputStream outputStream = new ByteArrayOutputStream();
        ZipOutputStream zipOutputStream = new ZipOutputStream(outputStream);
        addEntry(zipOutputStream, "demo/vlm/markdown/demo.md", "# nested");
        addEntry(zipOutputStream, "demo/vlm/layout/demo_middle.json", middleJson().toJSONString());
        addEntry(zipOutputStream, "demo/vlm/model/demo_model.json", nestedDocxContentList().toJSONString());
        addBinaryEntry(zipOutputStream, "demo/vlm/images/figure.png", "image".getBytes(StandardCharsets.UTF_8));
        zipOutputStream.close();
        return outputStream.toByteArray();
    }

    private JSONArray nestedDocxContentList() {
        JSONArray contentList = new JSONArray();

        JSONObject title = new JSONObject();
        title.put("type", "title");
        title.put("content", "二、技术要求");
        title.put("page_idx", 0);
        title.put("bbox", bbox());
        contentList.add(title);

        JSONObject text = new JSONObject();
        text.put("type", "text");
        text.put("content", "响应方式");
        text.put("page_idx", 0);
        text.put("bbox", bbox());
        contentList.add(text);

        JSONObject table = new JSONObject();
        table.put("type", "table");
        table.put("content", "<table></table>");
        table.put("table_body", "<table></table>");
        table.put("page_idx", 0);
        table.put("bbox", bbox());
        contentList.add(table);

        return contentList;
    }

    private void addEntry(ZipOutputStream zipOutputStream, String name, String content) throws IOException {
        addBinaryEntry(zipOutputStream, name, content.getBytes(StandardCharsets.UTF_8));
    }

    private void addBinaryEntry(ZipOutputStream zipOutputStream, String name, byte[] content) throws IOException {
        zipOutputStream.putNextEntry(new ZipEntry(name));
        zipOutputStream.write(content);
        zipOutputStream.closeEntry();
    }
}
