package com.ankurm.ragingestion;

import java.util.ArrayList;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;

import org.springframework.ai.document.Document;
import org.springframework.ai.document.DocumentTransformer;

/**
 * Groups whole sentences into chunks of about {@code targetTokens}, repeating the last
 * {@code overlapSentences} sentences at the start of the next chunk. Unlike TokenTextSplitter it never cuts
 * inside a sentence, so an answer that sits in one sentence cannot be torn in half by a chunk boundary.
 */
public class SentenceSplitter implements DocumentTransformer {

    /** A full stop after a letter, followed by a capital: "UTC. Daily" splits, "1. Backups" does not. */
    private static final String SENTENCE_END = "(?<=[a-z\\)][.!?])\\s+(?=[A-Z])";

    private final int targetTokens;
    private final int overlapSentences;

    public SentenceSplitter(int targetTokens, int overlapSentences) {
        this.targetTokens = targetTokens;
        this.overlapSentences = overlapSentences;
    }

    @Override
    public List<Document> apply(List<Document> documents) {
        List<Document> out = new ArrayList<>();
        for (Document doc : documents) {
            List<String> sentences = new ArrayList<>();
            for (String line : doc.getText().split("\\R")) {
                String clean = line.replaceAll("\\s+", " ").strip();
                if (!clean.isEmpty()) {
                    sentences.addAll(List.of(clean.split(SENTENCE_END)));
                }
            }
            List<String> window = new ArrayList<>();
            int tokens = 0;
            int index = 0;
            for (String sentence : sentences) {
                int t = Tokens.count(sentence);
                if (!window.isEmpty() && tokens + t > targetTokens) {
                    out.add(chunk(doc, window, index++));
                    List<String> keep = overlapSentences == 0 ? List.of()
                            : window.subList(Math.max(0, window.size() - overlapSentences), window.size());
                    window = new ArrayList<>(keep);
                    tokens = window.stream().mapToInt(Tokens::count).sum();
                }
                window.add(sentence);
                tokens += t;
            }
            if (!window.isEmpty()) {
                out.add(chunk(doc, window, index));
            }
        }
        return out;
    }

    private static Document chunk(Document parent, List<String> sentences, int index) {
        Map<String, Object> metadata = new LinkedHashMap<>(parent.getMetadata());
        metadata.put("chunk_index", index);
        return Document.builder().text(String.join(" ", sentences)).metadata(metadata).build();
    }
}