package com.easyagents.rag.ingestion.analysis;

import com.easyagents.core.util.StringUtil;
import com.easyagents.rag.ingestion.model.AnalysisResult;

import java.util.*;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

public class DocumentStructureAnalyzer {

    private static final Pattern MARKDOWN_HEADING = Pattern.compile("^#{1,6}\\s+\\S+.*$");
    private static final Pattern CHINESE_CHAPTER = Pattern.compile("^第[一二三四五六七八九十百零两0-9]+[章节篇部分卷]\\s*.*$");
    private static final Pattern CHINESE_SECTION = Pattern.compile("^[一二三四五六七八九十百零]+[、.．]\\s*\\S+.*$");
    private static final Pattern CHINESE_SUBSECTION = Pattern.compile("^[（(][一二三四五六七八九十百零0-9]+[)）]\\s*\\S+.*$");
    private static final Pattern NUMERIC_SECTION = Pattern.compile("^[0-9]+(\\.[0-9]+){0,4}\\s+\\S+.*$");
    private static final Pattern ENGLISH_SECTION = Pattern.compile("^(Chapter|Section|Part)\\s+[0-9IVXLC]+([.:\\-\\s].*)?$", Pattern.CASE_INSENSITIVE);
    private static final Pattern ENGLISH_ROMAN = Pattern.compile("^[IVXLC]+[.、)\\s-]+\\S+.*$");
    private static final Pattern QUESTION_LINE = Pattern.compile("^(Q|QUESTION|问|问题|FAQ\\s*[0-9]+)\\s*[.:：-]\\s*.+$", Pattern.CASE_INSENSITIVE);
    private static final Pattern ANSWER_LINE = Pattern.compile("^(A|ANSWER|答|答案)\\s*[.:：-]\\s*.+$", Pattern.CASE_INSENSITIVE);
    private static final Pattern PAGE_NUMBER = Pattern.compile("^(第?\\s*\\d+\\s*页|page\\s+\\d+|\\d+)\\s*$", Pattern.CASE_INSENSITIVE);
    private static final Pattern TOC_LINE = Pattern.compile("^.{2,80}[.·•…]{2,}\\s*\\d+\\s*$");

    public AnalysisResult analyze(String rawContent, String sourceFormat) {
        String normalizedContent = normalize(rawContent);
        List<String> lines = toLines(normalizedContent);

        int markdownHeadingCount = 0;
        int outlineHeadingCount = 0;
        int qaQuestionCount = 0;
        int qaAnswerCount = 0;
        int pairedQaCount = 0;
        int tocLineCount = 0;
        int shortLineCount = 0;
        Set<Integer> markdownLevels = new HashSet<Integer>();

        for (int i = 0; i < lines.size(); i++) {
            String line = lines.get(i);
            if (MARKDOWN_HEADING.matcher(line).matches()) {
                markdownHeadingCount++;
                markdownLevels.add(Integer.valueOf(countMarkdownLevel(line)));
            }
            if (isOutlineHeading(line)) {
                outlineHeadingCount++;
            }
            if (QUESTION_LINE.matcher(line).matches()) {
                qaQuestionCount++;
                if (hasAnswerNearby(lines, i)) {
                    pairedQaCount++;
                }
            }
            if (ANSWER_LINE.matcher(line).matches()) {
                qaAnswerCount++;
            }
            if (TOC_LINE.matcher(line).matches()) {
                tocLineCount++;
            }
            if (line.length() <= 20) {
                shortLineCount++;
            }
        }

        int paragraphCount = 0;
        int longParagraphCount = 0;
        List<String> paragraphs = splitParagraphs(normalizedContent);
        for (String paragraph : paragraphs) {
            if (StringUtil.hasText(paragraph)) {
                paragraphCount++;
                if (paragraph.length() > 800) {
                    longParagraphCount++;
                }
            }
        }

        Map<String, Object> features = new LinkedHashMap<String, Object>();
        features.put("lineCount", Integer.valueOf(lines.size()));
        features.put("paragraphCount", Integer.valueOf(paragraphCount));
        features.put("markdownHeadingCount", Integer.valueOf(markdownHeadingCount));
        features.put("markdownLevelVariety", Integer.valueOf(markdownLevels.size()));
        features.put("outlineHeadingCount", Integer.valueOf(outlineHeadingCount));
        features.put("qaQuestionCount", Integer.valueOf(qaQuestionCount));
        features.put("qaAnswerCount", Integer.valueOf(qaAnswerCount));
        features.put("pairedQaCount", Integer.valueOf(pairedQaCount));
        features.put("tocLineCount", Integer.valueOf(tocLineCount));
        features.put("shortLineRatio", lines.isEmpty() ? Double.valueOf(0D) : Double.valueOf((double) shortLineCount / (double) lines.size()));
        features.put("longParagraphCount", Integer.valueOf(longParagraphCount));

        AnalysisResult result = new AnalysisResult();
        result.setSourceFormat(sourceFormat == null ? "" : sourceFormat.toLowerCase(Locale.ROOT));
        result.setNormalizedContent(normalizedContent);
        result.setFeatures(features);
        return result;
    }

