package rag

import (
	"strings"
)

// 切块参数（针对中文 + embedding token 上限调过；后续可配置化）。
const (
	chunkTargetRunes  = 500  // 目标大小：贪心打包到接近此值就在句界收口
	chunkOverlapRunes = 80   // 块间重叠：保跨块上下文连续
	chunkMinRunes     = 100  // 最小块：低于则并入相邻，避免碎块稀释向量
	chunkMaxRunes     = 1000 // 原子硬上限：超大无标点段落兜底窗口切
)

// chunk 把文本切成检索友好的语义块：按段落/句界切成原子 → 贪心打包到目标大小（句末收口）
// → 块间加重叠。全程按 rune 操作，杜绝中文 UTF-8 被字节切碎。
func chunk(text string) []string {
	atoms := splitToAtoms(text)
	packed := packAtoms(atoms, chunkTargetRunes, chunkMinRunes)
	return addOverlap(packed, chunkOverlapRunes)
}

// splitToAtoms 把文本切成"原子"（句子/行）：在换行与句末标点处断开，去空白；
// 超大无标点原子按 rune 窗口兜底切到 ≤ 目标大小。原子是打包的最小不可分单元。
func splitToAtoms(text string) []string {
	runes := []rune(strings.ReplaceAll(text, "\r\n", "\n"))
	var atoms []string
	start := 0
	flush := func(end int) {
		if s := strings.TrimSpace(string(runes[start:end])); s != "" {
			atoms = append(atoms, s)
		}
		start = end
	}
	for i := 0; i < len(runes); i++ {
		r := runes[i]
		switch {
		case r == '\n' || isCJKEnd(r):
			flush(i + 1)
		case r == '.' || r == '!' || r == '?' || r == ';':
			// ASCII 句末：仅当其后为空白/行尾才断（避开缩写、小数点）。
			if i+1 >= len(runes) || runes[i+1] == ' ' || runes[i+1] == '\n' {
				flush(i + 1)
			}
		}
	}
	flush(len(runes))

	out := make([]string, 0, len(atoms))
	for _, a := range atoms {
		if runeLen(a) <= chunkMaxRunes {
			out = append(out, a)
		} else {
			out = append(out, splitByRuneWindow(a, chunkTargetRunes)...)
		}
	}
	return out
}

// packAtoms 贪心打包：把原子拼进当前块，直到再加会超 target 就收口（块在句界结束）。
// 末尾过小的块（< min）并入前一块，避免碎块。
func packAtoms(atoms []string, target, min int) []string {
	var chunks []string
	var cur strings.Builder
	curLen := 0
	closeCur := func() {
		if curLen > 0 {
			chunks = append(chunks, cur.String())
			cur.Reset()
			curLen = 0
		}
	}
	for _, a := range atoms {
		al := runeLen(a)
		if curLen > 0 && curLen+al > target {
			closeCur()
		}
		if curLen > 0 {
			cur.WriteByte('\n')
			curLen++
		}
		cur.WriteString(a)
		curLen += al
	}
	if curLen > 0 {
		if n := len(chunks); n > 0 && curLen < min {
			chunks[n-1] = chunks[n-1] + "\n" + cur.String() // 尾块太小 → 并入上一块
		} else {
			chunks = append(chunks, cur.String())
		}
	}
	return chunks
}

// addOverlap 给每块前缀上一块尾部的 overlap 个 rune，保跨块上下文连续（首块不加）。
func addOverlap(chunks []string, overlap int) []string {
	if overlap <= 0 || len(chunks) <= 1 {
		return chunks
	}
	out := make([]string, len(chunks))
	out[0] = chunks[0]
	for i := 1; i < len(chunks); i++ {
		prev := []rune(chunks[i-1])
		tail := prev
		if len(prev) > overlap {
			tail = prev[len(prev)-overlap:]
		}
		out[i] = strings.TrimSpace(string(tail)) + "\n" + chunks[i]
	}
	return out
}

func isCJKEnd(r rune) bool {
	switch r {
	case '。', '！', '？', '；', '…':
		return true
	}
	return false
}

// splitByRuneWindow 按 rune 窗口硬切（兜底：超大无标点原子），保证不切碎多字节字符。
func splitByRuneWindow(s string, size int) []string {
	r := []rune(s)
	var out []string
	for len(r) > size {
		if t := strings.TrimSpace(string(r[:size])); t != "" {
			out = append(out, t)
		}
		r = r[size:]
	}
	if t := strings.TrimSpace(string(r)); t != "" {
		out = append(out, t)
	}
	return out
}

func runeLen(s string) int { return len([]rune(s)) }