package tech.easyflow.ai.documentimport.task;

import com.alibaba.fastjson2.JSON;
import com.easyagents.rag.ingestion.model.AnalysisResult;
import org.springframework.stereotype.Service;
import tech.easyflow.ai.document.support.DocumentInputStreamSupport;
import tech.easyflow.ai.documentimport.DocumentImportDtos;
import tech.easyflow.ai.easyagents.CustomMultipartFile;
import tech.easyflow.ai.entity.DocumentChunk;
import tech.easyflow.common.filestorage.FileStorageService;
import tech.easyflow.common.util.StringUtil;
import tech.easyflow.common.web.exceptions.BusinessException;

import javax.annotation.Resource;
import java.io.IOException;
import java.io.InputStream;
import java.nio.charset.StandardCharsets;
import java.security.MessageDigest;
import java.security.NoSuchAlgorithmException;
import java.util.ArrayList;
import java.util.Date;
import java.util.HexFormat;
import java.util.List;
import java.util.UUID;
import java.util.function.Consumer;

/**
 * 自动导入分块快照持久化服务。
 *
 * <p>V2 快照使用“有界分片 + 最后写入清单”的提交协议，索引任务可以按批读取，
 * 同时保留对历史 V1 单 JSON 快照的读取兼容。</p>
 *
 * @author Codex
 * @since 2026-07-31
 */
@Service
public class DocumentImportChunkSnapshotService {

    private static final String SNAPSHOT_VERSION_V2 = "v2";
    private static final int DEFAULT_PART_CHUNK_COUNT = 128;
    private static final long MAX_MANIFEST_BYTES = 4L * 1024L * 1024L;
    private static final long MAX_PART_BYTES = 4L * 1024L * 1024L;
    private static final long MAX_LEGACY_SNAPSHOT_BYTES = 256L * 1024L * 1024L;

    @Resource(name = "default")
    private FileStorageService storageService;

    /**
     * 持久化预览会话及其最终分块。
     *
     * @param session 预览会话
     * @return V2 清单存储路径
     */
    public String save(DocumentImportDtos.PreviewSession session) {
        validateSession(session, true);
        try (SnapshotWriter writer = createWriter(session)) {
            List<DocumentChunk> chunks = session.getDocumentChunks();
            for (int start = 0; start < chunks.size(); start += DEFAULT_PART_CHUNK_COUNT) {
                int end = Math.min(start + DEFAULT_PART_CHUNK_COUNT, chunks.size());
                writer.append(new ArrayList<DocumentChunk>(chunks.subList(start, end)));
            }
            return writer.finish();
        }
    }

    /**
     * 创建流式快照写入器。
     *
     * @param session 仅需包含归属、策略和预览元信息的会话
     * @return 快照写入器
     */
    public SnapshotWriter createWriter(DocumentImportDtos.PreviewSession session) {
        validateSession(session, false);
        return new SnapshotWriter(copyHeader(session));
    }

    /**
     * 从稳定存储完整恢复预览会话。
     *
     * <p>该方法用于兼容现有调用和小型预览。大文件索引应使用
     * {@link #forEachBatch(String, int, Consumer)}。</p>
     *
     * @param path 快照路径
     * @return 预览会话
     */
    public DocumentImportDtos.PreviewSession load(String path) {
        SnapshotReadResult readResult = readManifestOrLegacy(path);
        if (readResult.legacySession != null) {
            return readResult.legacySession;
        }
        DocumentImportDtos.PreviewSession session = copyHeader(readResult.manifest.getSession());
        List<DocumentChunk> chunks = new ArrayList<DocumentChunk>();
        long loadedBytes = 0L;
        for (SnapshotPart part : readResult.manifest.getParts()) {
            PartPayload payload = readPart(part);
            loadedBytes += payload.byteLength;
            if (loadedBytes > MAX_LEGACY_SNAPSHOT_BYTES) {
                throw new BusinessException("分块快照过大，请使用分页或批量读取");
            }
            chunks.addAll(payload.chunks);
        }
        session.setDocumentChunks(chunks);
        session.setTotalChunks(chunks.size());
        session.setChunkSnapshotPath(path);
        return session;
    }