    private String normalize(String rawContent) {
        if (!StringUtil.hasText(rawContent)) {
            return "";
        }

        String content = rawContent
            .replace("\uFEFF", "")
            .replace("\u200B", "")
            .replace("\r\n", "\n")
            .replace('\r', '\n')
            .replace('\u00A0', ' ')
            .replace('：', ':');

        List<String> originalLines = toLines(content);
        Map<String, Integer> lineCounts = new HashMap<String, Integer>();
        for (String line : originalLines) {
            if (line.length() >= 4 && line.length() <= 60) {
                Integer count = lineCounts.get(line);
                lineCounts.put(line, count == null ? Integer.valueOf(1) : Integer.valueOf(count.intValue() + 1));
            }
        }

        List<String> filteredLines = new ArrayList<String>();
        for (String line : originalLines) {
            if (!StringUtil.hasText(line)) {
                filteredLines.add("");
                continue;
            }
            if (PAGE_NUMBER.matcher(line).matches()) {
                continue;
            }
            Integer repeated = lineCounts.get(line);
            if (repeated != null && repeated.intValue() >= 3 && line.length() <= 40) {
                continue;
            }
            filteredLines.add(line);
        }

        List<String> mergedLines = new ArrayList<String>();
        for (String line : filteredLines) {
            if (mergedLines.isEmpty()) {
                mergedLines.add(line);
                continue;
            }
            String previous = mergedLines.get(mergedLines.size() - 1);
            if (!StringUtil.hasText(previous) || !StringUtil.hasText(line)) {
                mergedLines.add(line);
                continue;
            }
            if (shouldMerge(previous, line)) {
                mergedLines.set(mergedLines.size() - 1, previous + joinToken(previous, line) + line);
            } else {
                mergedLines.add(line);
            }
        }

        return String.join("\n", mergedLines)
            .replaceAll("[ \\t]{2,}", " ")
            .replaceAll("\\n{3,}", "\n\n")
            .trim();
    }

    private boolean shouldMerge(String previous, String current) {
        if (isHeading(previous) || isHeading(current)) {
            return false;
        }
        if (QUESTION_LINE.matcher(current).matches() || ANSWER_LINE.matcher(current).matches()) {
            return false;
        }
        if (TOC_LINE.matcher(previous).matches() || TOC_LINE.matcher(current).matches()) {
            return false;
        }
        char previousChar = previous.charAt(previous.length() - 1);
        if ("。！？.!?:;：；".indexOf(previousChar) >= 0) {
            return false;
        }
        return current.length() < 80;
    }

    private String joinToken(String previous, String current) {
        char last = previous.charAt(previous.length() - 1);
        char first = current.charAt(0);
        if (Character.isLetterOrDigit(last) && Character.isLetterOrDigit(first)) {
            return " ";
        }
        return "";
    }

    private boolean hasAnswerNearby(List<String> lines, int index) {
        int end = Math.min(lines.size(), index + 4);
        for (int i = index + 1; i < end; i++) {
            if (ANSWER_LINE.matcher(lines.get(i)).matches()) {
                return true;
            }
        }
        return false;
    }

    private boolean isHeading(String line) {
        return MARKDOWN_HEADING.matcher(line).matches() || isOutlineHeading(line);
    }

    private boolean isOutlineHeading(String line) {
        return CHINESE_CHAPTER.matcher(line).matches()
            || CHINESE_SECTION.matcher(line).matches()
            || CHINESE_SUBSECTION.matcher(line).matches()
            || NUMERIC_SECTION.matcher(line).matches()
            || ENGLISH_SECTION.matcher(line).matches()
            || ENGLISH_ROMAN.matcher(line).matches();
    }

    private int countMarkdownLevel(String line) {
        Matcher matcher = Pattern.compile("^(#{1,6})\\s+").matcher(line);
        if (!matcher.find()) {
            return 0;
        }
        return matcher.group(1).length();
    }

    private List<String> splitParagraphs(String normalizedContent) {
        String[] parts = normalizedContent.split("\\n\\s*\\n");
        List<String> paragraphs = new ArrayList<String>();
        for (String part : parts) {
            String paragraph = part.trim();
            if (StringUtil.hasText(paragraph)) {
                paragraphs.add(paragraph);
            }
        }
        return paragraphs;
    }

    private List<String> toLines(String content) {
        String[] rawLines = content.split("\\n");
        List<String> lines = new ArrayList<String>(rawLines.length);
        for (String rawLine : rawLines) {
            lines.add(rawLine == null ? "" : rawLine.trim());
        }
        return lines;
    }
}
