package com.easyagents.document.pdf.mineru;

import com.alibaba.fastjson2.JSONObject;
import com.easyagents.document.core.mineru.MineruClient;
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.mineru.MineruTaskStatus;
import com.easyagents.document.core.entity.ParseFile;
import com.easyagents.document.core.entity.ParseRequest;
import com.easyagents.document.core.entity.ParseResponse;
import com.easyagents.document.core.entity.ParseTaskInfo;
import com.easyagents.document.core.entity.ParseTaskStatus;
import okhttp3.Request;
import okio.Buffer;
import org.junit.Assert;
import org.junit.Test;

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

/**
 * MinerU PDF 服务测试。
 *
 * @author Codex
 * @since 2026-04-14
 */
public class MineruPdfDocumentParseServiceTest {

    @Test
    public void shouldForceAsyncResultArtifacts() {
        RecordingClient client = new RecordingClient(defaultProperties());
        MineruMapper mapper = new MineruMapper(defaultProperties());
        MineruPdfDocumentParseService service = new MineruPdfDocumentParseService(defaultProperties(), client, mapper);

        ParseRequest request = buildRequest();
        request.setReturnMarkdown(false);
        request.setReturnMiddleJson(false);
        request.setReturnContentList(false);
        request.setReturnModelOutput(false);
        request.setReturnImages(false);

        ParseTaskStatus status = service.submit(request);

        Assert.assertEquals("task-1", status.getTaskId());
        Assert.assertTrue(client.lastSubmitRequest.getReturnMarkdown());
        Assert.assertTrue(client.lastSubmitRequest.getReturnMiddleJson());
        Assert.assertTrue(client.lastSubmitRequest.getReturnContentList());
        Assert.assertTrue(client.lastSubmitRequest.getReturnModelOutput());
        Assert.assertTrue(client.lastSubmitRequest.getReturnImages());
    }

    @Test
    public void shouldUseSyncResultFlagsDuringParse() {
        RecordingClient client = new RecordingClient(defaultProperties());
        MineruMapper mapper = new MineruMapper(defaultProperties());
        MineruPdfDocumentParseService service = new MineruPdfDocumentParseService(defaultProperties(), client, mapper);

        ParseRequest request = buildRequest();
        request.setReturnMarkdown(true);
        request.setReturnMiddleJson(false);
        request.setReturnContentList(true);
        request.setReturnModelOutput(false);
        request.setReturnImages(false);

        ParseResponse response = service.parse(request);

        Assert.assertEquals(1, response.getResults().size());
        Assert.assertFalse(client.lastParseRequest.getReturnMiddleJson());
        Assert.assertFalse(client.lastParseRequest.getReturnImages());
    }

    @Test
    public void shouldUseTaskMetadataWhenQueryingAsyncZipResult() {
        RecordingClient client = new RecordingClient(defaultProperties());
        MineruMapper mapper = new MineruMapper(defaultProperties());
        MineruPdfDocumentParseService service = new MineruPdfDocumentParseService(defaultProperties(), client, mapper);

        ParseResponse response = service.queryResult("task-1");

        Assert.assertEquals("vlm-http-client", response.getBackend());
        Assert.assertEquals("3.0.9", response.getVersion());
        Assert.assertEquals(1, response.getResults().size());
        Assert.assertEquals("demo", response.getResults().get(0).getFileName());
    }

    @Test
    public void shouldReturnPendingStatusWithoutFetchingResultInTaskInfo() {
        RecordingClient client = new RecordingClient(defaultProperties());
        client.taskStatusValue = "running";
        MineruMapper mapper = new MineruMapper(defaultProperties());
        MineruPdfDocumentParseService service = new MineruPdfDocumentParseService(defaultProperties(), client, mapper);

        ParseTaskInfo taskInfo = service.queryTaskInfo("task-1");

        Assert.assertEquals("running", taskInfo.getStatus());
        Assert.assertNull(taskInfo.getResult());
        Assert.assertEquals(0, client.queryResultZipCount);
    }

    @Test
    public void shouldReturnCompletedResultInTaskInfo() {
        RecordingClient client = new RecordingClient(defaultProperties());
        MineruMapper mapper = new MineruMapper(defaultProperties());
        MineruPdfDocumentParseService service = new MineruPdfDocumentParseService(defaultProperties(), client, mapper);

        ParseTaskInfo taskInfo = service.queryTaskInfo("task-1");

        Assert.assertEquals("completed", taskInfo.getStatus());
        Assert.assertNotNull(taskInfo.getResult());
        Assert.assertEquals(1, taskInfo.getResult().getResults().size());
        Assert.assertEquals(1, client.queryResultZipCount);
    }

