package com.yanzhu.module.ai.framework.ai.compat;

import dev.langchain4j.data.document.Document;

import java.util.ArrayList;
import java.util.List;

/**
 * 基于 Token 的文本切分器实现
 * 
 * @author xiaoxin
 */
public class TokenTextSplitter implements TextSplitter {

    private final int chunkSize;
    private final boolean keepSeparator;

    private TokenTextSplitter(Builder builder) {
        this.chunkSize = builder.chunkSize;
        this.keepSeparator = builder.keepSeparator;
    }

    @Override
    public List<Document> apply(List<Document> documents) {
        List<Document> result = new ArrayList<>();
        for (Document document : documents) {
            result.addAll(splitDocument(document));
        }
        return result;
    }

    private List<Document> splitDocument(Document document) {
        List<Document> result = new ArrayList<>();
        String content = document.text();
        if (content == null || content.isEmpty()) {
            return result;
        }

        // 简单的基于字符的切分实现
        // 注意：这里是一个简化的实现，实际的Token计算可能更复杂
        int start = 0;
        int end = Math.min(chunkSize, content.length());

        while (start < content.length()) {
            // 尝试在单词边界处切分
            if (end < content.length()) {
                while (end > start && !Character.isWhitespace(content.charAt(end))) {
                    end--;
                }
                if (end == start) {
                    // 如果找不到单词边界，强制切分
                    end = Math.min(start + chunkSize, content.length());
                }
            }

            String chunk = content.substring(start, end);
            result.add(new Document(chunk));

            start = end;
            if (start < content.length() && keepSeparator) {
                // 跳过分隔符
                while (start < content.length() && Character.isWhitespace(content.charAt(start))) {
                    start++;
                }
            }
            end = Math.min(start + chunkSize, content.length());
        }

        return result;
    }

    public static Builder builder() {
        return new Builder();
    }

    public static class Builder {
        private int chunkSize = 1000;
        private boolean keepSeparator = true;

        public Builder withChunkSize(int chunkSize) {
            this.chunkSize = chunkSize;
            return this;
        }

        public Builder withKeepSeparator(boolean keepSeparator) {
            this.keepSeparator = keepSeparator;
            return this;
        }

        // 其他配置方法，根据需要添加
        public Builder withMinChunkSizeChars(int minChunkSizeChars) {
            // 忽略此配置
            return this;
        }

        public Builder withMinChunkLengthToEmbed(int minChunkLengthToEmbed) {
            // 忽略此配置
            return this;
        }

        public Builder withMaxNumChunks(int maxNumChunks) {
            // 忽略此配置
            return this;
        }

        public TokenTextSplitter build() {
            return new TokenTextSplitter(this);
        }
    }

}