    /**
     * 仅读取快照头部元信息。
     *
     * @param path 快照路径
     * @return 不包含完整分块的预览会话
     */
    public DocumentImportDtos.PreviewSession loadHeader(String path) {
        SnapshotReadResult readResult = readManifestOrLegacy(path);
        if (readResult.legacySession != null) {
            DocumentImportDtos.PreviewSession header = copyHeader(readResult.legacySession);
            int total = readResult.legacySession.getDocumentChunks() == null
                ? 0
                : readResult.legacySession.getDocumentChunks().size();
            header.setTotalChunks(total);
            header.setChunkSnapshotPath(path);
            return header;
        }
        DocumentImportDtos.PreviewSession header = copyHeader(readResult.manifest.getSession());
        header.setTotalChunks(readResult.manifest.getTotalChunks());
        header.setChunkSnapshotPath(path);
        return header;
    }

    /**
     * 分页读取快照分块。
     *
     * @param path 快照路径
     * @param offset 零基偏移
     * @param limit 最大返回数
     * @return 指定页分块
     */
    public List<DocumentChunk> loadPage(String path, int offset, int limit) {
        if (offset < 0 || limit <= 0) {
            throw new BusinessException("分块快照分页参数不合法");
        }
        SnapshotReadResult readResult = readManifestOrLegacy(path);
        if (readResult.legacySession != null) {
            List<DocumentChunk> chunks = readResult.legacySession.getDocumentChunks();
            if (chunks == null || offset >= chunks.size()) {
                return new ArrayList<DocumentChunk>();
            }
            int end = Math.min(offset + limit, chunks.size());
            return new ArrayList<DocumentChunk>(chunks.subList(offset, end));
        }
        List<DocumentChunk> page = new ArrayList<DocumentChunk>(limit);
        int skipped = 0;
        for (SnapshotPart part : readResult.manifest.getParts()) {
            if (page.size() >= limit) {
                break;
            }
            if (skipped + part.getChunkCount() <= offset) {
                skipped += part.getChunkCount();
                continue;
            }
            List<DocumentChunk> chunks = readPart(part).chunks;
            int localStart = Math.max(0, offset - skipped);
            int localEnd = Math.min(chunks.size(), localStart + limit - page.size());
            page.addAll(chunks.subList(localStart, localEnd));
            skipped += chunks.size();
        }
        return page;
    }

    /**
     * 按固定上限顺序消费快照分块。
     *
     * @param path 快照路径
     * @param batchSize 单批最大分块数
     * @param consumer 批处理函数
     */
    public void forEachBatch(
        String path,
        int batchSize,
        Consumer<List<DocumentChunk>> consumer) {
        if (batchSize <= 0 || consumer == null) {
            throw new IllegalArgumentException("快照批量消费参数不合法");
        }
        SnapshotReadResult readResult = readManifestOrLegacy(path);
        if (readResult.legacySession != null) {
            consumeInBatches(readResult.legacySession.getDocumentChunks(), batchSize, consumer);
            return;
        }
        for (SnapshotPart part : readResult.manifest.getParts()) {
            consumeInBatches(readPart(part).chunks, batchSize, consumer);
        }
    }

    /**
     * 删除快照清单及其全部分片。
     *
     * @param path 快照清单或历史 V1 快照路径
     */
    public void delete(String path) {
        if (!StringUtil.hasText(path)) {
            return;
        }
        deleteParts(path);
        deleteManifest(path);
    }