    @Test
    public void shouldSendRepeatedLangListFields() {
        InspectingMultipartClient client = new InspectingMultipartClient(defaultProperties());
        ParseRequest request = buildRequest();
        request.setLanguages(java.util.Arrays.asList("zh", "en"));

        client.parse(request);

        Assert.assertEquals(2, countOccurrences(client.lastMultipartBody, "name=\"lang_list\""));
        Assert.assertTrue(client.lastMultipartBody.contains("\r\nzh\r\n"));
        Assert.assertTrue(client.lastMultipartBody.contains("\r\nen\r\n"));
    }

    private ParseRequest buildRequest() {
        ParseRequest request = new ParseRequest();
        request.addFile(ParseFile.of("demo.pdf", "pdf".getBytes(StandardCharsets.UTF_8)));
        return request;
    }

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

    private int countOccurrences(String source, String token) {
        int count = 0;
        int index = 0;
        while (source != null && token != null && !token.isEmpty() && (index = source.indexOf(token, index)) >= 0) {
            count++;
            index += token.length();
        }
        return count;
    }

    private static class RecordingClient extends MineruClient {

        private ParseRequest lastParseRequest;
        private ParseRequest lastSubmitRequest;
        private String taskStatusValue = "completed";
        private int queryResultZipCount;

        private RecordingClient(MineruProperties properties) {
            super(properties, new MineruMapper(properties));
        }

        @Override
        public MineruResultPayload parse(ParseRequest request) {
            this.lastParseRequest = request;
            return new MineruMapper(testProperties()).toResultPayload(syncPayload());
        }

        @Override
        public MineruTaskStatus submit(ParseRequest request) {
            this.lastSubmitRequest = request;
            MineruTaskStatus taskStatus = new MineruTaskStatus();
            taskStatus.setTaskId("task-1");
            taskStatus.setStatus("pending");
            return taskStatus;
        }

        @Override
        public MineruTaskStatus queryTask(String taskId) {
            MineruTaskStatus taskStatus = new MineruTaskStatus();
            taskStatus.setTaskId(taskId);
            taskStatus.setStatus(taskStatusValue);
            taskStatus.setBackend("vlm-http-client");
            taskStatus.setVersion("3.0.9");
            return taskStatus;
        }

        @Override
        public byte[] queryResultZip(String taskId) {
            queryResultZipCount++;
            try {
                return buildZipResult();
            } catch (IOException exception) {
                throw new IllegalStateException("Failed to build test ZIP", exception);
            }
        }

        private static byte[] buildZipResult() throws IOException {
            ByteArrayOutputStream outputStream = new ByteArrayOutputStream();
            try (ZipOutputStream zipOutputStream = new ZipOutputStream(outputStream)) {
                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());
            }
            return outputStream.toByteArray();
        }

        private static void addEntry(ZipOutputStream zipOutputStream, String name, String content) throws IOException {
            zipOutputStream.putNextEntry(new ZipEntry(name));
            zipOutputStream.write(content.getBytes(StandardCharsets.UTF_8));
            zipOutputStream.closeEntry();
        }

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

            JSONObject file = new JSONObject();
            file.put("md_content", "# title");
            JSONObject results = new JSONObject();
            results.put("demo", file);
            payload.put("results", results);
            return payload;
        }

        private static JSONObject middleJson() {
            JSONObject middleJson = new JSONObject();
            middleJson.put("_backend", "vlm");
            middleJson.put("_version_name", "3.0.9");
            middleJson.put("pdf_info", new com.alibaba.fastjson2.JSONArray());
            return middleJson;
        }

        private static com.alibaba.fastjson2.JSONArray contentList() {
            com.alibaba.fastjson2.JSONArray contentList = new com.alibaba.fastjson2.JSONArray();
            JSONObject text = new JSONObject();
            text.put("type", "text");
            text.put("text", "title");
            text.put("page_idx", 0);
            text.put("bbox", new com.alibaba.fastjson2.JSONArray());
            contentList.add(text);
            return contentList;
        }

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

    private static class InspectingMultipartClient extends MineruClient {

        private String lastMultipartBody;

        private InspectingMultipartClient(MineruProperties properties) {
            super(properties, new MineruMapper(properties));
        }

        @Override
        protected JSONObject executeJsonRequest(String path, Request request) {
            try {
                Buffer buffer = new Buffer();
                request.body().writeTo(buffer);
                this.lastMultipartBody = buffer.readUtf8();
            } catch (IOException exception) {
                throw new IllegalStateException("Failed to inspect multipart body", exception);
            }
            return new JSONObject();
        }
    }
}
