package tech.easyflow.core.runtime;

import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.Locale;

/**
 * 将旧模型混在 {@code content} 中的 {@code <think>} 或 {@code <thinking>} 内容拆分为思考与正文增量。
 *
 * <p>解析仅在单轮回复开头生效，避免误处理正文中的标签示例。该解析器保留少量标签前缀，
 * 因而可以正确处理开始或结束标签被拆分到多个流式增量中的情况。</p>
 */
public final class LegacyThinkingTagParser {

    private static final String[] OPEN_TAGS = {"<think>", "<thinking>"};
    private static final String[] CLOSE_TAGS = {"</think>", "</thinking>"};
    private static final int MAX_LEADING_WHITESPACE = 64;

    private final StringBuilder pending = new StringBuilder();
    private Mode mode = Mode.UNDECIDED;
    private boolean legacyFormatDetected;
    private String activeCloseTag;

    /**
     * 接收普通正文增量并按需拆分旧版思考标签。
     *
     * @param delta 普通正文增量
     * @return 可立即发送的思考或正文片段
     */
    public List<Segment> acceptContent(String delta) {
        if (delta == null || delta.isEmpty()) {
            return Collections.emptyList();
        }
        if (mode == Mode.CONTENT || mode == Mode.BYPASS) {
            return List.of(Segment.content(delta));
        }
        if (mode == Mode.THINKING) {
            return consumeThinking(delta);
        }
        pending.append(delta);
        return resolveUndecided();
    }

    /**
     * 接收模型已经结构化返回的思考增量，非空时关闭旧标签自动识别。
     *
     * @param delta 结构化思考增量
     * @return 待发送片段，包含必要的前置缓冲与当前思考增量
     */
    public List<Segment> acceptReasoning(String delta) {
        if (delta == null || delta.isBlank()) {
            return Collections.emptyList();
        }
        List<Segment> segments = new ArrayList<>();
        flushBeforeStructuredReasoning(segments);
        mode = Mode.BYPASS;
        addSegment(segments, SegmentType.REASONING, delta);
        return segments;
    }

    /**
     * 收口尚未发送的标签前缀或内容。
     *
     * @return 剩余的思考或正文片段
     */
    public List<Segment> finish() {
        if (pending.length() == 0) {
            return Collections.emptyList();
        }
        SegmentType type = mode == Mode.THINKING ? SegmentType.REASONING : SegmentType.CONTENT;
        String text = pending.toString();
        pending.setLength(0);
        mode = Mode.CONTENT;
        return List.of(new Segment(type, text));
    }

    /**
     * 重置单轮解析状态。
     */
    public void reset() {
        pending.setLength(0);
        mode = Mode.UNDECIDED;
        legacyFormatDetected = false;
        activeCloseTag = null;
    }

    /**
     * 返回当前轮次是否识别到旧版思考标签。
     *
     * @return 识别到回复开头的旧思考标签时为 {@code true}
     */
    public boolean isLegacyFormatDetected() {
        return legacyFormatDetected;
    }

    private List<Segment> resolveUndecided() {
        int contentStart = firstContentIndex(pending);
        if (contentStart == pending.length()) {
            if (pending.length() <= MAX_LEADING_WHITESPACE) {
                return Collections.emptyList();
            }
            mode = Mode.CONTENT;
            return drainPending(SegmentType.CONTENT);
        }

        String candidate = pending.substring(contentStart);
        String normalized = candidate.toLowerCase(Locale.ROOT);
        int matchedTagIndex = matchingOpenTagIndex(normalized);
        if (matchedTagIndex < 0 && isPossibleOpenTagPrefix(normalized)) {
            return Collections.emptyList();
        }
        if (matchedTagIndex < 0) {
            mode = Mode.CONTENT;
            return drainPending(SegmentType.CONTENT);
        }

        legacyFormatDetected = true;
        mode = Mode.THINKING;
        activeCloseTag = CLOSE_TAGS[matchedTagIndex];
        String remainder = candidate.substring(OPEN_TAGS[matchedTagIndex].length());
        pending.setLength(0);
        return consumeThinking(remainder);
    }