    /**
     * 删除快照分片并保留清单。
     *
     * <p>任一分片删除失败时立即抛出异常，清单继续作为精确重试依据。</p>
     *
     * @param path 快照清单或历史 V1 快照路径
     */
    public void deleteParts(String path) {
        if (!StringUtil.hasText(path)) {
            return;
        }
        SnapshotReadResult readResult = readManifestOrLegacy(path);
        if (readResult.manifest == null) {
            return;
        }
        for (SnapshotPart part : readResult.manifest.getParts()) {
            if (StringUtil.hasText(part.getPath())) {
                storageService.delete(part.getPath());
            }
        }
    }

    /**
     * 删除快照清单或历史 V1 单对象。
     *
     * @param path 快照清单或历史 V1 快照路径
     */
    public void deleteManifest(String path) {
        if (StringUtil.hasText(path)) {
            storageService.delete(path);
        }
    }

    /**
     * 校验预览会话的最小完整性。
     *
     * @param session 预览会话
     * @param requireChunks 是否要求已包含分块
     */
    private void validateSession(
        DocumentImportDtos.PreviewSession session,
        boolean requireChunks) {
        if (session == null || session.getKnowledgeId() == null || session.getDocumentId() == null) {
            throw new BusinessException("分块快照归属信息不完整");
        }
        if (requireChunks
            && (session.getDocumentChunks() == null || session.getDocumentChunks().isEmpty())) {
            throw new BusinessException("分块快照内容不完整");
        }
    }

    /**
     * 复制有界会话头，移除完整正文和分块数组。
     *
     * @param source 原始会话
     * @return 会话头
     */
    private DocumentImportDtos.PreviewSession copyHeader(
        DocumentImportDtos.PreviewSession source) {
        DocumentImportDtos.PreviewSession header = new DocumentImportDtos.PreviewSession();
        header.setSessionId(source.getSessionId());
        header.setKnowledgeId(source.getKnowledgeId());
        header.setDocumentId(source.getDocumentId());
        header.setFilePath(source.getFilePath());
        header.setFileName(source.getFileName());
        header.setSourceFormat(source.getSourceFormat());
        header.setStrategyConfig(source.getStrategyConfig());
        if (source.getAnalysis() != null) {
            AnalysisResult analysis = JSON.parseObject(
                JSON.toJSONBytes(source.getAnalysis()), AnalysisResult.class);
            analysis.setNormalizedContent(null);
            header.setAnalysis(analysis);
        }
        header.setDocument(null);
        header.setDocumentChunks(new ArrayList<DocumentChunk>());
        header.setPreviewChunks(new ArrayList<>());
        header.setChunkSnapshotPath(source.getChunkSnapshotPath());
        header.setTotalChunks(source.getTotalChunks());
        header.setCreatedAt(source.getCreatedAt());
        return header;
    }

    /**
     * 读取 V2 清单或历史 V1 会话。
     *
     * @param path 快照路径
     * @return 快照读取结果
     */
    private SnapshotReadResult readManifestOrLegacy(String path) {
        if (!StringUtil.hasText(path)) {
            throw new BusinessException("分块快照不存在，请重试");
        }
        byte[] payload;
        try (InputStream inputStream = storageService.readStream(path)) {
            payload = DocumentInputStreamSupport.readBytes(
                inputStream, MAX_LEGACY_SNAPSHOT_BYTES);
        } catch (IOException error) {
            throw new BusinessException("分块快照读取失败，请重试");
        }
        SnapshotManifest manifest = JSON.parseObject(payload, SnapshotManifest.class);
        if (manifest != null && SNAPSHOT_VERSION_V2.equals(manifest.getVersion())) {
            validateManifest(manifest);
            return SnapshotReadResult.forManifest(manifest);
        }
        DocumentImportDtos.PreviewSession session = JSON.parseObject(
            new String(payload, StandardCharsets.UTF_8),
            DocumentImportDtos.PreviewSession.class);
        if (session == null || session.getDocumentChunks() == null
            || session.getDocumentChunks().isEmpty()) {
            throw new BusinessException("分块快照无有效内容，请重试");
        }
        return SnapshotReadResult.forLegacy(session);
    }

