package tech.easyflow.ai.documentimport.task;

import org.junit.Assert;
import org.junit.Test;
import tech.easyflow.ai.documentimport.DocumentImportDtos;
import tech.easyflow.ai.entity.DocumentChunk;

import java.lang.reflect.Field;
import java.math.BigInteger;
import java.util.ArrayList;
import java.util.List;

/**
 * {@link DocumentImportChunkSnapshotService} 持久化恢复测试。
 */
public class DocumentImportChunkSnapshotServiceTest {

    /**
     * 验证分块快照可跨缓存写入并从文件存储完整恢复。
     *
     * @throws Exception 反射注入或文件读取异常
     */
    @Test
    public void shouldPersistAndRestoreChunkSnapshot() throws Exception {
        InMemoryFileStorageService storageService =
            new InMemoryFileStorageService();

        DocumentImportChunkSnapshotService service = new DocumentImportChunkSnapshotService();
        Field storageField = DocumentImportChunkSnapshotService.class
            .getDeclaredField("storageService");
        storageField.setAccessible(true);
        storageField.set(service, storageService);

        DocumentChunk chunk = new DocumentChunk();
        chunk.setId(BigInteger.valueOf(11));
        chunk.setDocumentId(BigInteger.valueOf(9));
        chunk.setDocumentCollectionId(BigInteger.valueOf(7));
        chunk.setContent("稳定分块");
        DocumentImportDtos.PreviewSession session = new DocumentImportDtos.PreviewSession();
        session.setKnowledgeId(BigInteger.valueOf(7));
        session.setDocumentId(BigInteger.valueOf(9));
        session.setDocumentChunks(List.of(chunk));

        String path = service.save(session);
        DocumentImportDtos.PreviewSession restored = service.load(path);

        Assert.assertTrue(path.endsWith("-manifest.json"));
        Assert.assertEquals(session.getKnowledgeId(), restored.getKnowledgeId());
        Assert.assertEquals(session.getDocumentId(), restored.getDocumentId());
        Assert.assertEquals(1, restored.getDocumentChunks().size());
        Assert.assertEquals("稳定分块", restored.getDocumentChunks().get(0).getContent());
    }

    /**
     * 验证 V2 快照可分页并按小批顺序消费，删除时同时清理清单和分片。
     *
     * @throws Exception 反射注入异常
     */
    @Test
    public void shouldReadV2SnapshotByPageAndBatch() throws Exception {
        InMemoryFileStorageService storageService =
            new InMemoryFileStorageService();
        DocumentImportChunkSnapshotService service =
            new DocumentImportChunkSnapshotService();
        Field storageField = DocumentImportChunkSnapshotService.class
            .getDeclaredField("storageService");
        storageField.setAccessible(true);
        storageField.set(service, storageService);

        DocumentImportDtos.PreviewSession session =
            new DocumentImportDtos.PreviewSession();
        session.setKnowledgeId(BigInteger.valueOf(7));
        session.setDocumentId(BigInteger.valueOf(9));
        List<DocumentChunk> chunks = new ArrayList<DocumentChunk>();
        for (int index = 0; index < 5; index++) {
            DocumentChunk chunk = new DocumentChunk();
            chunk.setId(BigInteger.valueOf(100 + index));
            chunk.setDocumentId(BigInteger.valueOf(9));
            chunk.setDocumentCollectionId(BigInteger.valueOf(7));
            chunk.setContent("chunk-" + index);
            chunks.add(chunk);
        }
        session.setDocumentChunks(chunks);

        String path = service.save(session);
        List<DocumentChunk> page = service.loadPage(path, 2, 2);
        List<String> consumed = new ArrayList<String>();
        service.forEachBatch(path, 2, batch -> {
            Assert.assertTrue(batch.size() <= 2);
            for (DocumentChunk chunk : batch) {
                consumed.add(chunk.getContent());
            }
        });

        Assert.assertEquals(List.of("chunk-2", "chunk-3"),
            page.stream().map(DocumentChunk::getContent).toList());
        Assert.assertEquals(
            List.of("chunk-0", "chunk-1", "chunk-2", "chunk-3", "chunk-4"),
            consumed);
        Assert.assertTrue(storageService.size() >= 2);
        service.delete(path);
        Assert.assertEquals(0, storageService.size());
    }

    /**
     * 验证分片删除失败时保留清单，使后续重试仍能获取精确对象集合。
     *
     * @throws Exception 反射注入异常
     */
    @Test
    public void shouldKeepManifestWhenPartDeletionFails() throws Exception {
        InMemoryFileStorageService storageService =
            new InMemoryFileStorageService();
        DocumentImportChunkSnapshotService service =
            new DocumentImportChunkSnapshotService();
        Field storageField = DocumentImportChunkSnapshotService.class
            .getDeclaredField("storageService");
        storageField.setAccessible(true);
        storageField.set(service, storageService);

        DocumentChunk chunk = new DocumentChunk();
        chunk.setId(BigInteger.valueOf(21));
        chunk.setDocumentId(BigInteger.valueOf(9));
        chunk.setDocumentCollectionId(BigInteger.valueOf(7));
        chunk.setContent("等待可靠清理");
        DocumentImportDtos.PreviewSession session =
            new DocumentImportDtos.PreviewSession();
        session.setKnowledgeId(BigInteger.valueOf(7));
        session.setDocumentId(BigInteger.valueOf(9));
        session.setDocumentChunks(List.of(chunk));

        String manifestPath = service.save(session);
        String partPath = storageService.paths().stream()
            .filter(path -> !path.equals(manifestPath))
            .findFirst()
            .orElseThrow();
        storageService.failDelete(partPath, 1);

        try {
            service.delete(manifestPath);
            Assert.fail("分片删除失败时应抛出异常");
        } catch (IllegalStateException expected) {
            Assert.assertTrue(expected.getMessage().contains("模拟删除失败"));
        }
        Assert.assertTrue(storageService.contains(manifestPath));
        Assert.assertTrue(storageService.contains(partPath));

        service.delete(manifestPath);
        Assert.assertEquals(0, storageService.size());
    }
}
