package com.easyagents.rag.core;

import java.text.Normalizer;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;

/**
 * Provides conservative BGE-M3 token estimation and hard-limit text splitting.
 *
 * <p>BGE-M3 uses an XLM-R SentencePiece tokenizer. Loading the model vocabulary
 * during ingestion would introduce runtime model downloads, so this guard counts
 * NFKC-normalized Unicode code points. The estimate is intentionally conservative
 * for ordinary Chinese and mixed-language knowledge documents.</p>
 */
public final class BgeM3ChunkSafety {

    private static final double PREFERRED_BOUNDARY_MIN_RATIO = 0.8D;

    private BgeM3ChunkSafety() {
    }

    /**
     * Estimates the content token count after tokenizer-compatible normalization.
     *
     * @param content chunk content
     * @return conservative token estimate, or zero for empty content
     */
    public static int estimateContentTokens(String content) {
        if (content == null || content.isEmpty()) {
            return 0;
        }
        String normalized = Normalizer.normalize(content, Normalizer.Form.NFKC);
        return normalized.codePointCount(0, normalized.length());
    }

    /**
     * Checks whether content stays within the shared BGE-M3 embedding budget.
     *
     * @param content chunk content
     * @return true when the content can be sent to BGE-M3 safely
     */
    public static boolean isWithinHardLimit(String content) {
        return estimateContentTokens(content)
            <= RagDefaults.BGE_M3_HARD_CHUNK_TOKEN_LIMIT;
    }

    /**
     * Splits content using the shared BGE-M3 hard token limit.
     *
     * @param content source content
     * @return ordered, non-overlapping source ranges
     */
    public static List<ChunkRange> splitToHardLimit(String content) {
        return splitToTokenLimit(content, RagDefaults.BGE_M3_HARD_CHUNK_TOKEN_LIMIT);
    }

    /**
     * Splits content to an explicit conservative token limit.
     *
     * @param content source content
     * @param maxTokens maximum estimated tokens per part
     * @return ordered, non-overlapping source ranges
     * @throws IllegalArgumentException when maxTokens is not positive
     */
    public static List<ChunkRange> splitToTokenLimit(String content, int maxTokens) {
        if (maxTokens <= 0) {
            throw new IllegalArgumentException("maxTokens must be positive");
        }
        if (content == null || content.isEmpty()) {
            return Collections.emptyList();
        }
        if (estimateContentTokens(content) <= maxTokens) {
            return Collections.singletonList(new ChunkRange(0, content.length()));
        }

        List<ChunkRange> parts = new ArrayList<ChunkRange>();
        int start = 0;
        while (start < content.length()) {
            int end = findSplitEnd(content, start, maxTokens);
            if (end <= start) {
                throw new IllegalStateException("BGE-M3 hard split did not advance the text cursor");
            }
            parts.add(new ChunkRange(start, end));
            start = end;
        }
        return parts;
    }

    /**
     * Finds a safe source boundary while preferring nearby sentence endings.
     *
     * @param content complete source content
     * @param start current part start
     * @param maxTokens maximum estimated tokens
     * @return exclusive end offset
     */
    private static int findSplitEnd(String content, int start, int maxTokens) {
        int cursor = start;
        int currentTokens = 0;
        int preferredEnd = -1;
        int preferredBoundaryThreshold = Math.max(
            1, (int) Math.floor(maxTokens * PREFERRED_BOUNDARY_MIN_RATIO));
        while (cursor < content.length()) {
            int codePoint = content.codePointAt(cursor);
            int next = cursor + Character.charCount(codePoint);
            int nextTokens = currentTokens + estimateCodePointTokens(codePoint);
            if (nextTokens > maxTokens) {
                break;
            }
            cursor = next;
            currentTokens = nextTokens;
            if (currentTokens >= preferredBoundaryThreshold
                && isPreferredBoundary(codePoint)) {
                preferredEnd = cursor;
            }
        }
        int candidateEnd = preferredEnd > start ? preferredEnd : cursor;
        if (candidateEnd <= start) {
            candidateEnd = start + Character.charCount(content.codePointAt(start));
        }
        return fitWithinTokenLimit(content, start, candidateEnd, maxTokens);
    }

    /**
     * Shrinks a candidate when Unicode normalization expands its token estimate.
     *
     * @param content complete source content
     * @param start part start
     * @param candidateEnd proposed exclusive end
     * @param maxTokens maximum estimated tokens
     * @return exclusive end within the requested token limit
     */
    private static int fitWithinTokenLimit(
        String content,
        int start,
        int candidateEnd,
        int maxTokens) {
        if (estimateContentTokens(content.substring(start, candidateEnd))
            <= maxTokens) {
            return candidateEnd;
        }
        int low = 1;
        int high = content.codePointCount(start, candidateEnd);
        while (low < high) {
            int middle = low + (high - low + 1) / 2;
            int end = content.offsetByCodePoints(start, middle);
            if (estimateContentTokens(content.substring(start, end)) <= maxTokens) {
                low = middle;
            } else {
                high = middle - 1;
            }
        }
        return content.offsetByCodePoints(start, low);
    }

    /**
     * Estimates the normalized token cost of one source code point.
     *
     * @param codePoint source Unicode code point
     * @return conservative token estimate for the code point
     */
    private static int estimateCodePointTokens(int codePoint) {
        if (codePoint >= 0 && codePoint <= 0x7F) {
            return 1;
        }
        String source = new String(Character.toChars(codePoint));
        String normalized = Normalizer.normalize(source, Normalizer.Form.NFKC);
        return Math.max(1, normalized.codePointCount(0, normalized.length()));
    }

    /**
     * Determines whether a code point is a suitable semantic boundary.
     *
     * @param codePoint current Unicode code point
     * @return true for sentence, paragraph, or whitespace boundaries
     */
    private static boolean isPreferredBoundary(int codePoint) {
        return Character.isWhitespace(codePoint)
            || codePoint == '。'
            || codePoint == '！'
            || codePoint == '？'
            || codePoint == '；'
            || codePoint == '.'
            || codePoint == '!'
            || codePoint == '?'
            || codePoint == ';';
    }

    /**
     * Immutable UTF-16 source range returned by the hard splitter.
     */
    public static final class ChunkRange {

        private final int start;
        private final int end;

        /**
         * Creates a source range.
         *
         * @param start inclusive UTF-16 offset
         * @param end exclusive UTF-16 offset
         */
        public ChunkRange(int start, int end) {
            if (start < 0 || end < start) {
                throw new IllegalArgumentException("Invalid chunk range");
            }
            this.start = start;
            this.end = end;
        }

        /**
         * Returns the inclusive start offset.
         *
         * @return inclusive UTF-16 offset
         */
        public int getStart() {
            return start;
        }

        /**
         * Returns the exclusive end offset.
         *
         * @return exclusive UTF-16 offset
         */
        public int getEnd() {
            return end;
        }
    }
}