    /**
     * 校验 V2 清单完整性。
     *
     * @param manifest 快照清单
     */
    private void validateManifest(SnapshotManifest manifest) {
        if (manifest.getSession() == null
            || manifest.getSession().getKnowledgeId() == null
            || manifest.getSession().getDocumentId() == null
            || manifest.getTotalChunks() <= 0
            || manifest.getParts() == null
            || manifest.getParts().isEmpty()) {
            throw new BusinessException("分块快照清单损坏，请重试");
        }
        int total = 0;
        for (SnapshotPart part : manifest.getParts()) {
            if (part == null || !StringUtil.hasText(part.getPath())
                || part.getChunkCount() <= 0
                || part.getByteLength() <= 0
                || !StringUtil.hasText(part.getSha256())) {
                throw new BusinessException("分块快照清单损坏，请重试");
            }
            total += part.getChunkCount();
        }
        if (total != manifest.getTotalChunks()) {
            throw new BusinessException("分块快照清单计数不一致，请重试");
        }
    }

    /**
     * 读取并校验一个分块分片。
     *
     * @param part 分片清单
     * @return 分片内容
     */
    private PartPayload readPart(SnapshotPart part) {
        try (InputStream inputStream = storageService.readStream(part.getPath())) {
            byte[] payload = DocumentInputStreamSupport.readBytes(inputStream, MAX_PART_BYTES);
            if (payload.length != part.getByteLength()
                || !sha256(payload).equals(part.getSha256())) {
                throw new BusinessException("分块快照校验失败，请重试");
            }
            List<DocumentChunk> chunks = JSON.parseArray(
                new String(payload, StandardCharsets.UTF_8), DocumentChunk.class);
            if (chunks == null || chunks.size() != part.getChunkCount()) {
                throw new BusinessException("分块快照分片损坏，请重试");
            }
            return new PartPayload(chunks, payload.length);
        } catch (IOException error) {
            throw new BusinessException("分块快照分片读取失败，请重试");
        }
    }

    /**
     * 将列表按固定大小交给消费者。
     *
     * @param chunks 分块
     * @param batchSize 批大小
     * @param consumer 消费者
     */
    private void consumeInBatches(
        List<DocumentChunk> chunks,
        int batchSize,
        Consumer<List<DocumentChunk>> consumer) {
        if (chunks == null || chunks.isEmpty()) {
            return;
        }
        for (int start = 0; start < chunks.size(); start += batchSize) {
            int end = Math.min(start + batchSize, chunks.size());
            consumer.accept(new ArrayList<DocumentChunk>(chunks.subList(start, end)));
        }
    }

    /**
     * 计算 SHA-256。
     *
     * @param payload 字节内容
     * @return 十六进制摘要
     */
    private String sha256(byte[] payload) {
        try {
            return HexFormat.of().formatHex(
                MessageDigest.getInstance("SHA-256").digest(payload));
        } catch (NoSuchAlgorithmException error) {
            throw new IllegalStateException("当前运行环境不支持 SHA-256", error);
        }
    }

    /**
     * V2 快照流式写入器。
     */
    public final class SnapshotWriter implements AutoCloseable {

        private final SnapshotManifest manifest;
        private final String storagePrefix;
        private boolean committed;

        /**
         * 创建写入器。
         *
         * @param session 会话头
         */
        private SnapshotWriter(DocumentImportDtos.PreviewSession session) {
            this.manifest = new SnapshotManifest();
            this.manifest.setVersion(SNAPSHOT_VERSION_V2);
            this.manifest.setSession(session);
            this.manifest.setCreatedAt(new Date());
            this.storagePrefix = "knowledge-import-snapshots/"
                + session.getKnowledgeId() + "/"
                + session.getDocumentId() + "/"
                + UUID.randomUUID();
        }