    private List<Segment> consumeThinking(String delta) {
        pending.append(delta);
        String normalized = pending.toString().toLowerCase(Locale.ROOT);
        int closingIndex = normalized.indexOf(activeCloseTag);
        List<Segment> segments = new ArrayList<>();
        if (closingIndex >= 0) {
            addSegment(segments, SegmentType.REASONING, pending.substring(0, closingIndex));
            String remainder = pending.substring(closingIndex + activeCloseTag.length());
            pending.setLength(0);
            mode = Mode.CONTENT;
            addSegment(segments, SegmentType.CONTENT, remainder);
            return segments;
        }

        int retainedLength = closingTagPrefixLength(normalized);
        int emittedLength = pending.length() - retainedLength;
        if (emittedLength > 0) {
            addSegment(segments, SegmentType.REASONING, pending.substring(0, emittedLength));
            String retained = pending.substring(emittedLength);
            pending.setLength(0);
            pending.append(retained);
        }
        return segments;
    }

    private void flushBeforeStructuredReasoning(List<Segment> segments) {
        if (pending.length() == 0) {
            return;
        }
        if (mode == Mode.UNDECIDED && pending.toString().isBlank()) {
            pending.setLength(0);
            return;
        }
        SegmentType type = mode == Mode.THINKING ? SegmentType.REASONING : SegmentType.CONTENT;
        addSegment(segments, type, pending.toString());
        pending.setLength(0);
    }

    private List<Segment> drainPending(SegmentType type) {
        String text = pending.toString();
        pending.setLength(0);
        return text.isEmpty() ? Collections.emptyList() : List.of(new Segment(type, text));
    }

    private int firstContentIndex(CharSequence value) {
        int index = 0;
        while (index < value.length()) {
            char current = value.charAt(index);
            if (!Character.isWhitespace(current) && current != '\uFEFF') {
                break;
            }
            index++;
        }
        return index;
    }

    private int closingTagPrefixLength(String value) {
        int maxLength = Math.min(value.length(), activeCloseTag.length() - 1);
        for (int length = maxLength; length > 0; length--) {
            if (activeCloseTag.startsWith(value.substring(value.length() - length))) {
                return length;
            }
        }
        return 0;
    }

    private int matchingOpenTagIndex(String value) {
        for (int index = 0; index < OPEN_TAGS.length; index++) {
            if (value.startsWith(OPEN_TAGS[index])) {
                return index;
            }
        }
        return -1;
    }

    private boolean isPossibleOpenTagPrefix(String value) {
        for (String openTag : OPEN_TAGS) {
            if (openTag.startsWith(value)) {
                return true;
            }
        }
        return false;
    }

    private void addSegment(List<Segment> segments, SegmentType type, String text) {
        if (text == null || text.isEmpty()) {
            return;
        }
        if (!segments.isEmpty() && segments.get(segments.size() - 1).getType() == type) {
            Segment previous = segments.remove(segments.size() - 1);
            segments.add(new Segment(type, previous.getText() + text));
            return;
        }
        segments.add(new Segment(type, text));
    }

    private enum Mode {
        UNDECIDED,
        THINKING,
        CONTENT,
        BYPASS
    }

    /**
     * 兼容解析后的片段类型。
     */
    public enum SegmentType {
        /** 思考增量。 */
        REASONING,
        /** 最终回答增量。 */
        CONTENT
    }

    /**
     * 兼容解析后的不可变文本片段。
     */
    public static final class Segment {

        private final SegmentType type;
        private final String text;

        /**
         * 创建解析片段。
         *
         * @param type 片段类型
         * @param text 片段文本
         */
        public Segment(SegmentType type, String text) {
            this.type = type;
            this.text = text;
        }

        /**
         * 创建正文片段。
         *
         * @param text 正文文本
         * @return 正文片段
         */
        public static Segment content(String text) {
            return new Segment(SegmentType.CONTENT, text);
        }

        /**
         * 获取片段类型。
         *
         * @return 片段类型
         */
        public SegmentType getType() {
            return type;
        }

        /**
         * 获取片段文本。
         *
         * @return 片段文本
         */
        public String getText() {
            return text;
        }
    }
}