        /**
         * 追加一批分块，超出单分片上限时自动二分。
         *
         * @param chunks 分块
         */
        public void append(List<DocumentChunk> chunks) {
            if (committed) {
                throw new IllegalStateException("分块快照已提交");
            }
            if (chunks == null || chunks.isEmpty()) {
                return;
            }
            writeBoundedPart(new ArrayList<DocumentChunk>(chunks));
        }

        /**
         * 最后写入清单并提交快照。
         *
         * @return 清单路径
         */
        public String finish() {
            if (committed) {
                throw new IllegalStateException("分块快照已提交");
            }
            if (manifest.getTotalChunks() <= 0 || manifest.getParts().isEmpty()) {
                throw new BusinessException("分块快照内容不完整");
            }
            byte[] payload = JSON.toJSONBytes(manifest);
            if (payload.length > MAX_MANIFEST_BYTES) {
                throw new BusinessException("分块快照清单过大，请调整文档后重试");
            }
            String fileName = manifest.getSession().getDocumentId() + "-manifest.json";
            String path = storageService.save(
                new CustomMultipartFile(payload, fileName, fileName, "application/json"),
                storagePrefix);
            if (!StringUtil.hasText(path)) {
                throw new BusinessException("分块快照清单保存失败");
            }
            committed = true;
            return path;
        }

        /**
         * 未提交时删除已经写入的分片。
         */
        @Override
        public void close() {
            if (committed) {
                return;
            }
            for (SnapshotPart part : manifest.getParts()) {
                if (StringUtil.hasText(part.getPath())) {
                    storageService.delete(part.getPath());
                }
            }
        }

        /**
         * 写入一个有界分片。
         *
         * @param chunks 分块
         */
        private void writeBoundedPart(List<DocumentChunk> chunks) {
            byte[] payload = JSON.toJSONBytes(chunks);
            if (payload.length > MAX_PART_BYTES && chunks.size() > 1) {
                int midpoint = chunks.size() / 2;
                writeBoundedPart(new ArrayList<DocumentChunk>(chunks.subList(0, midpoint)));
                writeBoundedPart(new ArrayList<DocumentChunk>(chunks.subList(midpoint, chunks.size())));
                return;
            }
            if (payload.length > MAX_PART_BYTES) {
                throw new BusinessException("单个分块快照超过 4 MiB，请缩小分块后重试");
            }
            int partNumber = manifest.getParts().size() + 1;
            String fileName = String.format("part-%06d.json", partNumber);
            String path = storageService.save(
                new CustomMultipartFile(payload, fileName, fileName, "application/json"),
                storagePrefix);
            if (!StringUtil.hasText(path)) {
                throw new BusinessException("分块快照分片保存失败");
            }
            SnapshotPart part = new SnapshotPart();
            part.setPath(path);
            part.setChunkCount(chunks.size());
            part.setByteLength(payload.length);
            part.setSha256(sha256(payload));
            manifest.getParts().add(part);
            manifest.setTotalChunks(manifest.getTotalChunks() + chunks.size());
        }
    }

    /**
     * V2 快照清单。
     */
    public static class SnapshotManifest {

        private String version;
        private DocumentImportDtos.PreviewSession session;
        private int totalChunks;
        private List<SnapshotPart> parts = new ArrayList<SnapshotPart>();
        private Date createdAt;

        /**
         * 返回版本。
         *
         * @return 版本
         */
        public String getVersion() {
            return version;
        }

        /**
         * 设置版本。
         *
         * @param version 版本
         */
        public void setVersion(String version) {
            this.version = version;
        }

        /**
         * 返回会话头。
         *
         * @return 会话头
         */
        public DocumentImportDtos.PreviewSession getSession() {
            return session;
        }

        /**
         * 设置会话头。
         *
         * @param session 会话头
         */
        public void setSession(DocumentImportDtos.PreviewSession session) {
            this.session = session;
        }

        /**
         * 返回总分块数。
         *
         * @return 总分块数
         */
        public int getTotalChunks() {
            return totalChunks;
        }

        /**
         * 设置总分块数。
         *
         * @param totalChunks 总分块数
         */
        public void setTotalChunks(int totalChunks) {
            this.totalChunks = totalChunks;
        }

        /**
         * 返回分片清单。
         *
         * @return 分片清单
         */
        public List<SnapshotPart> getParts() {
            return parts;
        }

        /**
         * 设置分片清单。
         *
         * @param parts 分片清单
         */
        public void setParts(List<SnapshotPart> parts) {
            this.parts = parts == null
                ? new ArrayList<SnapshotPart>()
                : parts;
        }

        /**
         * 返回创建时间。
         *
         * @return 创建时间
         */
        public Date getCreatedAt() {
            return createdAt;
        }

        /**
         * 设置创建时间。
         *
         * @param createdAt 创建时间
         */
        public void setCreatedAt(Date createdAt) {
            this.createdAt = createdAt;
        }
    }

    /**
     * V2 快照分片元信息。
     */
    public static class SnapshotPart {

        private String path;
        private int chunkCount;
        private long byteLength;
        private String sha256;

        /**
         * 返回存储路径。
         *
         * @return 存储路径
         */
        public String getPath() {
            return path;
        }

        /**
         * 设置存储路径。
         *
         * @param path 存储路径
         */
        public void setPath(String path) {
            this.path = path;
        }

        /**
         * 返回分块数。
         *
         * @return 分块数
         */
        public int getChunkCount() {
            return chunkCount;
        }

        /**
         * 设置分块数。
         *
         * @param chunkCount 分块数
         */
        public void setChunkCount(int chunkCount) {
            this.chunkCount = chunkCount;
        }

        /**
         * 返回字节数。
         *
         * @return 字节数
         */
        public long getByteLength() {
            return byteLength;
        }

        /**
         * 设置字节数。
         *
         * @param byteLength 字节数
         */
        public void setByteLength(long byteLength) {
            this.byteLength = byteLength;
        }

        /**
         * 返回 SHA-256。
         *
         * @return SHA-256
         */
        public String getSha256() {
            return sha256;
        }

        /**
         * 设置 SHA-256。
         *
         * @param sha256 SHA-256
         */
        public void setSha256(String sha256) {
            this.sha256 = sha256;
        }
    }

    /**
     * 内部读取结果。
     */
    private static final class SnapshotReadResult {

        private final SnapshotManifest manifest;
        private final DocumentImportDtos.PreviewSession legacySession;

        /**
         * 创建读取结果。
         *
         * @param manifest V2 清单
         * @param legacySession V1 会话
         */
        private SnapshotReadResult(
            SnapshotManifest manifest,
            DocumentImportDtos.PreviewSession legacySession) {
            this.manifest = manifest;
            this.legacySession = legacySession;
        }

        /**
         * 创建 V2 读取结果。
         *
         * @param manifest V2 清单
         * @return 读取结果
         */
        private static SnapshotReadResult forManifest(SnapshotManifest manifest) {
            return new SnapshotReadResult(manifest, null);
        }

        /**
         * 创建 V1 读取结果。
         *
         * @param session V1 会话
         * @return 读取结果
         */
        private static SnapshotReadResult forLegacy(
            DocumentImportDtos.PreviewSession session) {
            return new SnapshotReadResult(null, session);
        }
    }

    /**
     * 内部分片载荷。
     */
    private static final class PartPayload {

        private final List<DocumentChunk> chunks;
        private final long byteLength;

        /**
         * 创建分片载荷。
         *
         * @param chunks 分块
         * @param byteLength 字节数
         */
        private PartPayload(List<DocumentChunk> chunks, long byteLength) {
            this.chunks = chunks;
            this.byteLength = byteLength;
        }
    }
}
